def _test_gradients(self, device):
"""
Tests if gradients pass though without any problems (infs, nans etc) and
also performs gradcheck (compares numerical and analytical gradients)
"""
test_random_input = self.create_random_sym3x3(device, n=16)
test_diag_input = self.create_diag_sym3x3(device, n=16)
test_almost_diag_input = self.create_diag_sym3x3(device, n=16, noise=1e-4)
test_input = torch.cat(
(test_random_input, test_diag_input, test_almost_diag_input)
)
test_input.requires_grad = True
with torch.autograd.detect_anomaly():
eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
loss = eigenvalues.mean() + eigenvectors.mean()
loss.backward()
test_random_input.requires_grad = True
# Inputs are converted to double to increase the precision of gradcheck.
torch.autograd.gradcheck(
symeig3x3, test_random_input.double(), eps=1e-6, atol=1e-2, rtol=1e-2
)
def test_more_dimensions(self):
"""
Tests if function supports arbitrary leading dimensions
"""
repeat = 4
test_input = self.create_random_sym3x3(self._cpu, n=16)
test_input_4d = test_input[None, ...].expand((repeat,) + test_input.shape)
eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
eigenvalues_4d, eigenvectors_4d = symeig3x3(test_input_4d, eigenvectors=True)
eigenvalues_4d_gt = eigenvalues[None, ...].expand((repeat,) + eigenvalues.shape)
eigenvectors_4d_gt = eigenvectors[None, ...].expand(
(repeat,) + eigenvectors.shape
)
self.assertClose(eigenvalues_4d_gt, eigenvalues_4d)
self.assertClose(eigenvectors_4d_gt, eigenvectors_4d)
def _test_eigenvectors_are_orthonormal(self, test_input):
"""
Verify that eigenvectors are an orthonormal set
"""
eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
batched_eye = torch.zeros_like(test_input)
batched_eye[..., :, :] = torch.eye(3, device=batched_eye.device)
self.assertClose(
batched_eye, eigenvectors @ eigenvectors.transpose(-2, -1), atol=1e-06
)
def _test_is_eigen(self, test_input, atol=1e-04, rtol=1e-02):
"""
Verify that values and vectors produced are really eigenvalues and eigenvectors
and can restore the original input matrix with good precision
"""
eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
self.assertClose(
test_input,
eigenvectors @ eigenvalues.diag_embed() @ eigenvectors.transpose(-2, -1),
atol=atol,
rtol=rtol,
)
def _test_eigenvalues_and_eigenvectors(
self, test_eigenvectors, test_eigenvalues, atol=1e-04, rtol=1e-04
):
test_input = (
test_eigenvectors.transpose(-2, -1)
@ test_eigenvalues.diag_embed()
@ test_eigenvectors
)
test_eigenvalues_sorted, _ = torch.sort(test_eigenvalues, dim=-1)
eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
self.assertClose(
test_eigenvalues_sorted,
eigenvalues,
atol=atol,
rtol=rtol,
)
self._test_is_not_nan_or_inf(test_input)
self._test_is_eigen(test_input, atol=atol, rtol=rtol)
self._test_eigenvectors_are_orthonormal(test_input)
from test_symeig3x3 import TestSymEig3x3
torch.set_num_threads(1)
CUDA_DEVICE = get_random_cuda_device()
def create_traced_func(func, device, batch_size):
traced_func = torch.jit.trace(
func, (TestSymEig3x3.create_random_sym3x3(device, batch_size), ))
return traced_func
FUNC_NAME_TO_FUNC = {
"sym3x3eig": (lambda inputs: symeig3x3(inputs, eigenvectors=True)),
"sym3x3eig_traced_cuda":
create_traced_func((lambda inputs: symeig3x3(inputs, eigenvectors=True)),
CUDA_DEVICE, 1024),
"torch_symeig": (lambda inputs: torch.symeig(inputs, eigenvectors=True)),
"torch_linalg_eigh": (lambda inputs: torch.linalg.eigh(inputs)),
"torch_pca_lowrank":
(lambda inputs: torch.pca_lowrank(inputs, center=False, niter=1)),
"sym3x3eig_no_vecs":
(lambda inputs: symeig3x3(inputs, eigenvectors=False)),
"torch_symeig_no_vecs":
(lambda inputs: torch.symeig(inputs, eigenvectors=False)),
"torch_linalg_eigvalsh_no_vecs":
(lambda inputs: torch.linalg.eigvalsh(inputs)),
}
def estimate_pointcloud_local_coord_frames(
pointclouds: Union[torch.Tensor, "Pointclouds"],
neighborhood_size: int = 50,
disambiguate_directions: bool = True,
*,
use_symeig_workaround: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Estimates the principal directions of curvature (which includes normals)
of a batch of `pointclouds`.
The algorithm first finds `neighborhood_size` nearest neighbors for each
point of the point clouds, followed by obtaining principal vectors of
covariance matrices of each of the point neighborhoods.
The main principal vector corresponds to the normals, while the
other 2 are the direction of the highest curvature and the 2nd highest
curvature.
Note that each principal direction is given up to a sign. Hence,
the function implements `disambiguate_directions` switch that allows
to ensure consistency of the sign of neighboring normals. The implementation
follows the sign disabiguation from SHOT descriptors [1].
The algorithm also returns the curvature values themselves.
These are the eigenvalues of the estimated covariance matrices
of each point neighborhood.
Args:
**pointclouds**: Batch of 3-dimensional points of shape
`(minibatch, num_point, 3)` or a `Pointclouds` object.
**neighborhood_size**: The size of the neighborhood used to estimate the
geometry around each point.
**disambiguate_directions**: If `True`, uses the algorithm from [1] to
ensure sign consistency of the normals of neighboring points.
**use_symeig_workaround**: If `True`, uses a custom eigenvalue
calculation.
Returns:
**curvatures**: The three principal curvatures of each point
of shape `(minibatch, num_point, 3)`.
If `pointclouds` are of `Pointclouds` class, returns a padded tensor.
**local_coord_frames**: The three principal directions of the curvature
around each point of shape `(minibatch, num_point, 3, 3)`.
The principal directions are stored in columns of the output.
E.g. `local_coord_frames[i, j, :, 0]` is the normal of
`j`-th point in the `i`-th pointcloud.
If `pointclouds` are of `Pointclouds` class, returns a padded tensor.
References:
[1] Tombari, Salti, Di Stefano: Unique Signatures of Histograms for
Local Surface Description, ECCV 2010.
"""
points_padded, num_points = convert_pointclouds_to_tensor(pointclouds)
ba, N, dim = points_padded.shape
if dim != 3:
raise ValueError(
"The pointclouds argument has to be of shape (minibatch, N, 3)")
if (num_points <= neighborhood_size).any():
raise ValueError("The neighborhood_size argument has to be" +
" >= size of each of the point clouds.")
# undo global mean for stability
# TODO: replace with tutil.wmean once landed
pcl_mean = points_padded.sum(1) / num_points[:, None]
points_centered = points_padded - pcl_mean[:, None, :]
# get the per-point covariance and nearest neighbors used to compute it
cov, knns = get_point_covariances(points_centered, num_points,
neighborhood_size)
# get the local coord frames as principal directions of
# the per-point covariance
# this is done with torch.symeig / torch.linalg.eigh, which returns the
# eigenvectors (=principal directions) in an ascending order of their
# corresponding eigenvalues, and the smallest eigenvalue's eigenvector
# corresponds to the normal direction; or with a custom equivalent.
if use_symeig_workaround:
curvatures, local_coord_frames = symeig3x3(cov, eigenvectors=True)
else:
curvatures, local_coord_frames = eigh(cov)
# disambiguate the directions of individual principal vectors
if disambiguate_directions:
# disambiguate normal
n = _disambiguate_vector_directions(points_centered, knns,
local_coord_frames[:, :, :, 0])
# disambiguate the main curvature
z = _disambiguate_vector_directions(points_centered, knns,
local_coord_frames[:, :, :, 2])
# the secondary curvature is just a cross between n and z
y = torch.cross(n, z, dim=2)
# cat to form the set of principal directions
local_coord_frames = torch.stack((n, y, z), dim=3)
return curvatures, local_coord_frames
def _test_is_not_nan_or_inf(self, test_input):
eigenvalues, eigenvectors = symeig3x3(test_input, eigenvectors=True)
self.assertTrue(torch.isfinite(eigenvalues).all())
self.assertTrue(torch.isfinite(eigenvectors).all())