learning3d 0.0.7__py3-none-any.whl → 0.2.0__py3-none-any.whl

utils/model_common_utils.py

@@ -0,0 +1,156 @@
+import torch
+
+def knn(x, k, add_one_to_k=False):
+    if add_one_to_k: k = k + 1
+    inner = -2 * torch.matmul(x.transpose(2, 1).contiguous(), x)
+    xx = torch.sum(x**2, dim=1, keepdim=True)
+    pairwise_distance = -xx - inner - xx.transpose(2, 1).contiguous()
+    idx = pairwise_distance.topk(k=k, dim=-1)[1]  # (batch_size, num_points, k)
+    return idx
+
+def pc_normalize(pc):
+    l = pc.shape[0]
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
+    pc = pc / m
+    return pc
+
+def square_distance(src, dst):
+    """
+    Calculate Euclid distance between each two points.
+    src^T * dst = xn * xm + yn * ym + zn * zm；
+    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
+    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
+    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
+         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
+    Input:
+        src: source points, [B, N, C]
+        dst: target points, [B, M, C]
+    Output:
+        dist: per-point square distance, [B, N, M]
+    """
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src ** 2, -1).view(B, N, 1)
+    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
+    return dist
+
+def index_points(points, idx):
+    """
+    Input:
+        points: input points data, [B, N, C]
+        idx: sample index data, [B, S]
+    Return:
+        new_points:, indexed points data, [B, S, C]
+    """
+    device = points.device
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
+    new_points = points[batch_indices, idx, :]
+    return new_points
+
+def farthest_point_sample(xyz, npoint, start_with_first_point=False):
+    """
+    Input:
+        xyz: pointcloud data, [B, N, C]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [B, npoint]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
+    distance = torch.ones(B, N).to(device) * 1e10
+    if not start_with_first_point:
+        farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
+    else:
+        farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device) * 0
+    batch_indices = torch.arange(B, dtype=torch.long).to(device)
+    for i in range(npoint):
+        centroids[:, i] = farthest
+        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
+        dist = torch.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = torch.max(distance, -1)[1]
+    return centroids
+
+def knn_point(k, pos1, pos2):
+    '''
+    Input:
+        k: int32, number of k in k-nn search
+        pos1: (batch_size, ndataset, c) float32 array, input points
+        pos2: (batch_size, npoint, c) float32 array, query points
+    Output:
+        val: (batch_size, npoint, k) float32 array, L2 distances
+        idx: (batch_size, npoint, k) int32 array, indices to input points
+    '''
+    B, N, C = pos1.shape
+    M = pos2.shape[1]
+    pos1 = pos1.view(B,1,N,-1).repeat(1,M,1,1)
+    pos2 = pos2.view(B,M,1,-1).repeat(1,1,N,1)
+    dist = torch.sum(-(pos1-pos2)**2,-1)
+    val,idx = dist.topk(k=k,dim = -1)
+    return torch.sqrt(-val), idx
+
+def query_ball_point(radius, nsample, xyz, new_xyz, get_cnt=False):
+    """
+    Input:
+        radius: local region radius
+        nsample: max sample number in local region
+        xyz: all points, [B, N, C]
+        new_xyz: query points, [B, S, C]
+    Return:
+        group_idx: grouped points index, [B, S, nsample]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    _, S, _ = new_xyz.shape
+    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
+    sqrdists = square_distance(new_xyz, xyz)
+    group_idx[sqrdists > radius ** 2] = N
+    
+    if get_cnt:
+        mask = group_idx != N
+        cnt = mask.sum(dim=-1)
+
+    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
+    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
+    mask = group_idx == N
+    group_idx[mask] = group_first[mask]
+    if get_cnt:
+        return group_idx, cnt
+    else:
+        return group_idx
+
+def get_graph_feature(x, k=20, device=None):
+    # x = x.squeeze()
+    x = x.view(*x.size()[:3])
+    idx = knn(x, k=k)  # (batch_size, num_points, k)
+    batch_size, num_points, _ = idx.size()
+
+    if device is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
+
+    idx = idx + idx_base
+
+    idx = idx.view(-1)
+
+    _, num_dims, _ = x.size()
+
+    x = x.transpose(2, 1).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
+    feature = x.view(batch_size * num_points, -1)[idx, :]
+    feature = feature.view(batch_size, num_points, k, num_dims)
+    x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)
+
+    feature = torch.cat((feature, x), dim=3).permute(0, 3, 1, 2)
+
+    return feature

learning3d/losses/cuda/chamfer_distance/__init__.py

@@ -1 +0,0 @@
-from .chamfer_distance import ChamferDistance

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)