Spaces:
Running
Running
| # 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 pytorch3d.transforms import acos_linear_extrapolation | |
| from .common_testing import TestCaseMixin | |
| class TestAcosLinearExtrapolation(TestCaseMixin, unittest.TestCase): | |
| def setUp(self) -> None: | |
| super().setUp() | |
| torch.manual_seed(42) | |
| np.random.seed(42) | |
| def init_acos_boundary_values(batch_size: int = 10000): | |
| """ | |
| Initialize a tensor containing values close to the bounds of the | |
| domain of `acos`, i.e. close to -1 or 1; and random values between (-1, 1). | |
| """ | |
| device = torch.device("cuda:0") | |
| # one quarter are random values between -1 and 1 | |
| x_rand = 2 * torch.rand(batch_size // 4, dtype=torch.float32, device=device) - 1 | |
| x = [x_rand] | |
| for bound in [-1, 1]: | |
| for above_bound in [True, False]: | |
| for noise_std in [1e-4, 1e-2]: | |
| n_generate = (batch_size - batch_size // 4) // 8 | |
| x_add = ( | |
| bound | |
| + (2 * float(above_bound) - 1) | |
| * torch.randn( | |
| n_generate, device=device, dtype=torch.float32 | |
| ).abs() | |
| * noise_std | |
| ) | |
| x.append(x_add) | |
| x = torch.cat(x) | |
| return x | |
| def acos_linear_extrapolation(batch_size: int): | |
| x = TestAcosLinearExtrapolation.init_acos_boundary_values(batch_size) | |
| torch.cuda.synchronize() | |
| def compute_acos(): | |
| acos_linear_extrapolation(x) | |
| torch.cuda.synchronize() | |
| return compute_acos | |
| def _test_acos_outside_bounds(self, x, y, dydx, bound): | |
| """ | |
| Check that `acos_linear_extrapolation` yields points on a line with correct | |
| slope, and that the function is continuous around `bound`. | |
| """ | |
| bound_t = torch.tensor(bound, device=x.device, dtype=x.dtype) | |
| # fit a line: slope * x + bias = y | |
| x_1 = torch.stack([x, torch.ones_like(x)], dim=-1) | |
| slope, bias = torch.linalg.lstsq(x_1, y[:, None]).solution.view(-1)[:2] | |
| desired_slope = (-1.0) / torch.sqrt(1.0 - bound_t**2) | |
| # test that the desired slope is the same as the fitted one | |
| self.assertClose(desired_slope.view(1), slope.view(1), atol=1e-2) | |
| # test that the autograd's slope is the same as the desired one | |
| self.assertClose(desired_slope.expand_as(dydx), dydx, atol=1e-2) | |
| # test that the value of the fitted line at x=bound equals | |
| # arccos(x), i.e. the function is continuous around the bound | |
| y_bound_lin = (slope * bound_t + bias).view(1) | |
| y_bound_acos = bound_t.acos().view(1) | |
| self.assertClose(y_bound_lin, y_bound_acos, atol=1e-2) | |
| def _one_acos_test(self, x: torch.Tensor, lower_bound: float, upper_bound: float): | |
| """ | |
| Test that `acos_linear_extrapolation` returns correct values for | |
| `x` between/above/below `lower_bound`/`upper_bound`. | |
| """ | |
| x.requires_grad = True | |
| x.grad = None | |
| y = acos_linear_extrapolation(x, [lower_bound, upper_bound]) | |
| # compute the gradient of the acos w.r.t. x | |
| y.backward(torch.ones_like(y)) | |
| dacos_dx = x.grad | |
| x_lower = x <= lower_bound | |
| x_upper = x >= upper_bound | |
| x_mid = (~x_lower) & (~x_upper) | |
| # test that between bounds, the function returns plain acos | |
| self.assertClose(x[x_mid].acos(), y[x_mid]) | |
| # test that outside the bounds, the function is linear with the right | |
| # slope and continuous around the bound | |
| self._test_acos_outside_bounds( | |
| x[x_upper], y[x_upper], dacos_dx[x_upper], upper_bound | |
| ) | |
| self._test_acos_outside_bounds( | |
| x[x_lower], y[x_lower], dacos_dx[x_lower], lower_bound | |
| ) | |
| def test_acos(self, batch_size: int = 10000): | |
| """ | |
| Tests whether the function returns correct outputs | |
| inside/outside the bounds. | |
| """ | |
| x = TestAcosLinearExtrapolation.init_acos_boundary_values(batch_size) | |
| bounds = 1 - 10.0 ** torch.linspace(-1, -5, 5) | |
| for lower_bound in -bounds: | |
| for upper_bound in bounds: | |
| if upper_bound < lower_bound: | |
| continue | |
| self._one_acos_test(x, float(lower_bound), float(upper_bound)) | |
| def test_finite_gradient(self, batch_size: int = 10000): | |
| """ | |
| Tests whether gradients stay finite close to the bounds. | |
| """ | |
| x = TestAcosLinearExtrapolation.init_acos_boundary_values(batch_size) | |
| x.requires_grad = True | |
| bounds = 1 - 10.0 ** torch.linspace(-1, -5, 5) | |
| for lower_bound in -bounds: | |
| for upper_bound in bounds: | |
| if upper_bound < lower_bound: | |
| continue | |
| x.grad = None | |
| y = acos_linear_extrapolation( | |
| x, | |
| [float(lower_bound), float(upper_bound)], | |
| ) | |
| self.assertTrue(torch.isfinite(y).all()) | |
| loss = y.mean() | |
| loss.backward() | |
| self.assertIsNotNone(x.grad) | |
| self.assertTrue(torch.isfinite(x.grad).all()) | |