def backward(ctx, grad_dists): grad_dists = grad_dists.contiguous() points, tris, idxs = ctx.saved_tensors grad_points, grad_tris = _C.point_face_dist_backward( points, tris, idxs, grad_dists ) return grad_points, None, grad_tris, None, None
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)