Failing to create a gradiented rotation matrix

username: Omroth
cooked: <p>I’m trying to find the rotation and translation which best maps one set of 3D points to another.  My code:</p>
<pre><code class="lang-auto">    reg_verts_tensor = torch.from_numpy(reg_verts).float().to(device)
    verts_tensor = torch.from_numpy(verts).float().to(device)

    tensor_0 = torch.zeros(1).float().to(device)
    tensor_1 = torch.ones(1).float().to(device)


    criterion = torch.nn.MSELoss(reduction='sum')

    n_register_epochs = 1000
    register_initial_learning_rate = 0.001

    register_learning_rate = register_initial_learning_rate
    optimizer = torch.optim.SGD([translation, rotation_x, rotation_y, rotation_z], lr=register_learning_rate)
    
    loss_list = []

    for epoch in range(n_register_epochs):
        optimizer.zero_grad()

        epoch_loss = None

        transformed_verts = verts_tensor + translation

        rotation_matrix_x = torch.stack([torch.stack([tensor_1,tensor_0,tensor_0]),
                        torch.stack([tensor_0,torch.cos(rotation_x), -torch.sin(rotation_x)]),
                        torch.stack([tensor_0, torch.sin(rotation_x), torch.cos(rotation_x)])]).reshape(3,3)

        rotation_matrix_y = torch.stack([torch.stack([torch.cos(rotation_y), tensor_0, - torch.sin(rotation_y)]),
                            torch.stack([tensor_0, tensor_1, tensor_0]),
                            torch.stack([torch.sin(rotation_y), tensor_0, torch.cos(rotation_y)])]).reshape(3,3)

        rotation_matrix_z = torch.stack([torch.stack([torch.cos(rotation_z), torch.sin(rotation_z), tensor_0]),
                            torch.stack([-torch.sin(rotation_z), torch.cos(rotation_z), tensor_0]),
                            torch.stack([tensor_0,tensor_0,tensor_1])]).reshape(3,3)

        for i in range(transformed_verts.shape[0]):
            transformed_verts[i] = torch.matmul(rotation_matrix_x, transformed_verts[i])
            transformed_verts[i] = torch.matmul(rotation_matrix_y, transformed_verts[i])
            transformed_verts[i] = torch.matmul(rotation_matrix_z, transformed_verts[i])

        #

        loss = criterion(transformed_verts, reg_verts_tensor)


        epoch_loss = loss

        if epoch % 10 == 0:
            print(epoch_loss.item())

        epoch_loss.backward()
</code></pre>
<p>This fails on the “backwards” call with:</p>
<p>RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [3]], which is output 0 of AsStridedBackward0, is at version 1323; expected version 1322 instead.</p>
<p>What am I doing wrong?</p>

username: Omroth
cooked: <p>The detector says it’s here:</p>
<p><code>transformed_verts[i] = torch.matmul(rotation_matrix_z, transformed_verts[i])</code></p>
<p>What’s the gradientable way of doing that?</p>

username: Omroth
cooked: <p>This works:</p>
<pre><code class="lang-auto">        lines = []

        for i in range(transformed_verts.shape[0]):
            res = torch.matmul(rotation_matrix_x, transformed_verts[i])
            res = torch.matmul(rotation_matrix_y, res)
            res = torch.matmul(rotation_matrix_z, res)

            lines.append(res)

        transformed_verts = torch.stack(lines)
</code></pre>

username: KFrank
cooked: <p>Hi Ian!</p>
<aside class="quote no-group" data-username="Omroth" data-post="1" data-topic="144483" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/o/fbc32d/48.png" class="avatar"> Omroth:</div>
<blockquote>
<p>RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation:</p>
</blockquote>
</aside>
<p>As you’ve noted in your subsequent post, at issue is that you are modifying<br>
<code>transformed_verts</code> “inplace” by writing into it through an index.  Depending<br>
on what else is going on in autograd’s “computation graph,” this can block<br>
backpropagation.</p>
<aside class="quote no-group" data-username="Omroth" data-post="3" data-topic="144483">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/o/fbc32d/48.png" class="avatar"> Omroth:</div>
<blockquote>
<p><code>transformed_verts = torch.stack(lines)</code></p>
</blockquote>
</aside>
<p>This solution works.</p>
<p>But why?</p>
<p>Remember that python “variables” are references to python objects in<br>
memory.</p>
<p>You start with a python object – specifically a pytorch tensor – in memory,<br>
and it is referred to by the python “variable” <code>transformed_verts</code>.</p>
<p>In your original, non-working version, you modify this specific pytorch<br>
tensor (and pytorch rightly complains).</p>
<p>In your second, working version, you <em>don’t</em> modify the original pytorch<br>
tensor.  Instead you create a new pytorch tensor by calling <code>torch.stack()</code>.<br>
You then assign it to the python variable <code>transformed_verts</code>.  But this<br>
merely makes the variable <code>transformed_verts</code> refer to the new tensor.<br>
The old – and, importantly, unmodified – tensor still remains in memory<br>
and is referred to by some variable or another in autograd’s computation<br>
graph.</p>
<p>Finally (I’m assuming that <code>verts_tensor</code>, and hence <code>transformed_verts</code>,<br>
are batches of of vertices with shape <code>[nBatch, 3]</code>.), you don’t need to use<br>
a python loop over the batch elements to multiply the vertices by your<br>
rotation matrix one by one – you can do it all at once (more efficiently)<br>
with a single call to <code>matmul()</code> that broadcasts over the batch dimension.<br>
But you need to use <code>unsqueeze()</code> (or similar reshaping) to convert<br>
<code>transformed_verts</code> to a batch of “column vectors” so that <code>matmul()</code><br>
understands what you want to multiply.</p>
<p>(It’s also probably marginally more efficient to pre-multiply your x, y, and<br>
z rotation matrices into a single rotation matrix and use that to rotate your<br>
vertices.)</p>
<p>Here is an illustration:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.10.2'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; transformed_verts = torch.randn (5, 3)
&gt;&gt;&gt;
&gt;&gt;&gt; rotation_matrix_x = torch.randn (3, 3, requires_grad = True)
&gt;&gt;&gt; rotation_matrix_y = torch.randn (3, 3, requires_grad = True)
&gt;&gt;&gt; rotation_matrix_z = torch.randn (3, 3, requires_grad = True)
&gt;&gt;&gt;
&gt;&gt;&gt; # version that loops over batch dimension
...
&gt;&gt;&gt; lines = []
&gt;&gt;&gt;
&gt;&gt;&gt; for i in range(transformed_verts.shape[0]):
...     res = torch.matmul(rotation_matrix_x, transformed_verts[i])
...     res = torch.matmul(rotation_matrix_y, res)
...     res = torch.matmul(rotation_matrix_z, res)
...
...     lines.append(res)
...
&gt;&gt;&gt; transformed_vertsA = torch.stack(lines)
&gt;&gt;&gt;
&gt;&gt;&gt; transformed_vertsA.sum().backward()
&gt;&gt;&gt;
&gt;&gt;&gt; gradA = rotation_matrix_x.grad
&gt;&gt;&gt;
&gt;&gt;&gt; rotation_matrix_x.grad = None
&gt;&gt;&gt; rotation_matrix_y.grad = None
&gt;&gt;&gt; rotation_matrix_z.grad = None
&gt;&gt;&gt;
&gt;&gt;&gt; # version that calls matmul() once (with a single pre-multiplied rotation matrix)
...
&gt;&gt;&gt; rotation_matrix = rotation_matrix_z @ rotation_matrix_y @ rotation_matrix_x
&gt;&gt;&gt; transformed_vertsB = torch.matmul (rotation_matrix, transformed_verts.unsqueeze (2)).squeeze()
&gt;&gt;&gt;
&gt;&gt;&gt; transformed_vertsB.sum().backward()
&gt;&gt;&gt;
&gt;&gt;&gt; gradB = rotation_matrix_x.grad
&gt;&gt;&gt;
&gt;&gt;&gt; transformed_vertsA
tensor([[ -0.6688,   0.7012,  -0.2841],
        [ -3.5321,   4.6106,  -3.5176],
        [  7.7354, -11.2853,   6.8841],
        [  0.8983,  -0.3446,  -0.3980],
        [ -0.5089,   0.5077,  -0.4176]], grad_fn=&lt;StackBackward0&gt;)
&gt;&gt;&gt; transformed_vertsB
tensor([[ -0.6688,   0.7012,  -0.2841],
        [ -3.5321,   4.6106,  -3.5176],
        [  7.7354, -11.2853,   6.8841],
        [  0.8983,  -0.3446,  -0.3980],
        [ -0.5089,   0.5077,  -0.4176]], grad_fn=&lt;SqueezeBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; gradA
tensor([[-3.6647,  1.2165,  0.3677],
        [-1.7096,  0.5675,  0.1715],
        [-4.3082,  1.4301,  0.4323]])
&gt;&gt;&gt; gradB
tensor([[-3.6647,  1.2165,  0.3677],
        [-1.7096,  0.5675,  0.1715],
        [-4.3082,  1.4301,  0.4323]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>