Unexpected behaviour for affine_grid and grid_sample with 3D inputs

[EmilienDupont]

🐛 Bug

When performing affine transformations on 3D tensors I’ve noticed some unexpected behaviour which does not occur for 2D tensors.

Specifically, suppose we perform two 3D rotations with two rotation matrices A and B. Given a 3D input tensor x, we can first rotate it with B to get the rotated tensor x_B. We can then further rotate x_B by A to get the rotation corresponding to B followed by A, i.e. to the rotation matrix AB. We can denote this resulting tensor by x_AB.

Alternatively, we could have directly taken the matrix product C = AB and transformed the tensor with C, to obtain x_C. Now x_AB and x_C should be approximately equal up to sampling errors. However in practice this is not the case. As an example:

Clearly, the cube rotated by B followed by A is not the same as the one rotated by C=AB. Note that this does not happen in 2D. If we perform similar transformations to a 2D input, we obtain:

As can be seen, the transformations by AB and C are roughly the same (AB has more sampling noise since it corresponds to two affine transformations and two grid_samples), so the error does not occur in 2D.

Any idea why the difference in 3D might be occurring? Thank you!

To Reproduce

import torch

def get_3d_cube(grid_size=32, cube_size=16):
    """Returns a 3d cube which is useful for debugging and visualizing 3d
    rotations.

    Args:
        grid_size (int): Size of sides of 3d grid.
        cube_size (int): Size of sides of cube within grid.
    """
    cube = torch.zeros(1, 1, grid_size, grid_size, grid_size)
    hs = grid_size // 2  # Half grid size
    hcs = cube_size // 2  # Half cube size

    cube[:, :, -hcs + hs:hs + hcs, hs + hcs, hs + hcs] = 1
    cube[:, :, -hcs + hs:hs + hcs, hs + hcs, hs - hcs] = 1
    cube[:, :, -hcs + hs:hs + hcs, hs - hcs, hs + hcs] = 1
    cube[:, :, -hcs + hs:hs + hcs, hs - hcs, hs - hcs] = 1

    cube[:, :, hs + hcs, -hcs + hs:hs + hcs, hs + hcs] = 1
    cube[:, :, hs - hcs, -hcs + hs:hs + hcs, hs + hcs] = 1
    cube[:, :, hs + hcs, -hcs + hs:hs + hcs, hs - hcs] = 1
    cube[:, :, hs - hcs, -hcs + hs:hs + hcs, hs - hcs] = 1

    cube[:, :, hs + hcs, hs + hcs, -hcs + hs:hs + hcs] = 1
    cube[:, :, hs - hcs, hs + hcs, -hcs + hs:hs + hcs] = 1
    cube[:, :, hs + hcs, hs - hcs, -hcs + hs:hs + hcs] = 1
    cube[:, :, hs - hcs, hs - hcs, -hcs + hs:hs + hcs] = 1
    
    return cube

# Create a cube to use for debugging
cube = get_cube_3d(128, 64)

# 3D rotation matrices with shape (1, 3, 3)
rotation_A = torch.tensor([[[ 0.7198, -0.6428, -0.2620],
                            [ 0.5266,  0.7516, -0.3973],
                            [ 0.4523,  0.1480,  0.8795]]])

rotation_B = torch.tensor([[[ 0.6428, 0.0000,  0.7660],
                            [ 0.0000,  1.0000,  0.0000],
                            [-0.7660,  0.0000,  0.6428]]])

# C = AB, so C is a rotation by B, followed by a rotation by A
rotation_C = rotation_A @ rotation_B

# Convert to affine matrices of shape (1, 3, 4), where the 4th column
# corresponds to a zero translation
translations = torch.zeros(1, 3, 1)
affine_A = torch.cat([rotation_A, translations], dim=2)
affine_B = torch.cat([rotation_B, translations], dim=2)
affine_C = torch.cat([rotation_C, translations], dim=2)

# Get grid for each affine transform
grid_A = torch.nn.functional.affine_grid(affine_A, size=cube.shape)
grid_B = torch.nn.functional.affine_grid(affine_B, size=cube.shape)
grid_C = torch.nn.functional.affine_grid(affine_C, size=cube.shape)

# Rotate cube by B
cube_B = torch.nn.functional.grid_sample(cube, grid_B, mode="bilinear")
# Rotate resulting cube by A, to obtain AB transformation
cube_AB = torch.nn.functional.grid_sample(cube_B, grid_A, mode="bilinear")
# Rotate original cube by C
cube_C = torch.nn.functional.grid_sample(cube, grid_C, mode="bilinear")

print((cube_AB - cube_C).abs().sum())  # This is large

Expected behavior

Rotating a tensor by B followed by A should be approximately the same as rotating by AB.

Environment

• PyTorch Version (e.g., 1.0): 1.4.0
• OS (e.g., Linux): macOS Catalina
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source): N/A
• Python version: 3.6.8
• CUDA/cuDNN version: N/A
• GPU models and configuration: N/A
• Any other relevant information: N/A

cc @albanD @mruberry

[ezyang]

I'm not an expert on affine transformations so I can't tell exactly what the error is, but I would guess that the math for the 3d case is wrong somewhere. Perhaps someone on the forums would know.

[EmilienDupont]

Thanks! I'll try to post on the forums.

[EmilienDupont]

I've just tested this on GPU and the CUDA function also has the exact same problem.

This is the setup:

• PyTorch Version (e.g., 1.0): 1.3.1
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): conda
• Build command you used (if compiling from source): N/A
• Python version: 3.6.9
• CUDA/cuDNN version: 9.1
• GPU models and configuration: Tesla V100
• Any other relevant information: N/A

[mruberry]

Thanks for reporting this issue, @EmilienDupont. We you able to find a resolution on the forums, or have you been able to identify what you think is going on?

[EmilienDupont]

Hi @mruberry,

The problem ended up being that grid_sample performs an inverse warping, which means that passing an affine_grid for the matrix A actually corresponds to the transformation A^(-1). So in my example above, the transformation with B followed by A actually corresponds to A^(-1)B^(-1) = (BA)^(-1), which means I should use C = BA and not C = AB as would be the case if the original matrix was used by grid_sample.

The reason why it worked in 2D and not 3D is because 2D rotation matrices commute whereas 3D rotation matrices do not.

In addition to this, it seems like the order of the coordinates is also inverted, i.e. for an (N, C, D, H, W) tensor, the (D, H, W) dimensions correspond to (z, y, x) instead of (x, y, z). I'm not sure exactly why this is.

There's also some discussion on the forums.

Thank you!

[mruberry]

Thanks for the update, @EmilienDupont. What you're saying makes a lot of sense. Would you be interested in helping improve our documentation so there's less confusion in the future?