mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2026-10-17 03:35:43 +08:00
716c594eab
Summary: ## Motivation `symeig3x3` uses a fixed absolute epsilon (`torch.finfo(torch.float).eps`) in several places where the magnitude of the input actually matters: - `p1.clamp(self._eps)` - the `diag_soft_cond` test, which decides whether a matrix is treated as diagonal - the `+= eps * sign(...)` regularizers in `_get_ev0` and `_get_ev1` Since those thresholds are absolute, the routine only behaves as intended for matrices whose entries are of order 1. Once the entries drop below roughly `1e-3` the regularizers dominate the data and the output stops being an eigendecomposition at all. At the other end, `1e10` is already wrong and `1e30` returns NaN. This is not an exotic corner. `estimate_pointcloud_normals` and `estimate_pointcloud_local_coord_frames` call `symeig3x3` by default (`use_symeig_workaround=True`) on per-point covariance matrices, and those scale with the square of the local neighbourhood size. The denser the point cloud, the smaller the covariances, and the worse the normals get, silently. Points sampled on a sphere, `neighborhood_size=50`, median angle between the estimated normal and the true radial normal: | points | radius | mean abs cov entry | main | this PR | `torch.linalg.eigh` | |--------|--------|--------------------|------|---------|---------------------| | 4 000 | 1.0 | 4.5e-03 | 1.60° | 1.61° | 1.61° | | 20 000 | 1.0 | 9.1e-04 | 1.55° | 0.68° | 0.68° | | 100 000 | 1.0 | 1.8e-04 | 31.86° | 0.31° | 0.31° | | 4 000 | 0.1 | 4.5e-05 | 50.24° | 1.54° | 1.54° | This looks like the cause of https://github.com/facebookresearch/pytorch3d/issues/1577, where the reporter narrowed bad normals down to `symeig3x3` and found that passing `use_symeig_workaround=False` fixed them. ## Fix The decomposition is homogeneous: `eig(s * A) == s * eig(A)` and the eigenvectors do not change. So preconditioning each matrix to unit maximum magnitude and undoing the scaling on the eigenvalues afterwards is exact, and it puts the absolute thresholds back in the regime they were written for. The scale is detached. For any fixed scale the expression is exactly the eigendecomposition of the input, so gradients are unaffected, and this avoids differentiating through `amax`. Max relative reconstruction error `max|V diag(w) Vᵀ - A| / max|A|` over 2 000 random symmetric matrices: | input scale | main | this PR | |-------------|---------|---------| | 1e30 | nan | 9.9e-07 | | 1e10 | 1.00 | 1.0e-06 | | 1e0 | 1.6e-06 | 1.0e-06 | | 1e-10 | 1.74 | 9.7e-07 | | 1e-30 | 1.73 | 9.7e-07 | float64 inputs break on main below about `1e-8` and are accurate here down to `1e-60`. Exactly-zero, identity, diagonal and rank-1 inputs, and batches mixing magnitudes within one call, all behave. ## Tests Added `test_is_eigen_small_magnitude`, which runs the existing eigen check on the same random matrices scaled by `1e-2`, `1e-4` and `1e-6` with a correspondingly scaled `atol`. On main it fails with a max relative difference of 101.6; with the fix it passes. The 13 existing tests in `tests/test_symeig3x3.py` pass unchanged both before and after, on CPU and on GPU, `gradcheck` included. Cost is one `amax` and one division: no measurable change on GPU, around 2% on CPU at 1e6 matrices. Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/2045 Reviewed By: MichaelRamamonjisoa Differential Revision: D119970731 Pulled By: bottler fbshipit-source-id: 35d85a342513c11faca4e6d8ee98052be90172a6
320 lines
11 KiB
Python
320 lines
11 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 torch
|
|
from pytorch3d.common.workaround import symeig3x3
|
|
from pytorch3d.transforms.rotation_conversions import random_rotations
|
|
|
|
from .common_testing import get_random_cuda_device, TestCaseMixin
|
|
|
|
|
|
class TestSymEig3x3(TestCaseMixin, unittest.TestCase):
|
|
TEST_BATCH_SIZE = 1024
|
|
|
|
@staticmethod
|
|
def create_random_sym3x3(device, n):
|
|
random_3x3 = torch.randn((n, 3, 3), device=device)
|
|
random_3x3_T = torch.transpose(random_3x3, 1, 2)
|
|
random_sym_3x3 = (random_3x3 * random_3x3_T).contiguous()
|
|
|
|
return random_sym_3x3
|
|
|
|
@staticmethod
|
|
def create_diag_sym3x3(device, n, noise=0.0):
|
|
# Create purly diagonal matrices
|
|
random_diag_3x3 = torch.randn((n, 3), device=device).diag_embed()
|
|
|
|
# Make them 'almost' diagonal
|
|
random_diag_3x3 += noise * TestSymEig3x3.create_random_sym3x3(device, n)
|
|
|
|
return random_diag_3x3
|
|
|
|
def setUp(self) -> None:
|
|
super().setUp()
|
|
torch.manual_seed(42)
|
|
|
|
self._gpu = get_random_cuda_device()
|
|
self._cpu = torch.device("cpu")
|
|
|
|
def test_is_eigen_gpu(self):
|
|
test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)
|
|
|
|
self._test_is_eigen(test_input)
|
|
|
|
def test_is_eigen_cpu(self):
|
|
test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)
|
|
|
|
self._test_is_eigen(test_input)
|
|
|
|
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_is_eigen_scaled_gpu(self):
|
|
self._test_is_eigen_scaled(self._gpu)
|
|
|
|
def test_is_eigen_scaled_cpu(self):
|
|
self._test_is_eigen_scaled(self._cpu)
|
|
|
|
def _test_is_eigen_scaled(self, device):
|
|
"""
|
|
The decomposition is homogeneous, so scaling the input must not affect
|
|
the relative accuracy of the result.
|
|
"""
|
|
test_input = self.create_random_sym3x3(device, n=self.TEST_BATCH_SIZE)
|
|
|
|
for scale in (1e-6, 1e-4, 1e-2, 1e2, 1e4, 1e6):
|
|
scaled_input = test_input * scale
|
|
|
|
self._test_is_not_nan_or_inf(scaled_input)
|
|
self._test_is_eigen(scaled_input, atol=1e-04 * scale, rtol=1e-02)
|
|
self._test_eigenvectors_are_orthonormal(scaled_input)
|
|
|
|
def test_mixed_magnitudes_gpu(self):
|
|
self._test_mixed_magnitudes(self._gpu)
|
|
|
|
def test_mixed_magnitudes_cpu(self):
|
|
self._test_mixed_magnitudes(self._cpu)
|
|
|
|
def _test_mixed_magnitudes(self, device):
|
|
"""
|
|
Each matrix in a batch must be conditioned on its own magnitude, so
|
|
matrices spanning many orders of magnitude can be batched together.
|
|
"""
|
|
test_input = self.create_random_sym3x3(device, n=self.TEST_BATCH_SIZE)
|
|
scales = torch.logspace(-6, 6, self.TEST_BATCH_SIZE, device=device)[
|
|
:, None, None
|
|
]
|
|
scaled_input = test_input * scales
|
|
|
|
eigenvalues, eigenvectors = symeig3x3(scaled_input, eigenvectors=True)
|
|
reconstruction = (
|
|
eigenvectors @ eigenvalues.diag_embed() @ eigenvectors.transpose(-2, -1)
|
|
)
|
|
|
|
self.assertClose(test_input, reconstruction / scales, atol=1e-04, rtol=1e-02)
|
|
self._test_eigenvectors_are_orthonormal(scaled_input)
|
|
|
|
def test_eigenvectors_are_orthonormal_gpu(self):
|
|
test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)
|
|
|
|
self._test_eigenvectors_are_orthonormal(test_input)
|
|
|
|
def test_eigenvectors_are_orthonormal_cpu(self):
|
|
test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)
|
|
|
|
self._test_eigenvectors_are_orthonormal(test_input)
|
|
|
|
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_not_nan_or_inf_gpu(self):
|
|
test_input = self.create_random_sym3x3(self._gpu, n=self.TEST_BATCH_SIZE)
|
|
|
|
self._test_is_not_nan_or_inf(test_input)
|
|
|
|
def test_is_not_nan_or_inf_cpu(self):
|
|
test_input = self.create_random_sym3x3(self._cpu, n=self.TEST_BATCH_SIZE)
|
|
|
|
self._test_is_not_nan_or_inf(test_input)
|
|
|
|
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())
|
|
|
|
def test_degenerate_inputs_gpu(self):
|
|
self._test_degenerate_inputs(self._gpu)
|
|
|
|
def test_degenerate_inputs_cpu(self):
|
|
self._test_degenerate_inputs(self._cpu)
|
|
|
|
def _test_degenerate_inputs(self, device):
|
|
"""
|
|
Test degenerate case when input matrices are diagonal or near-diagonal
|
|
"""
|
|
|
|
# Purely diagonal case
|
|
test_input = self.create_diag_sym3x3(device, self.TEST_BATCH_SIZE)
|
|
|
|
self._test_is_not_nan_or_inf(test_input)
|
|
self._test_is_eigen(test_input)
|
|
self._test_eigenvectors_are_orthonormal(test_input)
|
|
|
|
# Almost-diagonal case
|
|
test_input = self.create_diag_sym3x3(device, self.TEST_BATCH_SIZE, noise=1e-4)
|
|
|
|
self._test_is_not_nan_or_inf(test_input)
|
|
self._test_is_eigen(test_input)
|
|
self._test_eigenvectors_are_orthonormal(test_input)
|
|
|
|
def test_gradients_cpu(self):
|
|
self._test_gradients(self._cpu)
|
|
|
|
def test_gradients_gpu(self):
|
|
self._test_gradients(self._gpu)
|
|
|
|
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_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)
|
|
|
|
def test_degenerate_eigenvalues_gpu(self):
|
|
self._test_degenerate_eigenvalues(self._gpu)
|
|
|
|
def test_degenerate_eigenvalues_cpu(self):
|
|
self._test_degenerate_eigenvalues(self._cpu)
|
|
|
|
def test_degenerate_eigenvalues_scaled_gpu(self):
|
|
for scale in (1e-6, 1e6):
|
|
self._test_degenerate_eigenvalues(self._gpu, scale=scale)
|
|
|
|
def test_degenerate_eigenvalues_scaled_cpu(self):
|
|
for scale in (1e-6, 1e6):
|
|
self._test_degenerate_eigenvalues(self._cpu, scale=scale)
|
|
|
|
def _test_degenerate_eigenvalues(self, device, scale: float = 1.0):
|
|
"""
|
|
Test degenerate eigenvalues like zero-valued and with 2-/3-multiplicity.
|
|
These are the cases held together by the regularizers inside the
|
|
decomposition, which only bite at the right input magnitude.
|
|
"""
|
|
# Error tolerances for degenerate values are increased as things might become
|
|
# numerically unstable
|
|
deg_atol = 1e-3 * scale
|
|
deg_rtol = 1.0
|
|
|
|
# Construct random orthonormal sets
|
|
test_eigenvecs = random_rotations(n=self.TEST_BATCH_SIZE, device=device)
|
|
|
|
# Construct random eigenvalues
|
|
test_eigenvals = (
|
|
torch.randn((self.TEST_BATCH_SIZE, 3), device=test_eigenvecs.device) * scale
|
|
)
|
|
self._test_eigenvalues_and_eigenvectors(
|
|
test_eigenvecs, test_eigenvals, atol=deg_atol, rtol=deg_rtol
|
|
)
|
|
|
|
# First eigenvalue is always 0.0 here: [0.0 X Y]
|
|
test_eigenvals_with_zero = test_eigenvals.clone()
|
|
test_eigenvals_with_zero[..., 0] = 0.0
|
|
self._test_eigenvalues_and_eigenvectors(
|
|
test_eigenvecs, test_eigenvals_with_zero, atol=deg_atol, rtol=deg_rtol
|
|
)
|
|
|
|
# First two eigenvalues are always the same here: [X X Y]
|
|
test_eigenvals_with_multiplicity2 = test_eigenvals.clone()
|
|
test_eigenvals_with_multiplicity2[..., 1] = test_eigenvals_with_multiplicity2[
|
|
..., 0
|
|
]
|
|
self._test_eigenvalues_and_eigenvectors(
|
|
test_eigenvecs,
|
|
test_eigenvals_with_multiplicity2,
|
|
atol=deg_atol,
|
|
rtol=deg_rtol,
|
|
)
|
|
|
|
# All three eigenvalues are the same here: [X X X]
|
|
test_eigenvals_with_multiplicity3 = test_eigenvals_with_multiplicity2.clone()
|
|
test_eigenvals_with_multiplicity3[..., 2] = test_eigenvals_with_multiplicity2[
|
|
..., 0
|
|
]
|
|
self._test_eigenvalues_and_eigenvectors(
|
|
test_eigenvecs,
|
|
test_eigenvals_with_multiplicity3,
|
|
atol=deg_atol,
|
|
rtol=deg_rtol,
|
|
)
|
|
|
|
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)
|