Unexpected behaviour for affine_grid and grid_sample with 3D inputs

EmilienDupont: <p>Hi, I posted this issue on <a href="https://github.com/pytorch/pytorch/issues/35775#issuecomment-607251114" rel="nofollow noopener">github</a> and was told to ask about it on here. Thanks for your help!</p>
<h2>
<img src="https://discuss.pytorch.org/images/emoji/apple/bug.png?v=9" title=":bug:" class="emoji" alt=":bug:"> Bug</h2>
<p>When performing affine transformations on 3D tensors I’ve noticed some unexpected behaviour which does not occur for 2D tensors.</p>
<p>Specifically, suppose we perform two 3D rotations with two rotation matrices <code>A</code> and <code>B</code>. Given a 3D input tensor <code>x</code>, we can first rotate it with <code>B</code> to get the rotated tensor <code>x_B</code>. We can then further rotate <code>x_B</code> by <code>A</code> to get the rotation corresponding to <code>B</code> followed by <code>A</code>, i.e. to the rotation matrix <code>AB</code>. We can denote this resulting tensor by <code>x_AB</code>.</p>
<p>Alternatively, we could have directly taken the matrix product <code>C = AB</code> and transformed the tensor with <code>C</code>, to obtain <code>x_C</code>. Now <code>x_AB</code> and <code>x_C</code> should be approximately equal up to sampling errors. However in practice this is not the case. As an example:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://user-images.githubusercontent.com/12076034/78072507-29fe2980-7397-11ea-96c5-3157040f239b.png" title="bug-3d-rotation" rel="nofollow noopener"><img src="https://user-images.githubusercontent.com/12076034/78072507-29fe2980-7397-11ea-96c5-3157040f239b.png" alt="bug-3d-rotation" width="690" height="206"><div class="meta">
<svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use xlink:href="#far-image"></use></svg><span class="filename">bug-3d-rotation</span><span class="informations">1072×321</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use xlink:href="#discourse-expand"></use></svg>
</div></a></div></p>
<p>Clearly, the cube rotated by <code>B</code> followed by <code>A</code> is not the same as the one rotated by <code>C=AB</code>. Note that this does not happen in 2D. If we perform similar transformations to a 2D input, we obtain:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://user-images.githubusercontent.com/12076034/78072540-35e9eb80-7397-11ea-9bde-a90ae241c999.png" title="bug-3d-rotation-2d-example" rel="nofollow noopener"><img src="https://user-images.githubusercontent.com/12076034/78072540-35e9eb80-7397-11ea-9bde-a90ae241c999.png" alt="bug-3d-rotation-2d-example" width="690" height="217"><div class="meta">
<svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use xlink:href="#far-image"></use></svg><span class="filename">bug-3d-rotation-2d-example</span><span class="informations">1040×328</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use xlink:href="#discourse-expand"></use></svg>
</div></a></div></p>
<p>As can be seen, the transformations by <code>AB</code> and <code>C</code> are roughly the same (<code>AB</code> has more sampling noise since it corresponds to two affine transformations and two grid_samples), so the error does not occur in 2D.</p>
<p>Any idea why the difference in 3D might be occurring? Thank you!</p>
<h2>To Reproduce</h2>
<pre><code class="lang-python">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_matrixA @ rotation_matrixB

# 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
</code></pre>
<h2>Expected behavior</h2>
<p>Rotating a tensor by <code>B</code> followed by <code>A</code> should be approximately the same as rotating by <code>AB</code>.</p>
<h2>Environment</h2>
<ul>
<li>PyTorch Version (e.g., 1.0): 1.4.0</li>
<li>OS (e.g., Linux): macOS Catalina</li>
<li>How you installed PyTorch (<code>conda</code>, <code>pip</code>, source): pip</li>
<li>Build command you used (if compiling from source): N/A</li>
<li>Python version: 3.6.8</li>
<li>CUDA/cuDNN version: N/A</li>
<li>GPU models and configuration: N/A</li>
<li>Any other relevant information: N/A</li>
</ul>

EmilienDupont: <p>I’ve just tested this on GPU and the CUDA function also has the exact same problem.</p>
<p>This is the setup:</p>
<ul>
<li>PyTorch Version (e.g., 1.0): 1.3.1</li>
<li>OS (e.g., Linux): Linux</li>
<li>How you installed PyTorch (conda, pip, source): conda</li>
<li>Build command you used (if compiling from source): N/A</li>
<li>Python version: 3.6.9</li>
<li>CUDA/cuDNN version: 9.1</li>
<li>GPU models and configuration: Tesla V100</li>
<li>Any other relevant information: N/A</li>
</ul>

MicPie: <p>Hi Emilien,</p>
<p>I am currently also playing with <code>affine_grid</code> and <code>grid_sample</code> and found your post via the forum search.</p>
<p>When I rotate 3D data in my setup I get a rotated <em><strong>and sheared</strong></em> output.</p>
<p>So far I am still trying to find the root cause but I expect that something unexpected happens in <code>grid_sample</code> I am not taking into account (currently I’m looking into the implications of this <a href="https://github.com/pytorch/pytorch/issues/24870#issuecomment-573701812" rel="nofollow noopener">github comment</a>, but I still have to think about it).</p>
<p>Have you tried to see if you rotation is preserving edge lengths of your cube, i.e., it only rotates and does not shear too?  (I guess a small shear can be quite tricky to see in the 3D visualizations.)</p>
<p>Maybe a more basic setup where you only rotate in one dimension can help you pinpoint the problem better?</p>
<p>Kind regards<br>
Michael</p>

MicPie: <p>I was able to track down my errors by carrying everything out in “world” coordinates and then creating the affine grid from them.</p>
<p>I had to do a lot of “visual” debugging with <a href="https://github.com/K3D-tools/K3D-jupyter" rel="nofollow noopener">https://github.com/K3D-tools/K3D-jupyter</a> which I used for rendering my point clouds and input/output data (really a great tool for that).</p>

EmilienDupont: <p>Hi Michael,</p>
<p>Thanks for your reply!</p>
<p>The problem ended up being that <code>grid_sample</code> performs an <strong>inverse</strong> warping, which means that passing an <code>affine_grid</code> 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 <code>grid_sample</code>.</p>
<p>The reason why it worked in 2D and not 3D is because 2D rotation matrices commute whereas 3D rotation matrices do not.</p>
<p>It looks like you basically found the same fix by using world coordinates!</p>
<p>Thanks again for your replies and good to hear you fixed your problems too.</p>