learning3d 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

learning3d/losses/cuda/chamfer_distance/chamfer_distance.cpp

@@ -1,185 +0,0 @@
-#include <torch/torch.h>
-
-// CUDA forward declarations
-int ChamferDistanceKernelLauncher(
-    const int b, const int n,
-    const float* xyz,
-    const int m,
-    const float* xyz2,
-    float* result,
-    int* result_i,
-    float* result2,
-    int* result2_i);
-
-int ChamferDistanceGradKernelLauncher(
-    const int b, const int n,
-    const float* xyz1,
-    const int m,
-    const float* xyz2,
-    const float* grad_dist1,
-    const int* idx1,
-    const float* grad_dist2,
-    const int* idx2,
-    float* grad_xyz1,
-    float* grad_xyz2);
-
-
-void chamfer_distance_forward_cuda(
-    const at::Tensor xyz1, 
-    const at::Tensor xyz2, 
-    const at::Tensor dist1, 
-    const at::Tensor dist2, 
-    const at::Tensor idx1, 
-    const at::Tensor idx2) 
-{
-    ChamferDistanceKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data<float>(),
-                                            xyz2.size(1), xyz2.data<float>(),
-                                            dist1.data<float>(), idx1.data<int>(),
-                                            dist2.data<float>(), idx2.data<int>());
-}
-
-void chamfer_distance_backward_cuda(
-    const at::Tensor xyz1,
-    const at::Tensor xyz2, 
-    at::Tensor gradxyz1, 
-    at::Tensor gradxyz2, 
-    at::Tensor graddist1, 
-    at::Tensor graddist2, 
-    at::Tensor idx1, 
-    at::Tensor idx2)
-{
-    ChamferDistanceGradKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data<float>(),
-                                           xyz2.size(1), xyz2.data<float>(),
-                                           graddist1.data<float>(), idx1.data<int>(),
-                                           graddist2.data<float>(), idx2.data<int>(),
-                                           gradxyz1.data<float>(), gradxyz2.data<float>());
-}
-
-
-void nnsearch(
-    const int b, const int n, const int m,
-    const float* xyz1,
-    const float* xyz2,
-    float* dist,
-    int* idx)
-{
-    for (int i = 0; i < b; i++) {
-        for (int j = 0; j < n; j++) {
-            const float x1 = xyz1[(i*n+j)*3+0];
-            const float y1 = xyz1[(i*n+j)*3+1];
-            const float z1 = xyz1[(i*n+j)*3+2];
-            double best = 0;
-            int besti = 0;
-            for (int k = 0; k < m; k++) {
-                const float x2 = xyz2[(i*m+k)*3+0] - x1;
-                const float y2 = xyz2[(i*m+k)*3+1] - y1;
-                const float z2 = xyz2[(i*m+k)*3+2] - z1;
-                const double d=x2*x2+y2*y2+z2*z2;
-                if (k==0 || d < best){
-                    best = d;
-                    besti = k;
-                }
-            }
-            dist[i*n+j] = best;
-            idx[i*n+j] = besti;
-        }
-    }
-}
-
-
-void chamfer_distance_forward(
-    const at::Tensor xyz1, 
-    const at::Tensor xyz2, 
-    const at::Tensor dist1, 
-    const at::Tensor dist2, 
-    const at::Tensor idx1, 
-    const at::Tensor idx2) 
-{
-    const int batchsize = xyz1.size(0);
-    const int n = xyz1.size(1);
-    const int m = xyz2.size(1);
-
-    const float* xyz1_data = xyz1.data<float>();
-    const float* xyz2_data = xyz2.data<float>();
-    float* dist1_data = dist1.data<float>();
-    float* dist2_data = dist2.data<float>();
-    int* idx1_data = idx1.data<int>();
-    int* idx2_data = idx2.data<int>();
-
-    nnsearch(batchsize, n, m, xyz1_data, xyz2_data, dist1_data, idx1_data);
-    nnsearch(batchsize, m, n, xyz2_data, xyz1_data, dist2_data, idx2_data);
-}
-
-
-void chamfer_distance_backward(
-    const at::Tensor xyz1, 
-    const at::Tensor xyz2, 
-    at::Tensor gradxyz1, 
-    at::Tensor gradxyz2, 
-    at::Tensor graddist1, 
-    at::Tensor graddist2, 
-    at::Tensor idx1, 
-    at::Tensor idx2) 
-{
-    const int b = xyz1.size(0);
-    const int n = xyz1.size(1);
-    const int m = xyz2.size(1);
-
-    const float* xyz1_data = xyz1.data<float>();
-    const float* xyz2_data = xyz2.data<float>();
-    float* gradxyz1_data = gradxyz1.data<float>();
-    float* gradxyz2_data = gradxyz2.data<float>();
-    float* graddist1_data = graddist1.data<float>();
-    float* graddist2_data = graddist2.data<float>();
-    const int* idx1_data = idx1.data<int>();
-    const int* idx2_data = idx2.data<int>();
-
-    for (int i = 0; i < b*n*3; i++)
-        gradxyz1_data[i] = 0;
-    for (int i = 0; i < b*m*3; i++)
-        gradxyz2_data[i] = 0;
-    for (int i = 0;i < b; i++) {
-        for (int j = 0; j < n; j++) {
-            const float x1 = xyz1_data[(i*n+j)*3+0];
-            const float y1 = xyz1_data[(i*n+j)*3+1];
-            const float z1 = xyz1_data[(i*n+j)*3+2];
-            const int j2 = idx1_data[i*n+j];
-
-            const float x2 = xyz2_data[(i*m+j2)*3+0];
-            const float y2 = xyz2_data[(i*m+j2)*3+1];
-            const float z2 = xyz2_data[(i*m+j2)*3+2];
-            const float g = graddist1_data[i*n+j]*2;
-
-            gradxyz1_data[(i*n+j)*3+0] += g*(x1-x2);
-            gradxyz1_data[(i*n+j)*3+1] += g*(y1-y2);
-            gradxyz1_data[(i*n+j)*3+2] += g*(z1-z2);
-            gradxyz2_data[(i*m+j2)*3+0] -= (g*(x1-x2));
-            gradxyz2_data[(i*m+j2)*3+1] -= (g*(y1-y2));
-            gradxyz2_data[(i*m+j2)*3+2] -= (g*(z1-z2));
-        }
-        for (int j = 0; j < m; j++) {
-            const float x1 = xyz2_data[(i*m+j)*3+0];
-            const float y1 = xyz2_data[(i*m+j)*3+1];
-            const float z1 = xyz2_data[(i*m+j)*3+2];
-            const int j2 = idx2_data[i*m+j];
-            const float x2 = xyz1_data[(i*n+j2)*3+0];
-            const float y2 = xyz1_data[(i*n+j2)*3+1];
-            const float z2 = xyz1_data[(i*n+j2)*3+2];
-            const float g = graddist2_data[i*m+j]*2;
-            gradxyz2_data[(i*m+j)*3+0] += g*(x1-x2);
-            gradxyz2_data[(i*m+j)*3+1] += g*(y1-y2);
-            gradxyz2_data[(i*m+j)*3+2] += g*(z1-z2);
-            gradxyz1_data[(i*n+j2)*3+0] -= (g*(x1-x2));
-            gradxyz1_data[(i*n+j2)*3+1] -= (g*(y1-y2));
-            gradxyz1_data[(i*n+j2)*3+2] -= (g*(z1-z2));
-        }
-    }
-}
-
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("forward", &chamfer_distance_forward, "ChamferDistance forward");
-    m.def("forward_cuda", &chamfer_distance_forward_cuda, "ChamferDistance forward (CUDA)");
-    m.def("backward", &chamfer_distance_backward, "ChamferDistance backward");
-    m.def("backward_cuda", &chamfer_distance_backward_cuda, "ChamferDistance backward (CUDA)");
-}

learning3d/losses/cuda/chamfer_distance/chamfer_distance.cu

@@ -1,209 +0,0 @@
-#include <ATen/ATen.h>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-__global__ 
-void ChamferDistanceKernel(
-	int b,
-	int n,
-	const float* xyz,
-	int m,
-	const float* xyz2,
-	float* result,
-	int* result_i)
-{
-	const int batch=512;
-	__shared__ float buf[batch*3];
-	for (int i=blockIdx.x;i<b;i+=gridDim.x){
-		for (int k2=0;k2<m;k2+=batch){
-			int end_k=min(m,k2+batch)-k2;
-			for (int j=threadIdx.x;j<end_k*3;j+=blockDim.x){
-				buf[j]=xyz2[(i*m+k2)*3+j];
-			}
-			__syncthreads();
-			for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
-				float x1=xyz[(i*n+j)*3+0];
-				float y1=xyz[(i*n+j)*3+1];
-				float z1=xyz[(i*n+j)*3+2];
-				int best_i=0;
-				float best=0;
-				int end_ka=end_k-(end_k&3);
-				if (end_ka==batch){
-					for (int k=0;k<batch;k+=4){
-						{
-							float x2=buf[k*3+0]-x1;
-							float y2=buf[k*3+1]-y1;
-							float z2=buf[k*3+2]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (k==0 || d<best){
-								best=d;
-								best_i=k+k2;
-							}
-						}
-						{
-							float x2=buf[k*3+3]-x1;
-							float y2=buf[k*3+4]-y1;
-							float z2=buf[k*3+5]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+1;
-							}
-						}
-						{
-							float x2=buf[k*3+6]-x1;
-							float y2=buf[k*3+7]-y1;
-							float z2=buf[k*3+8]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+2;
-							}
-						}
-						{
-							float x2=buf[k*3+9]-x1;
-							float y2=buf[k*3+10]-y1;
-							float z2=buf[k*3+11]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+3;
-							}
-						}
-					}
-				}else{
-					for (int k=0;k<end_ka;k+=4){
-						{
-							float x2=buf[k*3+0]-x1;
-							float y2=buf[k*3+1]-y1;
-							float z2=buf[k*3+2]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (k==0 || d<best){
-								best=d;
-								best_i=k+k2;
-							}
-						}
-						{
-							float x2=buf[k*3+3]-x1;
-							float y2=buf[k*3+4]-y1;
-							float z2=buf[k*3+5]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+1;
-							}
-						}
-						{
-							float x2=buf[k*3+6]-x1;
-							float y2=buf[k*3+7]-y1;
-							float z2=buf[k*3+8]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+2;
-							}
-						}
-						{
-							float x2=buf[k*3+9]-x1;
-							float y2=buf[k*3+10]-y1;
-							float z2=buf[k*3+11]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+3;
-							}
-						}
-					}
-				}
-				for (int k=end_ka;k<end_k;k++){
-					float x2=buf[k*3+0]-x1;
-					float y2=buf[k*3+1]-y1;
-					float z2=buf[k*3+2]-z1;
-					float d=x2*x2+y2*y2+z2*z2;
-					if (k==0 || d<best){
-						best=d;
-						best_i=k+k2;
-					}
-				}
-				if (k2==0 || result[(i*n+j)]>best){
-					result[(i*n+j)]=best;
-					result_i[(i*n+j)]=best_i;
-				}
-			}
-			__syncthreads();
-		}
-	}
-}
-
-void ChamferDistanceKernelLauncher(
-    const int b, const int n,
-    const float* xyz,
-    const int m,
-    const float* xyz2,
-    float* result,
-    int* result_i,
-    float* result2,
-    int* result2_i)
-{
-	ChamferDistanceKernel<<<dim3(32,16,1),512>>>(b, n, xyz, m, xyz2, result, result_i);
-	ChamferDistanceKernel<<<dim3(32,16,1),512>>>(b, m, xyz2, n, xyz, result2, result2_i);
-
-	cudaError_t err = cudaGetLastError();
-	if (err != cudaSuccess)
-	    printf("error in chamfer distance updateOutput: %s\n", cudaGetErrorString(err));
-}
-
-
-__global__ 
-void ChamferDistanceGradKernel(
-	int b, int n,
-	const float* xyz1,
-	int m,
-	const float* xyz2,
-	const float* grad_dist1,
-	const int* idx1,
-	float* grad_xyz1,
-	float* grad_xyz2)
-{
-	for (int i = blockIdx.x; i<b; i += gridDim.x) {
-		for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; j += blockDim.x*gridDim.y) {
-			float x1=xyz1[(i*n+j)*3+0];
-			float y1=xyz1[(i*n+j)*3+1];
-			float z1=xyz1[(i*n+j)*3+2];
-			int j2=idx1[i*n+j];
-			float x2=xyz2[(i*m+j2)*3+0];
-			float y2=xyz2[(i*m+j2)*3+1];
-			float z2=xyz2[(i*m+j2)*3+2];
-			float g=grad_dist1[i*n+j]*2;
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+0]),g*(x1-x2));
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+1]),g*(y1-y2));
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+2]),g*(z1-z2));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+0]),-(g*(x1-x2)));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+1]),-(g*(y1-y2)));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+2]),-(g*(z1-z2)));
-		}
-	}
-}
-
-void ChamferDistanceGradKernelLauncher(
-    const int b, const int n,
-    const float* xyz1,
-    const int m,
-    const float* xyz2,
-    const float* grad_dist1,
-    const int* idx1,
-    const float* grad_dist2,
-    const int* idx2,
-    float* grad_xyz1,
-    float* grad_xyz2)
-{
-	cudaMemset(grad_xyz1, 0, b*n*3*4);
-	cudaMemset(grad_xyz2, 0, b*m*3*4);
-	ChamferDistanceGradKernel<<<dim3(1,16,1), 256>>>(b, n, xyz1, m, xyz2, grad_dist1, idx1, grad_xyz1, grad_xyz2);
-	ChamferDistanceGradKernel<<<dim3(1,16,1), 256>>>(b, m, xyz2, n, xyz1, grad_dist2, idx2, grad_xyz2, grad_xyz1);
-
-	cudaError_t err = cudaGetLastError();
-  	if (err != cudaSuccess)
-	    printf("error in chamfer distance get grad: %s\n", cudaGetErrorString(err));
-}

learning3d/losses/cuda/chamfer_distance/chamfer_distance.py

@@ -1,66 +0,0 @@
-import torch
-from torch.utils.cpp_extension import load
-import os
-
-script_dir = os.path.dirname(__file__)
-sources = [
-    os.path.join(script_dir, "chamfer_distance.cpp"),
-    os.path.join(script_dir, "chamfer_distance.cu"),
-]
-
-cd = load(name="cd", sources=sources)
-
-
-class ChamferDistanceFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, xyz1, xyz2):
-        batchsize, n, _ = xyz1.size()
-        _, m, _ = xyz2.size()
-        xyz1 = xyz1.contiguous()
-        xyz2 = xyz2.contiguous()
-        dist1 = torch.zeros(batchsize, n)
-        dist2 = torch.zeros(batchsize, m)
-
-        idx1 = torch.zeros(batchsize, n, dtype=torch.int)
-        idx2 = torch.zeros(batchsize, m, dtype=torch.int)
-
-        if not xyz1.is_cuda:
-            cd.forward(xyz1, xyz2, dist1, dist2, idx1, idx2)
-        else:
-            dist1 = dist1.cuda()
-            dist2 = dist2.cuda()
-            idx1 = idx1.cuda()
-            idx2 = idx2.cuda()
-            cd.forward_cuda(xyz1, xyz2, dist1, dist2, idx1, idx2)
-
-        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
-
-        return dist1, dist2
-
-    @staticmethod
-    def backward(ctx, graddist1, graddist2):
-        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
-
-        graddist1 = graddist1.contiguous()
-        graddist2 = graddist2.contiguous()
-
-        gradxyz1 = torch.zeros(xyz1.size())
-        gradxyz2 = torch.zeros(xyz2.size())
-
-        if not graddist1.is_cuda:
-            cd.backward(
-                xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2
-            )
-        else:
-            gradxyz1 = gradxyz1.cuda()
-            gradxyz2 = gradxyz2.cuda()
-            cd.backward_cuda(
-                xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2
-            )
-
-        return gradxyz1, gradxyz2
-
-
-class ChamferDistance(torch.nn.Module):
-    def forward(self, xyz1, xyz2):
-        return ChamferDistanceFunction.apply(xyz1, xyz2)

learning3d/losses/cuda/emd_torch/pkg/emd_loss_layer.py

@@ -1,41 +0,0 @@
-import torch
-import torch.nn as nn
-
-import _emd_ext._emd as emd
-
-
-class EMDFunction(torch.autograd.Function):
-	@staticmethod
-	def forward(self, xyz1, xyz2):
-		cost, match = emd.emd_forward(xyz1, xyz2)
-		self.save_for_backward(xyz1, xyz2, match)
-		return cost
-
-
-	@staticmethod
-	def backward(self, grad_output):
-		xyz1, xyz2, match = self.saved_tensors
-		grad_xyz1, grad_xyz2 = emd.emd_backward(xyz1, xyz2, match)
-		return grad_xyz1, grad_xyz2
-
-
-
-
-class EMDLoss(nn.Module):
-	'''
-	Computes the (approximate) Earth Mover's Distance between two point sets. 
-
-	IMPLEMENTATION LIMITATIONS:
-	- Double tensors must have <=11 dimensions
-	- Float tensors must have <=23 dimensions
-	This is due to the use of CUDA shared memory in the computation. This shared memory is limited by the hardware to 48kB.
-	'''
-
-	def __init__(self):
-		super(EMDLoss, self).__init__()
-
-	def forward(self, xyz1, xyz2):
-
-		assert xyz1.shape[-1] == xyz2.shape[-1], 'Both point sets must have the same dimensions!'
-		assert xyz1.shape[1] == xyz2.shape[1], 'Both Point Clouds must have same number of points in it.'
-		return EMDFunction.apply(xyz1, xyz2)