mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
34bbb3ad32
Summary: Applies new import merging and sorting from µsort v1.0. When merging imports, µsort will make a best-effort to move associated comments to match merged elements, but there are known limitations due to the diynamic nature of Python and developer tooling. These changes should not produce any dangerous runtime changes, but may require touch-ups to satisfy linters and other tooling. Note that µsort uses case-insensitive, lexicographical sorting, which results in a different ordering compared to isort. This provides a more consistent sorting order, matching the case-insensitive order used when sorting import statements by module name, and ensures that "frog", "FROG", and "Frog" always sort next to each other. For details on µsort's sorting and merging semantics, see the user guide: https://usort.readthedocs.io/en/stable/guide.html#sorting Reviewed By: bottler Differential Revision: D35553814 fbshipit-source-id: be49bdb6a4c25264ff8d4db3a601f18736d17be1
938 lines
35 KiB
Python
938 lines
35 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import torch
|
|
from common_testing import get_random_cuda_device, TestCaseMixin
|
|
from pytorch3d import _C
|
|
from pytorch3d.loss import point_mesh_edge_distance, point_mesh_face_distance
|
|
from pytorch3d.structures import Meshes, packed_to_list, Pointclouds
|
|
|
|
|
|
class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
np.random.seed(42)
|
|
torch.manual_seed(42)
|
|
|
|
@staticmethod
|
|
def eps():
|
|
return 1e-8
|
|
|
|
@staticmethod
|
|
def min_triangle_area():
|
|
return 5e-3
|
|
|
|
@staticmethod
|
|
def init_meshes_clouds(
|
|
batch_size: int = 10,
|
|
num_verts: int = 1000,
|
|
num_faces: int = 3000,
|
|
num_points: int = 3000,
|
|
device: str = "cuda:0",
|
|
):
|
|
device = torch.device(device)
|
|
nump = torch.randint(low=1, high=num_points, size=(batch_size,))
|
|
numv = torch.randint(low=3, high=num_verts, size=(batch_size,))
|
|
numf = torch.randint(low=1, high=num_faces, size=(batch_size,))
|
|
verts_list = []
|
|
faces_list = []
|
|
points_list = []
|
|
for i in range(batch_size):
|
|
# Randomly choose vertices
|
|
verts = torch.rand((numv[i], 3), dtype=torch.float32, device=device)
|
|
verts.requires_grad_(True)
|
|
|
|
# Randomly choose faces. Our tests below compare argmin indices
|
|
# over faces and edges. Argmin is sensitive even to small numeral variations
|
|
# thus we make sure that faces are valid
|
|
# i.e. a face f = (i0, i1, i2) s.t. i0 != i1 != i2,
|
|
# otherwise argmin due to numeral sensitivities cannot be resolved
|
|
faces, allf = [], 0
|
|
validf = numv[i].item() - numv[i].item() % 3
|
|
while allf < numf[i]:
|
|
ff = torch.randperm(numv[i], device=device)[:validf].view(-1, 3)
|
|
faces.append(ff)
|
|
allf += ff.shape[0]
|
|
faces = torch.cat(faces, 0)
|
|
if faces.shape[0] > numf[i]:
|
|
faces = faces[: numf[i]]
|
|
|
|
verts_list.append(verts)
|
|
faces_list.append(faces)
|
|
|
|
# Randomly choose points
|
|
points = torch.rand((nump[i], 3), dtype=torch.float32, device=device)
|
|
points.requires_grad_(True)
|
|
|
|
points_list.append(points)
|
|
|
|
meshes = Meshes(verts_list, faces_list)
|
|
pcls = Pointclouds(points_list)
|
|
|
|
return meshes, pcls
|
|
|
|
@staticmethod
|
|
def _point_to_bary(point: torch.Tensor, tri: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Computes the barycentric coordinates of point wrt triangle (tri)
|
|
Note that point needs to live in the space spanned by tri = (a, b, c),
|
|
i.e. by taking the projection of an arbitrary point on the space spanned by tri
|
|
|
|
Args:
|
|
point: FloatTensor of shape (3)
|
|
tri: FloatTensor of shape (3, 3)
|
|
Returns:
|
|
bary: FloatTensor of shape (3)
|
|
"""
|
|
assert point.dim() == 1 and point.shape[0] == 3
|
|
assert tri.dim() == 2 and tri.shape[0] == 3 and tri.shape[1] == 3
|
|
|
|
a, b, c = tri.unbind(0)
|
|
|
|
v0 = b - a
|
|
v1 = c - a
|
|
v2 = point - a
|
|
|
|
d00 = v0.dot(v0)
|
|
d01 = v0.dot(v1)
|
|
d11 = v1.dot(v1)
|
|
d20 = v2.dot(v0)
|
|
d21 = v2.dot(v1)
|
|
|
|
denom = d00 * d11 - d01 * d01 + TestPointMeshDistance.eps()
|
|
s2 = (d11 * d20 - d01 * d21) / denom
|
|
s3 = (d00 * d21 - d01 * d20) / denom
|
|
s1 = 1.0 - s2 - s3
|
|
|
|
bary = torch.tensor([s1, s2, s3])
|
|
return bary
|
|
|
|
@staticmethod
|
|
def _is_inside_triangle(point: torch.Tensor, tri: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Computes whether point is inside triangle tri
|
|
Note that point needs to live in the space spanned by tri = (a, b, c)
|
|
i.e. by taking the projection of an arbitrary point on the space spanned by tri
|
|
|
|
Args:
|
|
point: FloatTensor of shape (3)
|
|
tri: FloatTensor of shape (3, 3)
|
|
Returns:
|
|
inside: BoolTensor of shape (1)
|
|
"""
|
|
v0 = tri[1] - tri[0]
|
|
v1 = tri[2] - tri[0]
|
|
area = torch.cross(v0, v1).norm() / 2.0
|
|
|
|
# check if triangle is a line or a point. In that case, return False
|
|
if area < 5e-3:
|
|
return False
|
|
bary = TestPointMeshDistance._point_to_bary(point, tri)
|
|
inside = ((bary >= 0.0) * (bary <= 1.0)).all()
|
|
return inside
|
|
|
|
@staticmethod
|
|
def _point_to_edge_distance(
|
|
point: torch.Tensor, edge: torch.Tensor
|
|
) -> torch.Tensor:
|
|
"""
|
|
Computes the squared euclidean distance of points to edges
|
|
Args:
|
|
point: FloatTensor of shape (3)
|
|
edge: FloatTensor of shape (2, 3)
|
|
Returns:
|
|
dist: FloatTensor of shape (1)
|
|
|
|
If a, b are the start and end points of the segments, we
|
|
parametrize a point p as
|
|
x(t) = a + t * (b - a)
|
|
To find t which describes p we minimize (x(t) - p) ^ 2
|
|
Note that p does not need to live in the space spanned by (a, b)
|
|
"""
|
|
s0, s1 = edge.unbind(0)
|
|
|
|
s01 = s1 - s0
|
|
norm_s01 = s01.dot(s01)
|
|
|
|
same_edge = norm_s01 < TestPointMeshDistance.eps()
|
|
if same_edge:
|
|
dist = 0.5 * (point - s0).dot(point - s0) + 0.5 * (point - s1).dot(
|
|
point - s1
|
|
)
|
|
return dist
|
|
|
|
t = s01.dot(point - s0) / norm_s01
|
|
t = torch.clamp(t, min=0.0, max=1.0)
|
|
x = s0 + t * s01
|
|
dist = (x - point).dot(x - point)
|
|
return dist
|
|
|
|
@staticmethod
|
|
def _point_to_tri_distance(point: torch.Tensor, tri: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Computes the squared euclidean distance of points to edges
|
|
Args:
|
|
point: FloatTensor of shape (3)
|
|
tri: FloatTensor of shape (3, 3)
|
|
Returns:
|
|
dist: FloatTensor of shape (1)
|
|
"""
|
|
a, b, c = tri.unbind(0)
|
|
cross = torch.cross(b - a, c - a)
|
|
norm = cross.norm()
|
|
normal = torch.nn.functional.normalize(cross, dim=0)
|
|
|
|
# p0 is the projection of p onto the plane spanned by (a, b, c)
|
|
# p0 = p + tt * normal, s.t. (p0 - a) is orthogonal to normal
|
|
# => tt = dot(a - p, n)
|
|
tt = normal.dot(a) - normal.dot(point)
|
|
p0 = point + tt * normal
|
|
dist_p = tt * tt
|
|
|
|
# Compute the distance of p to all edge segments
|
|
e01_dist = TestPointMeshDistance._point_to_edge_distance(point, tri[[0, 1]])
|
|
e02_dist = TestPointMeshDistance._point_to_edge_distance(point, tri[[0, 2]])
|
|
e12_dist = TestPointMeshDistance._point_to_edge_distance(point, tri[[1, 2]])
|
|
|
|
with torch.no_grad():
|
|
inside_tri = TestPointMeshDistance._is_inside_triangle(p0, tri)
|
|
|
|
if inside_tri and (norm > TestPointMeshDistance.eps()):
|
|
return dist_p
|
|
else:
|
|
if e01_dist.le(e02_dist) and e01_dist.le(e12_dist):
|
|
return e01_dist
|
|
elif e02_dist.le(e01_dist) and e02_dist.le(e12_dist):
|
|
return e02_dist
|
|
else:
|
|
return e12_dist
|
|
|
|
def test_point_edge_array_distance(self):
|
|
"""
|
|
Test CUDA implementation for PointEdgeArrayDistanceForward
|
|
& PointEdgeArrayDistanceBackward
|
|
"""
|
|
P, E = 16, 32
|
|
device = get_random_cuda_device()
|
|
points = torch.rand((P, 3), dtype=torch.float32, device=device)
|
|
edges = torch.rand((E, 2, 3), dtype=torch.float32, device=device)
|
|
|
|
# randomly make some edge points equal
|
|
same = torch.rand((E,), dtype=torch.float32, device=device) > 0.5
|
|
edges[same, 1] = edges[same, 0].clone().detach()
|
|
|
|
points_cpu = points.clone().cpu()
|
|
edges_cpu = edges.clone().cpu()
|
|
|
|
points.requires_grad = True
|
|
edges.requires_grad = True
|
|
grad_dists = torch.rand((P, E), dtype=torch.float32, device=device)
|
|
|
|
# Naive python implementation
|
|
dists_naive = torch.zeros((P, E), dtype=torch.float32, device=device)
|
|
for p in range(P):
|
|
for e in range(E):
|
|
dist = self._point_to_edge_distance(points[p], edges[e])
|
|
dists_naive[p, e] = dist
|
|
|
|
# Cuda Forward Implementation
|
|
dists_cuda = _C.point_edge_array_dist_forward(points, edges)
|
|
dists_cpu = _C.point_edge_array_dist_forward(points_cpu, edges_cpu)
|
|
|
|
# Compare
|
|
self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
|
|
self.assertClose(dists_naive.cpu(), dists_cpu)
|
|
|
|
# CUDA Bacwkard Implementation
|
|
grad_points_cuda, grad_edges_cuda = _C.point_edge_array_dist_backward(
|
|
points, edges, grad_dists
|
|
)
|
|
grad_points_cpu, grad_edges_cpu = _C.point_edge_array_dist_backward(
|
|
points_cpu, edges_cpu, grad_dists.cpu()
|
|
)
|
|
|
|
dists_naive.backward(grad_dists)
|
|
grad_points_naive = points.grad.cpu()
|
|
grad_edges_naive = edges.grad.cpu()
|
|
|
|
# Compare
|
|
self.assertClose(grad_points_naive, grad_points_cuda.cpu())
|
|
self.assertClose(grad_edges_naive, grad_edges_cuda.cpu())
|
|
self.assertClose(grad_points_naive, grad_points_cpu)
|
|
self.assertClose(grad_edges_naive, grad_edges_cpu)
|
|
|
|
def test_point_edge_distance(self):
|
|
"""
|
|
Test CUDA implementation for PointEdgeDistanceForward
|
|
& PointEdgeDistanceBackward
|
|
"""
|
|
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()
|
|
edges_packed = verts_packed[meshes.edges_packed()] # (E, 2, 3)
|
|
edges_packed = edges_packed.clone().detach()
|
|
|
|
edges_first_idx = meshes.mesh_to_edges_packed_first_idx()
|
|
|
|
# leaf nodes
|
|
points_packed.requires_grad = True
|
|
edges_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_edge_dist_forward(
|
|
points_packed, points_first_idx, edges_packed, edges_first_idx, max_p
|
|
)
|
|
# Cuda Implementation: backward
|
|
grad_points_cuda, grad_edges_cuda = _C.point_edge_dist_backward(
|
|
points_packed, edges_packed, idx_cuda, grad_dists
|
|
)
|
|
# Cpu Implementation: forward
|
|
dists_cpu, idx_cpu = _C.point_edge_dist_forward(
|
|
points_packed.cpu(),
|
|
points_first_idx.cpu(),
|
|
edges_packed.cpu(),
|
|
edges_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_edges_cpu = _C.point_edge_dist_backward(
|
|
points_packed.cpu(), edges_packed.cpu(), idx_cuda.cpu(), grad_dists.cpu()
|
|
)
|
|
|
|
# Naive Implementation: forward
|
|
edges_list = packed_to_list(edges_packed, meshes.num_edges_per_mesh().tolist())
|
|
dists_naive = []
|
|
for i in range(N):
|
|
points = pcls.points_list()[i]
|
|
edges = edges_list[i]
|
|
dists_temp = torch.zeros(
|
|
(points.shape[0], edges.shape[0]), dtype=torch.float32, device=device
|
|
)
|
|
for p in range(points.shape[0]):
|
|
for e in range(edges.shape[0]):
|
|
dist = self._point_to_edge_distance(points[p], edges[e])
|
|
dists_temp[p, e] = 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_edges 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[start:end] - edges_first_idx[i]
|
|
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_edges_naive = edges_packed.grad.cpu()
|
|
|
|
# Compare
|
|
self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
|
|
self.assertClose(grad_edges_naive, grad_edges_cuda.cpu(), atol=5e-7)
|
|
self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
|
|
self.assertClose(grad_edges_naive, grad_edges_cpu, atol=5e-7)
|
|
|
|
def test_edge_point_distance(self):
|
|
"""
|
|
Test CUDA implementation for EdgePointDistanceForward
|
|
& EdgePointDistanceBackward
|
|
"""
|
|
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()
|
|
|
|
# make edges packed a leaf node
|
|
verts_packed = meshes.verts_packed()
|
|
edges_packed = verts_packed[meshes.edges_packed()] # (E, 2, 3)
|
|
edges_packed = edges_packed.clone().detach()
|
|
|
|
edges_first_idx = meshes.mesh_to_edges_packed_first_idx()
|
|
max_e = meshes.num_edges_per_mesh().max().item()
|
|
|
|
# leaf nodes
|
|
points_packed.requires_grad = True
|
|
edges_packed.requires_grad = True
|
|
grad_dists = torch.rand(
|
|
(edges_packed.shape[0],), dtype=torch.float32, device=device
|
|
)
|
|
|
|
# Cuda Implementation: forward
|
|
dists_cuda, idx_cuda = _C.edge_point_dist_forward(
|
|
points_packed, points_first_idx, edges_packed, edges_first_idx, max_e
|
|
)
|
|
|
|
# Cuda Implementation: backward
|
|
grad_points_cuda, grad_edges_cuda = _C.edge_point_dist_backward(
|
|
points_packed, edges_packed, idx_cuda, grad_dists
|
|
)
|
|
|
|
# Cpu Implementation: forward
|
|
dists_cpu, idx_cpu = _C.edge_point_dist_forward(
|
|
points_packed.cpu(),
|
|
points_first_idx.cpu(),
|
|
edges_packed.cpu(),
|
|
edges_first_idx.cpu(),
|
|
max_e,
|
|
)
|
|
|
|
# Cpu Implementation: backward
|
|
grad_points_cpu, grad_edges_cpu = _C.edge_point_dist_backward(
|
|
points_packed.cpu(), edges_packed.cpu(), idx_cpu, grad_dists.cpu()
|
|
)
|
|
|
|
# Naive Implementation: forward
|
|
edges_list = packed_to_list(edges_packed, meshes.num_edges_per_mesh().tolist())
|
|
dists_naive = []
|
|
for i in range(N):
|
|
points = pcls.points_list()[i]
|
|
edges = edges_list[i]
|
|
dists_temp = torch.zeros(
|
|
(edges.shape[0], points.shape[0]), dtype=torch.float32, device=device
|
|
)
|
|
for e in range(edges.shape[0]):
|
|
for p in range(points.shape[0]):
|
|
dist = self._point_to_edge_distance(points[p], edges[e])
|
|
dists_temp[e, p] = 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_edges 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 = edges_first_idx[i]
|
|
end = edges_first_idx[i + 1] if i < N - 1 else edges_packed.shape[0]
|
|
|
|
min_idx = idx_cuda.cpu()[start:end] - points_first_idx[i].cpu()
|
|
iidx = torch.arange(edges.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_edges_naive = edges_packed.grad.cpu()
|
|
|
|
# Compare
|
|
self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
|
|
self.assertClose(grad_edges_naive, grad_edges_cuda.cpu(), atol=5e-7)
|
|
self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
|
|
self.assertClose(grad_edges_naive, grad_edges_cpu, atol=5e-7)
|
|
|
|
def test_point_mesh_edge_distance(self):
|
|
"""
|
|
Test point_mesh_edge_distance from pytorch3d.loss
|
|
"""
|
|
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)
|
|
|
|
# clone and detach for another backward pass through the op
|
|
verts_op = [verts.clone().detach() for verts in meshes.verts_list()]
|
|
for i in range(N):
|
|
verts_op[i].requires_grad = True
|
|
|
|
faces_op = [faces.clone().detach() for faces in meshes.faces_list()]
|
|
meshes_op = Meshes(verts=verts_op, faces=faces_op)
|
|
points_op = [points.clone().detach() for points in pcls.points_list()]
|
|
for i in range(N):
|
|
points_op[i].requires_grad = True
|
|
pcls_op = Pointclouds(points_op)
|
|
|
|
# Cuda implementation: forward & backward
|
|
loss_op = point_mesh_edge_distance(meshes_op, pcls_op)
|
|
|
|
# Naive implementation: forward & backward
|
|
edges_packed = meshes.edges_packed()
|
|
edges_list = packed_to_list(edges_packed, meshes.num_edges_per_mesh().tolist())
|
|
loss_naive = torch.zeros(N, dtype=torch.float32, device=device)
|
|
for i in range(N):
|
|
points = pcls.points_list()[i]
|
|
verts = meshes.verts_list()[i]
|
|
v_first_idx = meshes.mesh_to_verts_packed_first_idx()[i]
|
|
edges = verts[edges_list[i] - v_first_idx]
|
|
|
|
num_p = points.shape[0]
|
|
num_e = edges.shape[0]
|
|
dists = torch.zeros((num_p, num_e), dtype=torch.float32, device=device)
|
|
for p in range(num_p):
|
|
for e in range(num_e):
|
|
dist = self._point_to_edge_distance(points[p], edges[e])
|
|
dists[p, e] = dist
|
|
|
|
min_dist_p, min_idx_p = dists.min(1)
|
|
min_dist_e, min_idx_e = dists.min(0)
|
|
|
|
loss_naive[i] = min_dist_p.mean() + min_dist_e.mean()
|
|
loss_naive = loss_naive.mean()
|
|
|
|
# NOTE that hear the comparison holds despite the discrepancy
|
|
# due to the argmin indices returned by min(). This is because
|
|
# we don't will compare gradients on the verts and not on the
|
|
# edges or faces.
|
|
|
|
# Compare forward pass
|
|
self.assertClose(loss_op, loss_naive)
|
|
|
|
# Compare backward pass
|
|
rand_val = torch.rand(1).item()
|
|
grad_dist = torch.tensor(rand_val, dtype=torch.float32, device=device)
|
|
|
|
loss_naive.backward(grad_dist)
|
|
loss_op.backward(grad_dist)
|
|
|
|
# check verts grad
|
|
for i in range(N):
|
|
self.assertClose(
|
|
meshes.verts_list()[i].grad, meshes_op.verts_list()[i].grad
|
|
)
|
|
self.assertClose(pcls.points_list()[i].grad, pcls_op.points_list()[i].grad)
|
|
|
|
def test_point_face_array_distance(self):
|
|
"""
|
|
Test CUDA implementation for PointFaceArrayDistanceForward
|
|
& PointFaceArrayDistanceBackward
|
|
"""
|
|
P, T = 16, 32
|
|
device = get_random_cuda_device()
|
|
points = torch.rand((P, 3), dtype=torch.float32, device=device)
|
|
tris = torch.rand((T, 3, 3), dtype=torch.float32, device=device)
|
|
points_cpu = points.clone().cpu()
|
|
tris_cpu = tris.clone().cpu()
|
|
|
|
points.requires_grad = True
|
|
tris.requires_grad = True
|
|
grad_dists = torch.rand((P, T), dtype=torch.float32, device=device)
|
|
|
|
points_temp = points.clone().detach()
|
|
points_temp.requires_grad = True
|
|
tris_temp = tris.clone().detach()
|
|
tris_temp.requires_grad = True
|
|
|
|
# Naive python implementation
|
|
dists_naive = torch.zeros((P, T), dtype=torch.float32, device=device)
|
|
for p in range(P):
|
|
for t in range(T):
|
|
dist = self._point_to_tri_distance(points[p], tris[t])
|
|
dists_naive[p, t] = dist
|
|
|
|
# Naive Backward
|
|
dists_naive.backward(grad_dists)
|
|
grad_points_naive = points.grad.cpu()
|
|
grad_tris_naive = tris.grad.cpu()
|
|
|
|
# Cuda Forward Implementation
|
|
dists_cuda = _C.point_face_array_dist_forward(
|
|
points, tris, TestPointMeshDistance.min_triangle_area()
|
|
)
|
|
dists_cpu = _C.point_face_array_dist_forward(
|
|
points_cpu, tris_cpu, TestPointMeshDistance.min_triangle_area()
|
|
)
|
|
|
|
# Compare
|
|
self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
|
|
self.assertClose(dists_naive.cpu(), dists_cpu)
|
|
|
|
# CUDA Backward Implementation
|
|
grad_points_cuda, grad_tris_cuda = _C.point_face_array_dist_backward(
|
|
points, tris, grad_dists, TestPointMeshDistance.min_triangle_area()
|
|
)
|
|
grad_points_cpu, grad_tris_cpu = _C.point_face_array_dist_backward(
|
|
points_cpu,
|
|
tris_cpu,
|
|
grad_dists.cpu(),
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# Compare
|
|
self.assertClose(grad_points_naive, grad_points_cuda.cpu())
|
|
self.assertClose(grad_tris_naive, grad_tris_cuda.cpu(), atol=5e-6)
|
|
self.assertClose(grad_points_naive, grad_points_cpu)
|
|
self.assertClose(grad_tris_naive, grad_tris_cpu, atol=5e-6)
|
|
|
|
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,
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# Cuda Implementation: backward
|
|
grad_points_cuda, grad_faces_cuda = _C.point_face_dist_backward(
|
|
points_packed,
|
|
faces_packed,
|
|
idx_cuda,
|
|
grad_dists,
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# 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,
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# 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(),
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# 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)
|
|
|
|
def test_face_point_distance(self):
|
|
"""
|
|
Test CUDA implementation for FacePointDistanceForward
|
|
& FacePointDistanceBackward
|
|
"""
|
|
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()
|
|
|
|
# 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()
|
|
max_f = meshes.num_faces_per_mesh().max().item()
|
|
|
|
# leaf nodes
|
|
points_packed.requires_grad = True
|
|
faces_packed.requires_grad = True
|
|
grad_dists = torch.rand(
|
|
(faces_packed.shape[0],), dtype=torch.float32, device=device
|
|
)
|
|
|
|
# Cuda Implementation: forward
|
|
dists_cuda, idx_cuda = _C.face_point_dist_forward(
|
|
points_packed,
|
|
points_first_idx,
|
|
faces_packed,
|
|
faces_first_idx,
|
|
max_f,
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# Cuda Implementation: backward
|
|
grad_points_cuda, grad_faces_cuda = _C.face_point_dist_backward(
|
|
points_packed,
|
|
faces_packed,
|
|
idx_cuda,
|
|
grad_dists,
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# Cpu Implementation: forward
|
|
dists_cpu, idx_cpu = _C.face_point_dist_forward(
|
|
points_packed.cpu(),
|
|
points_first_idx.cpu(),
|
|
faces_packed.cpu(),
|
|
faces_first_idx.cpu(),
|
|
max_f,
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# Cpu Implementation: backward
|
|
grad_points_cpu, grad_faces_cpu = _C.face_point_dist_backward(
|
|
points_packed.cpu(),
|
|
faces_packed.cpu(),
|
|
idx_cpu,
|
|
grad_dists.cpu(),
|
|
TestPointMeshDistance.min_triangle_area(),
|
|
)
|
|
|
|
# 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(
|
|
(tris.shape[0], points.shape[0]), dtype=torch.float32, device=device
|
|
)
|
|
for t in range(tris.shape[0]):
|
|
for p in range(points.shape[0]):
|
|
dist = self._point_to_tri_distance(points[p], tris[t])
|
|
dists_temp[t, p] = 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 = faces_first_idx[i]
|
|
end = faces_first_idx[i + 1] if i < N - 1 else faces_packed.shape[0]
|
|
|
|
min_idx = idx_cuda.cpu()[start:end] - points_first_idx[i].cpu()
|
|
iidx = torch.arange(tris.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
|
|
|
|
# Compare
|
|
self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
|
|
self.assertClose(grad_faces_naive.cpu(), grad_faces_cuda.cpu(), atol=5e-7)
|
|
self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
|
|
self.assertClose(grad_faces_naive.cpu(), grad_faces_cpu, atol=5e-7)
|
|
|
|
def test_point_mesh_face_distance(self):
|
|
"""
|
|
Test point_mesh_face_distance from pytorch3d.loss
|
|
"""
|
|
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)
|
|
|
|
# clone and detach for another backward pass through the op
|
|
verts_op = [verts.clone().detach() for verts in meshes.verts_list()]
|
|
for i in range(N):
|
|
verts_op[i].requires_grad = True
|
|
|
|
faces_op = [faces.clone().detach() for faces in meshes.faces_list()]
|
|
meshes_op = Meshes(verts=verts_op, faces=faces_op)
|
|
points_op = [points.clone().detach() for points in pcls.points_list()]
|
|
for i in range(N):
|
|
points_op[i].requires_grad = True
|
|
pcls_op = Pointclouds(points_op)
|
|
|
|
# naive implementation
|
|
loss_naive = torch.zeros(N, dtype=torch.float32, device=device)
|
|
for i in range(N):
|
|
points = pcls.points_list()[i]
|
|
verts = meshes.verts_list()[i]
|
|
faces = meshes.faces_list()[i]
|
|
tris = verts[faces]
|
|
|
|
num_p = points.shape[0]
|
|
num_t = tris.shape[0]
|
|
dists = torch.zeros((num_p, num_t), dtype=torch.float32, device=device)
|
|
for p in range(num_p):
|
|
for t in range(num_t):
|
|
dist = self._point_to_tri_distance(points[p], tris[t])
|
|
dists[p, t] = dist
|
|
|
|
min_dist_p, min_idx_p = dists.min(1)
|
|
min_dist_t, min_idx_t = dists.min(0)
|
|
|
|
loss_naive[i] = min_dist_p.mean() + min_dist_t.mean()
|
|
loss_naive = loss_naive.mean()
|
|
|
|
# Op
|
|
loss_op = point_mesh_face_distance(meshes_op, pcls_op)
|
|
|
|
# Compare forward pass
|
|
self.assertClose(loss_op, loss_naive)
|
|
|
|
# Compare backward pass
|
|
rand_val = torch.rand(1).item()
|
|
grad_dist = torch.tensor(rand_val, dtype=torch.float32, device=device)
|
|
|
|
loss_naive.backward(grad_dist)
|
|
loss_op.backward(grad_dist)
|
|
|
|
# check verts grad
|
|
for i in range(N):
|
|
self.assertClose(
|
|
meshes.verts_list()[i].grad, meshes_op.verts_list()[i].grad
|
|
)
|
|
self.assertClose(pcls.points_list()[i].grad, pcls_op.points_list()[i].grad)
|
|
|
|
def test_small_faces_case(self):
|
|
for device in [torch.device("cpu"), torch.device("cuda:0")]:
|
|
mesh_vertices = torch.tensor(
|
|
[
|
|
[-0.0021, -0.3769, 0.7146],
|
|
[-0.0161, -0.3771, 0.7146],
|
|
[-0.0021, -0.3771, 0.7147],
|
|
],
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
mesh1_faces = torch.tensor([[0, 2, 1]], device=device)
|
|
mesh2_faces = torch.tensor([[2, 0, 1]], device=device)
|
|
pcd_points = torch.tensor([[-0.3623, -0.5340, 0.7727]], device=device)
|
|
mesh1 = Meshes(verts=[mesh_vertices], faces=[mesh1_faces])
|
|
mesh2 = Meshes(verts=[mesh_vertices], faces=[mesh2_faces])
|
|
pcd = Pointclouds(points=[pcd_points])
|
|
|
|
loss1 = point_mesh_face_distance(mesh1, pcd)
|
|
loss2 = point_mesh_face_distance(mesh2, pcd)
|
|
self.assertClose(loss1, loss2)
|
|
|
|
@staticmethod
|
|
def point_mesh_edge(N: int, V: int, F: int, P: int, device: str):
|
|
device = torch.device(device)
|
|
meshes, pcls = TestPointMeshDistance.init_meshes_clouds(
|
|
N, V, F, P, device=device
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
def loss():
|
|
point_mesh_edge_distance(meshes, pcls)
|
|
torch.cuda.synchronize()
|
|
|
|
return loss
|
|
|
|
@staticmethod
|
|
def point_mesh_face(N: int, V: int, F: int, P: int, device: str):
|
|
device = torch.device(device)
|
|
meshes, pcls = TestPointMeshDistance.init_meshes_clouds(
|
|
N, V, F, P, device=device
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
def loss():
|
|
point_mesh_face_distance(meshes, pcls)
|
|
torch.cuda.synchronize()
|
|
|
|
return loss
|