def forward(ctx, points, points_first_idx, tris, tris_first_idx, max_points): """ Args: ctx: Context object used to calculate gradients. points: FloatTensor of shape `(P, 3)` points_first_idx: LongTensor of shape `(N,)` indicating the first point index in each example in the batch tris: FloatTensor of shape `(T, 3, 3)` of triangular faces. The `t`-th triangular face is spanned by `(tris[t, 0], tris[t, 1], tris[t, 2])` tris_first_idx: LongTensor of shape `(N,)` indicating the first face index in each example in the batch max_points: Scalar equal to maximum number of points in the batch Returns: dists: FloatTensor of shape `(P,)`, where `dists[p]` is the squared euclidean distance of `p`-th point to the closest triangular face in the corresponding example in the batch idxs: LongTensor of shape `(P,)` indicating the closest triangular face in the corresponding example in the batch. `dists[p]` is `d(points[p], tris[idxs[p], 0], tris[idxs[p], 1], tris[idxs[p], 2])` where `d(u, v0, v1, v2)` is the distance of point `u` from the triangular face `(v0, v1, v2)` """ dists, idxs = _C.point_face_dist_forward(points, points_first_idx, tris, tris_first_idx, max_points) ctx.save_for_backward(points, tris, idxs) return dists
def test_point_face_distance(self): """ Test CUDA implementation for PointFaceDistanceForward & PointFaceDistanceBackward """ device = get_random_cuda_device() N, V, F, P = 4, 32, 16, 24 meshes, pcls = self.init_meshes_clouds(N, V, F, P, device=device) # make points packed a leaf node points_packed = pcls.points_packed().detach().clone() # (P, 3) points_first_idx = pcls.cloud_to_packed_first_idx() max_p = pcls.num_points_per_cloud().max().item() # make edges packed a leaf node verts_packed = meshes.verts_packed() faces_packed = verts_packed[meshes.faces_packed()] # (T, 3, 3) faces_packed = faces_packed.clone().detach() faces_first_idx = meshes.mesh_to_faces_packed_first_idx() # leaf nodes points_packed.requires_grad = True faces_packed.requires_grad = True grad_dists = torch.rand( (points_packed.shape[0],), dtype=torch.float32, device=device ) # Cuda Implementation: forward dists_cuda, idx_cuda = _C.point_face_dist_forward( points_packed, points_first_idx, faces_packed, faces_first_idx, max_p ) # Cuda Implementation: backward grad_points_cuda, grad_faces_cuda = _C.point_face_dist_backward( points_packed, faces_packed, idx_cuda, grad_dists ) # Cpu Implementation: forward dists_cpu, idx_cpu = _C.point_face_dist_forward( points_packed.cpu(), points_first_idx.cpu(), faces_packed.cpu(), faces_first_idx.cpu(), max_p, ) # Cpu Implementation: backward # Note that using idx_cpu doesn't pass - there seems to be a problem with tied results. grad_points_cpu, grad_faces_cpu = _C.point_face_dist_backward( points_packed.cpu(), faces_packed.cpu(), idx_cuda.cpu(), grad_dists.cpu() ) # Naive Implementation: forward faces_list = packed_to_list(faces_packed, meshes.num_faces_per_mesh().tolist()) dists_naive = [] for i in range(N): points = pcls.points_list()[i] tris = faces_list[i] dists_temp = torch.zeros( (points.shape[0], tris.shape[0]), dtype=torch.float32, device=device ) for p in range(points.shape[0]): for t in range(tris.shape[0]): dist = self._point_to_tri_distance(points[p], tris[t]) dists_temp[p, t] = dist # torch.min() doesn't necessarily return the first index of the # smallest value, our warp_reduce does. So it's not straightforward # to directly compare indices, nor the gradients of grad_tris which # also depend on the indices of the minimum value. # To be able to compare, we will compare dists_temp.min(1) and # then feed the cuda indices to the naive output start = points_first_idx[i] end = points_first_idx[i + 1] if i < N - 1 else points_packed.shape[0] min_idx = idx_cuda.cpu()[start:end] - faces_first_idx[i].cpu() iidx = torch.arange(points.shape[0], device=device) min_dist = dists_temp[iidx, min_idx] dists_naive.append(min_dist) dists_naive = torch.cat(dists_naive) # Compare self.assertClose(dists_naive.cpu(), dists_cuda.cpu()) self.assertClose(dists_naive.cpu(), dists_cpu) # Naive Implementation: backward dists_naive.backward(grad_dists) grad_points_naive = torch.cat([cloud.grad for cloud in pcls.points_list()]) grad_faces_naive = faces_packed.grad.cpu() # Compare self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7) self.assertClose(grad_faces_naive, grad_faces_cuda.cpu(), atol=5e-7) self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7) self.assertClose(grad_faces_naive, grad_faces_cpu, atol=5e-7)