Increase performance for conversions including axis angles

[alex-bene]

This is an extension of #1544 with various speed, stability, and readability improvements. (I could not find a way to make a commit to the existing PR). This PR is still based on the Rodrigues' rotation formula.

The motivation is the same; this change speeds up the conversions up to 10x, depending on the device, batch size, etc.

Notes

• As the angles get very close to π, the existing implementation and the proposed one start to differ. However, (my understanding is that) this is not a problem as the axis can not be stably inferred from the rotation matrix in this case in general.
• @bottler , I tried to follow similar conventions as existing functions to deal with weird angles, let me know if something needs to be changed to merge this.

[facebook-github-bot]

Hi @alex-bene!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[facebook-github-bot]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[alex-bene]

@bottler let's continue the conversation here.
Regarding your comment on #1544 , do you want to add the "fast" argument to both matrix->axis angle and axis angle-> matrix conversions? Also, all the rest of the conversions (e.g. quaternion_to_axis_angle) use a fixed eps=1e-6. While I do agree this should be an input argument, this is different from the existing convention. Is it okay to add this argument to all functions that use it?

[bottler]

Yes let's add the "fast" to both cases. And I don't mind the eps being inconsistent: actually feel free to add it everywhere or in these two places or nowhere.

[alex-bene]

Okay, also, one final question @bottler . On the implementation of quaternion_to_axis_angle and axis_angle_to_quaternion I can see that you take a 2nd-order Taylor approximation for the calculation of torch.sin(x)/x. Is this better than using the torch-native torch.sinc(x/torch.pi)? If not, we could simplify these functions a lot, e.g.:

def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = angles * 0.5
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions

could become

def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    sin_half_angles_over_angles = 0.5 * torch.sinc(0.5 * angles / torch.pi)
    return torch.cat(
        [torch.cos(0.5 * angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )

and this also avoids setting the eps in general.

From my testing, there does not seem to be a downside to using torch.sinc as seen here:

import math

for aprx in range(5):
    sinx_div_x_taylor = lambda x: sum([
        ((-1 if n%2 == 1 else 1) / math.factorial(n)) * (x**(2*n)) for n in range(aprx)
    ])

    res = torch.tensor([
        sinx_div_x_taylor(i) - torch.sinc(torch.tensor(i)/torch.pi).item()
        for i in np.linspace(0, 1e-6, 1000)
    ])
    print(f"taylor approximation: {aprx} - max error: {res.abs().max()} - mean error: {res.abs().mean()}")

which outputs:

taylor approximation: 0 - max error: 1.0 - mean error: 1.0
taylor approximation: 1 - max error: 1.66644475996236e-13 - mean error: 5.558120630411167e-14
taylor approximation: 2 - max error: 8.333334022836425e-13 - mean error: 2.779195762414588e-13
taylor approximation: 3 - max error: 8.333334022836425e-13 - mean error: 2.779195762414588e-13
taylor approximation: 4 - max error: 8.333334022836425e-13 - mean error: 2.779195762414588e-13

[bottler]

Yes please! I think the change to torch.sinc would be a good idea. It wasn't tried before. Thanks!

[bottler]

Prefer torch.special.sinc to torch.sinc .

[alex-bene]

torch.special.sinc is the same as torch.sinc as far as I understand.
Also, FYI, sinc generates the exact same results as a naive implementation of sinx/x:

import torch

naive = lambda x: torch.sin(x) / x if x != 0 else 1

res = torch.tensor([
    naive(i) - torch.sinc(i/torch.pi).item()
    for i in torch.linspace(0, 1e-6, 100000)
])

assert res.abs().max() == 0.0
assert res.abs().mean() == 0.0

[bottler]

So the whole taylor series thing probably wasn't useful at all!

[alex-bene]

So the whole taylor series thing probably wasn't useful at all!

It seems so

So, I spent quite a lot of time to figure out what to do in the conversion matrix_to_axis_angle when the angle is near pi or -pi. Eventually, it seems like the solution is to do an intermediate conversion to quaternions.

In total, I feel quite confident to set fast=True in general, but take a look and let me know if there's anything else that seems like it could create a problem.

[bottler]

So, I spent quite a lot of time to figure out what to do in the conversion matrix_to_axis_angle when the angle is near pi or -pi. Eventually, it seems like the solution is to do an intermediate conversion to quaternions.

In which case, is this change to that function actually a speed improvement? I suggest either not bothering, or, I guess, doing something less precise and acknowledging that in the docstring. I think the thing this PR is really trying to do is speed up a case where you know have a simple speedup(s) with an approach already thought to be accurate enough for your application. That might only be true in one direction.

[alex-bene]

In which case, is this change to that function actually a speed improvement

Well, considering a uniform distribution of angles, this should still be a speed improvement (unless all angles were near pi for example).

However, calling the functions for the intermediate conversion even with empty tables led to very slow performance (not sure why). So, I just committed a change to avoid this. Now there is a total improvement in speed across all cases plus the function is stable in all edge cases.

[facebook-github-bot]

@bottler has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[alex-bene]

Hey @bottler , I can see that some tests are falling, however, I do not have access to see exactly what fails. Is it possible to share the error logs to fix this?

[facebook-github-bot]

@bottler merged this pull request in 7a3c0cb.

[bottler]

Thanks for this work! Tests are fine. Probably you saw some reflection of longstanding failures.

[alex-bene]

Okay, I see. Glad I could help. 🙌