Jeremy Reizenstein 2bce7110d5 Fix and speed up the fast axis-angle conversions
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
2026-08-17 07:59:33 -07:00
2026-06-01 06:08:12 -07:00
2025-08-27 06:55:50 -07:00
2024-11-20 09:15:51 -08:00
2022-01-04 11:43:38 -08:00
2020-06-09 13:20:47 -07:00
2023-12-04 13:43:34 -08:00
2026-06-01 06:08:12 -07:00
2024-09-13 02:07:25 -07:00
2022-01-04 11:43:38 -08:00
2024-02-07 11:56:52 -08:00
2022-01-04 11:43:38 -08:00
2026-06-01 06:08:12 -07:00

CircleCI Anaconda-Server Badge

Introduction

PyTorch3D provides efficient, reusable components for 3D Computer Vision research with PyTorch.

Key features include:

  • Data structure for storing and manipulating triangle meshes
  • Efficient operations on triangle meshes (projective transformations, graph convolution, sampling, loss functions)
  • A differentiable mesh renderer
  • Implicitron, see its README, a framework for new-view synthesis via implicit representations. (blog post)

PyTorch3D is designed to integrate smoothly with deep learning methods for predicting and manipulating 3D data. For this reason, all operators in PyTorch3D:

  • Are implemented using PyTorch tensors
  • Can handle minibatches of hetereogenous data
  • Can be differentiated
  • Can utilize GPUs for acceleration

Within FAIR, PyTorch3D has been used to power research projects such as Mesh R-CNN.

See our blog post to see more demos and learn about PyTorch3D.

Installation

For detailed instructions refer to INSTALL.md.

License

PyTorch3D is released under the BSD License.

Tutorials

Get started with PyTorch3D by trying one of the tutorial notebooks.

Deform a sphere mesh to dolphin Bundle adjustment
Render textured meshes Camera position optimization
Render textured pointclouds Fit a mesh with texture
Render DensePose data Load & Render ShapeNet data
Fit Textured Volume Fit A Simple Neural Radiance Field
Fit Textured Volume in Implicitron Implicitron Config System

Documentation

Learn more about the API by reading the PyTorch3D documentation.

We also have deep dive notes on several API components:

Overview Video

We have created a short (~14 min) video tutorial providing an overview of the PyTorch3D codebase including several code examples. Click on the image below to watch the video on YouTube:

Development

We welcome new contributions to PyTorch3D and we will be actively maintaining this library! Please refer to CONTRIBUTING.md for full instructions on how to run the code, tests and linter, and submit your pull requests.

Development and Compatibility

  • main branch: actively developed, without any guarantee, Anything can be broken at any time
    • REMARK: this includes nightly builds which are built from main
    • HINT: the commit history can help locate regressions or changes
  • backward-compatibility between releases: no guarantee. Best efforts to communicate breaking changes and facilitate migration of code or data (incl. models).

Contributors

PyTorch3D is written and maintained by the Facebook AI Research Computer Vision Team.

In alphabetical order:

  • Amitav Baruah
  • Steve Branson
  • Krzysztof Chalupka
  • Jiali Duan
  • Luya Gao
  • Georgia Gkioxari
  • Taylor Gordon
  • Justin Johnson
  • Patrick Labatut
  • Christoph Lassner
  • Wan-Yen Lo
  • David Novotny
  • Nikhila Ravi
  • Jeremy Reizenstein
  • Dave Schnizlein
  • Roman Shapovalov
  • Olivia Wiles

Citation

If you find PyTorch3D useful in your research, please cite our tech report:

@article{ravi2020pytorch3d,
    author = {Nikhila Ravi and Jeremy Reizenstein and David Novotny and Taylor Gordon
                  and Wan-Yen Lo and Justin Johnson and Georgia Gkioxari},
    title = {Accelerating 3D Deep Learning with PyTorch3D},
    journal = {arXiv:2007.08501},
    year = {2020},
}

If you are using the pulsar backend for sphere-rendering (the PulsarPointRenderer or pytorch3d.renderer.points.pulsar.Renderer), please cite the tech report:

@article{lassner2020pulsar,
    author = {Christoph Lassner and Michael Zollh\"ofer},
    title = {Pulsar: Efficient Sphere-based Neural Rendering},
    journal = {arXiv:2004.07484},
    year = {2020},
}

News

Please see below for a timeline of the codebase updates in reverse chronological order. We are sharing updates on the releases as well as research projects which are built with PyTorch3D. The changelogs for the releases are available under Releases, and the builds can be installed using conda as per the instructions in INSTALL.md.

[Oct 31st 2023]: PyTorch3D v0.7.5 released.

[May 10th 2023]: PyTorch3D v0.7.4 released.

[Apr 5th 2023]: PyTorch3D v0.7.3 released.

[Dec 19th 2022]: PyTorch3D v0.7.2 released.

[Oct 23rd 2022]: PyTorch3D v0.7.1 released.

[Aug 10th 2022]: PyTorch3D v0.7.0 released with Implicitron and MeshRasterizerOpenGL.

[Apr 28th 2022]: PyTorch3D v0.6.2 released

[Dec 16th 2021]: PyTorch3D v0.6.1 released

[Oct 6th 2021]: PyTorch3D v0.6.0 released

[Aug 5th 2021]: PyTorch3D v0.5.0 released

[Feb 9th 2021]: PyTorch3D v0.4.0 released with support for implicit functions, volume rendering and a reimplementation of NeRF.

[November 2nd 2020]: PyTorch3D v0.3.0 released, integrating the pulsar backend.

[Aug 28th 2020]: PyTorch3D v0.2.5 released

[July 17th 2020]: PyTorch3D tech report published on ArXiv: https://arxiv.org/abs/2007.08501

[April 24th 2020]: PyTorch3D v0.2.0 released

[March 25th 2020]: SynSin codebase released using PyTorch3D: https://github.com/facebookresearch/synsin

[March 8th 2020]: PyTorch3D v0.1.1 bug fix release

[Jan 23rd 2020]: PyTorch3D v0.1.0 released. Mesh R-CNN codebase released: https://github.com/facebookresearch/meshrcnn

Description
PyTorch3D is FAIR's library of reusable components for deep learning with 3D data
Readme BSD-3-Clause 93 MiB
Languages
Python 80.9%
C++ 10.2%
Cuda 6.3%
C 0.9%
Shell 0.8%
Other 0.9%