mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2026-08-18 13:55:43 +08:00
2bce7110d5
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