mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
d7d740abe9
Summary: Symmetric eigenvalues 3x3 implementation from https://github.com/fairinternal/denseposeslim/blob/roman_c3dpo/tools/functions.py#L612 based on https://en.wikipedia.org/wiki/Eigenvalue_algorithm#3.C3.973_matrices and https://www.geometrictools.com/Documentation/RobustEigenSymmetric3x3.pdf Benchmarks show significant outperformance of symeig3x3 in comparison with torch implementations (torch.symeig and torch.linalg.eigh) on GPU (P100), especially for large batches: 70-280ns per sample vs 3400ns per sample for torch_linalg_eigh_1048576_cpu It's worth mentioning that torch.linalg.eigh is still comparably fast for batches up to 8192 on CPU. Some tests are still failing as the error thresholds need to be adjusted appropriately. Reviewed By: patricklabatut Differential Revision: D29915453 fbshipit-source-id: 7c1b062da631c57c4e22a42dd0027ea5e205f1b5
264 lines
9.1 KiB
Python
264 lines
9.1 KiB
Python
# Copyright (c) Facebook, Inc. and its 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 common_testing import TestCaseMixin, get_random_cuda_device
|
|
from pytorch3d.common.workaround import symeig3x3
|
|
from pytorch3d.transforms.rotation_conversions import random_rotations
|
|
|
|
|
|
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_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(self, device):
|
|
"""
|
|
Test degenerate eigenvalues like zero-valued and with 2-/3-multiplicity
|
|
"""
|
|
# Error tolerances for degenerate values are increased as things might become
|
|
# numerically unstable
|
|
deg_atol = 1e-3
|
|
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
|
|
)
|
|
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)
|