423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2026-08-18 13:55:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix and speed up the fast axis-angle conversions

Browse Source 
Summary:
GitHub issue #2002 (https://github.com/facebookresearch/pytorch3d/issues/2002)
points out that the near-pi branch of `matrix_to_axis_angle(..., fast=True)`
normalizes with `torch.norm(n)`, which reduces over the whole batch. That is
real, and there were two further problems next to it.

`matrix_to_axis_angle(..., fast=True)`:

- The axis was read from row 0 of `(R + I) / 2`, which is `n_x * n`. That is
  exactly zero whenever the axis is perpendicular to x, giving `nan` (for
  example for an exact rotation by pi about y), and is dominated by rounding
  when `abs(n_x)` is small. It is now read from
  `R + R^T - 2*cos(angle)*I == 2*(1 - cos(angle))*nnT`, taking the column whose
  diagonal entry is largest so that the multiplier is at least `1/sqrt(3)`.
  Symmetrizing also makes the identity exact at every angle rather than only at
  pi. The sign, which `nnT` does not determine, comes from `omegas`.
- `torch.norm(n)` becomes a per-row `torch.linalg.vector_norm`, which is the
  reported bug: with more than one near-pi rotation in a batch, every one of
  them was scaled wrongly.
- The branch threshold was `isclose(angle, pi)`. `omegas` is `2*sin(angle)`
  times the axis, so the other branch loses relative precision like
  `1/(pi - angle)`, and just outside `isclose` the float32 round-trip error
  reached 9.6e-4. The threshold is now `pi - 1e-2`, which leaves under 1% of
  uniformly random rotations on the more expensive branch.
- The `torch.isclose(angles, 0)` guard on `omegas` is removed. `torch.sinc(0)`
  is 1, so a zero angle was never a special case, and `torch.norm` has a zero
  rather than `nan` gradient at zero. The guard only zeroed the answer for
  angles below its `atol` of 1e-8, returning 0 instead of an exact 1e-9
  rotation, and cost 24% at a batch of 100k.

Worst-case float32 round-trip matrix error over random axes, before -> after:
exactly pi `nan` or 2.0 -> 3.6e-07; `pi - 1e-5` 2.0 -> 4.2e-07; `pi - 1e-4`
9.6e-04 -> 3.6e-07. The worst case at any angle is now 9.9e-06, at the branch
boundary, against 4.8e-07 for `fast=False`.

On the choice of `pi - 1e-2`: the new near-pi branch is accurate at any
threshold, so the threshold only decides how much of the batch takes the
slower branch, and the worst case is always the angle just below it. Measured
on one batch of 100k random float32 rotations on CPU, where `fast=False` took
11.4 ms, threshold against fraction of uniformly random rotations selected,
time, and worst-case round-trip error:

- `isclose(angle, pi)`, the previous behaviour: 0.003%, 8.6 ms, 1e-03
- `pi - 1e-3`: 0.1%, 8.9 ms, 9e-05
- `pi - 1e-2`, chosen: 0.7%, 8.9 ms, 1e-05
- `pi - 0.05`: 3.2%, 10.2 ms, 3e-06
- `3.0` radians: 9.0%, 17.6 ms, 9e-07

So `pi - 1e-2` buys two orders of magnitude of accuracy over the old threshold
for no measurable time, and it is the last threshold that is free; at `3.0`
radians `fast=True` would be slower than `fast=False`. The errors are sampled
over random axes, so they move in the last digit between runs.

`axis_angle_to_matrix(..., fast=True)` is Rodrigues' formula rewritten with
`cross_product_matrix^2 == outer_product - angle^2 * I` and
`(1 - cos(angle)) / angle^2 == sinc(angle / (2*pi))^2 / 2`. The first removes a
batched 3x3 matrix multiplication, which `bmm` serves poorly at that size, and
folds the leftover `-angle^2 * I` into the identity term as `cos(angle) * I`;
the second is defined at zero, so the `angles_sqrd == 0` special case goes
away. Values are unchanged to 8.9e-16.

`matrix_to_axis_angle` also used eight separate boolean mask indexes, each of
which re-runs `nonzero` over the whole batch and, on CUDA, synchronizes. The
cheap branch is now evaluated densely and the near-pi minority is selected
with a single `nonzero`.

Evaluating the near-pi branch densely too, which would remove the last
`nonzero` and so every data-dependent shape, was tried and rejected: against
this version it is 1.24x to 1.60x faster on CUDA at batches of 1k to 100k but
0.88x at 1M, and about 2x slower on CPU, and it does not survive
`torch.jit.script` without further work. It is the version to revisit if
export or `torch.compile` friendliness ever matters more than speed.

Speedups, before -> after:

- `matrix_to_axis_angle`: CUDA 1.88x at a batch of 1k, 1.66x at 100k, 1.47x at
  1M; CPU 1.26x at 1k, 2.01x at 100k, 1.71x at 1M.
- `axis_angle_to_matrix`: CUDA 1.26x at 1k, 1.81x at 100k, 5.49x at 1M; CPU
  unchanged.

Before this change `matrix_to_axis_angle(..., fast=True)` was slower than
`fast=False` on CUDA at a batch of 100k, 0.587 against 0.463 ms, because the
selection synchronized.

Rewriting `axis_angle_to_matrix` to compute all nine entries in one
`torch.stack`, as `quaternion_to_matrix` does, was also tried and is not
faster: 0.96x, 0.74x and 1.10x on CUDA at 1k, 100k and 1M, because the extra
kernel launches cost about what the saved intermediates gain.

Reviewed By: MichaelRamamonjisoa

Differential Revision: D115714860

fbshipit-source-id: cf19695f67bf2e2e6f8719419e9f902d5c58c309
This commit is contained in:
Jeremy Reizenstein
2026-08-17 07:59:33 -07:00
committed by meta-codesync[bot]
parent a5093b158b
commit 2bce7110d5
2 changed files with 208 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
pytorch3d/transforms/rotation_conversions.py
View File

@@ -489,15 +489,18 @@ def axis_angle_to_matrix(axis_angle: torch.Tensor, fast: bool = False) -> torch.
cross_product_matrix = torch.stack(
[zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=-1
).view(shape + (3,))
cross_product_matrix_sqrd = cross_product_matrix @ cross_product_matrix
outer_product = axis_angle.unsqueeze(-1) * axis_angle.unsqueeze(-2)
# This is the Rodrigues formula written with the identities
# cross_product_matrix^2 == outer_product - angle^2 * I
# (1 - cos(angle)) / angle^2 == sinc(angle / (2*pi))^2 / 2
# which save a matrix multiplication and, as sinc is defined at zero,
# remove the need to special-case a zero angle.
identity = torch.eye(3, dtype=dtype, device=device)
angles_sqrd = angles * angles
angles_sqrd = torch.where(angles_sqrd == 0, 1, angles_sqrd)
return (
identity.expand(cross_product_matrix.shape)
torch.cos(angles) * identity
+ torch.sinc(angles / torch.pi) * cross_product_matrix
+ ((1 - torch.cos(angles)) / angles_sqrd) * cross_product_matrix_sqrd
+ 0.5 * torch.sinc(angles / (2 * torch.pi)) ** 2 * outer_product
)
@@ -536,24 +539,52 @@ def matrix_to_axis_angle(matrix: torch.Tensor, fast: bool = False) -> torch.Tens
traces = torch.diagonal(matrix, dim1=-2, dim2=-1).sum(-1).unsqueeze(-1)
angles = torch.atan2(norms, traces - 1)
zeros = torch.zeros(3, dtype=matrix.dtype, device=matrix.device)
omegas = torch.where(torch.isclose(angles, torch.zeros_like(angles)), zeros, omegas)
# omegas is 2*sin(angle) times the axis, so the relative precision of the
# axis it gives decays like 1/(pi - angle). Above the following angle the
# axis is taken from the symmetric part of the matrix instead. Switching
# over 0.01 radians before pi keeps the error close to that of the slow
# implementation, while leaving under 1% of uniformly random rotations on
# the more expensive branch.
near_pi = angles > torch.pi - 1e-2
near_pi = angles.isclose(angles.new_full((1,), torch.pi)).squeeze(-1)
# An angle of zero needs no special case: torch.sinc(0) is 1, and omegas
# is then zero, which is the right answer. torch.sinc is instead zero at
# an angle of pi, so those entries are given a harmless denominator here
# and overwritten below.
sincs = torch.sinc(angles / torch.pi)
axis_angles = omegas * (0.5 / torch.where(near_pi, 1.0, sincs))
axis_angles = torch.empty_like(omegas)
axis_angles[~near_pi] = (
0.5 * omegas[~near_pi] / torch.sinc(angles[~near_pi] / torch.pi)
# Rotations by nearly pi are a small minority, so it is worth selecting
# them rather than evaluating the following densely. The selection index is
# computed once, because each boolean mask index would recompute it (and,
# on the GPU, synchronize).
flat_axis_angles = axis_angles.reshape(-1, 3)
index = near_pi.reshape(-1).nonzero().squeeze(-1)
near_pi_matrix = matrix.reshape(-1, 3, 3).index_select(0, index)
near_pi_omegas = omegas.reshape(-1, 3).index_select(0, index)
# this derives from: R + R^T - 2*cos(angle)*I = 2*(1 - cos(angle))*nnT,
# together with trace(R) - 1 == 2*cos(angle)
double_cosines = traces.reshape(-1, 1, 1).index_select(0, index) - 1
nnT = (
near_pi_matrix
+ near_pi_matrix.transpose(-1, -2)
- double_cosines * torch.eye(3, dtype=matrix.dtype, device=matrix.device)
)
# Every column of nnT is a multiple of n. Taking the one whose diagonal
# entry is largest means the multiplier has absolute value at least
# 1/sqrt(3), so the column is never degenerate.
largest = torch.diagonal(nnT, dim1=-2, dim2=-1).argmax(dim=-1)
n = nnT.gather(-1, largest.reshape(-1, 1, 1).expand(-1, 3, 1)).squeeze(-1)
n = n / torch.linalg.vector_norm(n, dim=-1, keepdim=True)
# nnT determines n only up to sign, which omegas resolves. (At exactly pi
# the sign is arbitrary, and there omegas is zero.)
n = torch.where((n * near_pi_omegas).sum(-1, keepdim=True) < 0, -n, n)
# this derives from: nnT = (R + 1) / 2
n = 0.5 * (
matrix[near_pi][..., 0, :]
+ torch.eye(1, 3, dtype=matrix.dtype, device=matrix.device)
near_pi_angles = angles.reshape(-1, 1).index_select(0, index)
return torch.index_copy(flat_axis_angles, 0, index, near_pi_angles * n).reshape(
axis_angles.shape
)
axis_angles[near_pi] = angles[near_pi] * n / torch.norm(n)
return axis_angles
def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:

159
tests/test_rotation_conversions.py
View File

@@ -228,6 +228,165 @@ class TestRotationConversion(TestCaseMixin, unittest.TestCase):
self.assertClose(data, axis_angle_to_matrix(euler_angles_fast))
self.assertClose(data, axis_angle_to_matrix(euler_angles, fast=True))
def test_axis_angle_by_pi(self):
"""Rotations by exactly pi are recovered, and don't depend on the batch."""
options = [0.0, -1.0, 1.0]
axes = torch.nn.functional.normalize(
torch.stack(
[
torch.tensor(vec, dtype=torch.float64)
for vec in itertools.islice( # exclude [0, 0, 0]
itertools.product(options, options, options), 1, None
)
]
),
dim=-1,
)
# Rotation by pi around unit vector x is given by 2 x x^T - Id.
R = 2 * torch.matmul(axes[..., None], axes[..., None, :]) - torch.eye(
3, dtype=torch.float64
)
for fast in [False, True]:
axis_angles = matrix_to_axis_angle(R, fast=fast)
self.assertClose(
axis_angles.norm(dim=-1),
torch.full(axes.shape[:1], math.pi, dtype=torch.float64),
)
self.assertClose(axis_angle_to_matrix(axis_angles), R)
# The axis of one rotation cannot depend on the others.
singly = torch.stack(
[matrix_to_axis_angle(r, fast=fast) for r in R.unbind()]
)
self.assertClose(singly, axis_angles)
def test_axis_angle_near_pi(self):
"""The fast path stays accurate as the angle approaches pi."""
for dtype, atol in [(torch.float64, 1e-9), (torch.float32, 1e-4)]:
axes = torch.nn.functional.normalize(
torch.randn(50, 3, dtype=dtype), dim=-1
)
for angle in [
math.pi,
math.pi - 1e-7,
math.pi - 1e-5,
math.pi - 1e-3,
math.pi - 1e-2,
math.pi - 0.1,
2.0,
]:
data = axes * angle
R = axis_angle_to_matrix(data)
axis_angles = matrix_to_axis_angle(R, fast=True)
self.assertClose(axis_angle_to_matrix(axis_angles), R, atol=atol)
self.assertClose(
axis_angles.norm(dim=-1),
torch.full(axes.shape[:1], angle, dtype=dtype),
atol=atol,
)
if angle < math.pi - 1e-2:
# Closer to pi than this, the sign of the axis is not
# determined to float32 precision (and both signs describe
# the same rotation to within the tolerance above).
self.assertClose(axis_angles, data, atol=atol)
def test_axis_angle_to_matrix_fast(self):
"""The Rodrigues implementation agrees with the quaternion one."""
data = torch.randn(100, 3, dtype=torch.float64)
normalized = torch.nn.functional.normalize(data, dim=-1)
data[:5] = 0.0
data[5:10] = normalized[5:10] * math.pi
data[10:15] = normalized[10:15] * 1e-12
self.assertClose(
axis_angle_to_matrix(data, fast=True),
axis_angle_to_matrix(data),
atol=1e-14,
)
# A zero angle is not a special case, but must still have gradients.
zeros = torch.zeros(1, 3, dtype=torch.float64, requires_grad=True)
matrices = axis_angle_to_matrix(zeros, fast=True)
self.assertClose(matrices, torch.eye(3, dtype=torch.float64)[None])
(grad,) = torch.autograd.grad(matrices.sum(), zeros)
self.assertTrue(torch.isfinite(grad).all())
def test_axis_angle_tiny(self):
"""Tiny angles are preserved, not flattened to zero, and have grads."""
axis = torch.nn.functional.normalize(
torch.tensor([[0.3, -0.5, 0.81]], dtype=torch.float64), dim=-1
)
for angle in [1e-3, 1e-8, 1e-9, 1e-12]:
data = axis * angle
R = axis_angle_to_matrix(data)
for fast in [False, True]:
# atol must be 0 here: the default would accept zero.
self.assertClose(matrix_to_axis_angle(R, fast=fast), data, atol=0)
identity = torch.eye(3, dtype=torch.float64)
for fast in [False, True]:
R = identity.clone().requires_grad_(True)
axis_angles = matrix_to_axis_angle(R, fast=fast)
self.assertClose(axis_angles, torch.zeros(3, dtype=torch.float64))
(grad,) = torch.autograd.grad(axis_angles.sum(), R)
self.assertTrue(torch.isfinite(grad).all())
def test_axis_angle_batch_independence(self):
"""Every matrix in a batch is converted independently of the others."""
angles = torch.tensor(
[math.pi, math.pi, math.pi - 1e-9, math.pi - 1e-3, 3.0, 1.0, 0.5, 0.0],
dtype=torch.float64,
)
axes = torch.nn.functional.normalize(
torch.randn(angles.shape[0], 3, dtype=torch.float64), dim=-1
)
R = axis_angle_to_matrix(axes * angles[:, None])
for fast in [False, True]:
batched = matrix_to_axis_angle(R, fast=fast)
singly = torch.stack(
[matrix_to_axis_angle(r, fast=fast) for r in R.unbind()]
)
self.assertClose(batched, singly)
def test_axis_angle_shapes(self):
"""Arbitrary leading dimensions, including none and zero."""
R = random_rotations(24, dtype=torch.float64).reshape(2, 3, 4, 3, 3)
# A rotation by pi, so that the near-pi branch is used too.
axis = torch.tensor([0.0, 0.6, -0.8], dtype=torch.float64)
R[1, 2, 0] = 2 * torch.outer(axis, axis) - torch.eye(3, dtype=torch.float64)
for fast in [False, True]:
axis_angles = matrix_to_axis_angle(R, fast=fast)
self.assertEqual(axis_angles.shape, (2, 3, 4, 3))
self.assertClose(axis_angle_to_matrix(axis_angles), R)
self.assertEqual(matrix_to_axis_angle(R[0, 0, 0], fast=fast).shape, (3,))
empty = torch.zeros(0, 3, 3, dtype=torch.float64)
self.assertEqual(matrix_to_axis_angle(empty, fast=fast).shape, (0, 3))
def test_axis_angle_near_pi_grad(self):
"""Round tripping is the identity, so its derivative is too."""
axis = torch.nn.functional.normalize(
torch.tensor([0.3, -0.5, 0.81], dtype=torch.float64), dim=-1
)
eye = torch.eye(3, dtype=torch.float64)
for angle in [0.5, 2.0, math.pi - 0.02, math.pi - 1e-4]:
data = axis * angle
for fast in [False, True]:
jacobian = torch.autograd.functional.jacobian(
lambda x, fast=fast: matrix_to_axis_angle(
axis_angle_to_matrix(x), fast=fast
),
data,
)
self.assertClose(jacobian, eye, atol=1e-8)
# At exactly pi the axis, and so the derivative, is ambiguous, but the
# result must still be usable.
R = (2 * torch.outer(axis, axis) - eye).requires_grad_(True)
for fast in [False, True]:
(grad,) = torch.autograd.grad(
matrix_to_axis_angle(R, fast=fast).sum(), R, retain_graph=True
)
self.assertTrue(torch.isfinite(grad).all())
def test_quaternion_application(self):
"""Applying a quaternion is the same as applying the matrix."""
quaternions = random_quaternions(3, torch.float64)