Inplace Operation error with my Forward Kinematic function

kamat_mayur: <pre><code class="lang-auto">class DeepR_v1(nn.Module):
    def __init__(self, input_features, output_features, rest_pose, parent_indices, device):
        super(DeepR_v1, self).__init__()
        self.input_features = input_features
        self.output_features = output_features
        self.rest_pose = rest_pose
        self.parent_indices = parent_indices
        self.device = device
        
        self.converter = nn.Sequential(
            nn.Linear(input_features, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),  # ReLU activation
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),  # ReLU activation
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),  # ReLU activation
            nn.Linear(128, output_features),
            nn.Tanh()  # Tanh activation
)

    def axis_angle_to_quaternion(self, axis_angle: torch.Tensor) -&gt; torch.Tensor:...

    def quaternion_to_matrix(self, quaternions: torch.Tensor) -&gt; torch.Tensor:...
      
    def axis_angle_to_matrix(self, axis_angle: torch.Tensor) -&gt; torch.Tensor:...
    
    def make_4x4_transforms(self, rot, pelv_pos):...
      
    def fk(self, rest_rel_local_transforms):
        """
        Compute the global transforms for multiple frames given the rest-relative local transforms, 
        rest pose, and parent indices for each bone.

        Args:
        rest_rel_local_transforms (torch.Tensor): The rest-relative local transforms with shape (num_frames, num_bones, 4, 4).
        rest_pose (torch.Tensor): The rest pose transform with shape (num_bones, 4, 4).
        parent_indices (torch.Tensor): The parent indices for each bone with shape (num_bones).

        Returns:
        torch.Tensor: The global transforms with shape (num_frames, num_bones, 4, 4).
        """

        # Get the number of frames and bones from the shape of the input transforms
        num_frames, num_bones, _, _ = rest_rel_local_transforms.shape
        
        # Initialize the global transforms tensor with the same shape as the input transforms
        global_transforms = torch.zeros_like(rest_rel_local_transforms)

        # Compute the local transforms for all frames by multiplying the rest pose with the rest-relative local transforms
        local_transforms = self.rest_pose.unsqueeze(0).repeat(num_frames, 1, 1, 1) @ rest_rel_local_transforms

        # Initialize the global transform for the first bone (assuming it has no parent)
        global_transforms[:, 0] = local_transforms[:, 0]  # Assuming the first bone has no parent (parent_indices[0] == -1)
        
        # Use a loop to compute global transforms for the remaining bones for all frames
        for bone_idx in range(1, num_bones):
            # Get the parent index for the current bone
            parent_idx = self.parent_indices[bone_idx]
            
            # Compute the global transform for the current bone by multiplying the parent's global transform with the current local transform
            t = global_transforms[:, parent_idx] @ local_transforms[:, bone_idx]
            global_transforms[:, bone_idx] = t



        return global_transforms

    def forward(self, x):
        y = self.converter(x)

        r = y[:, :-3]
        rot = r.reshape(r.shape[0], r.shape[1]//3, 3)
        pelv_pos = y[:, -3:]

        r_mat = self.axis_angle_to_matrix(rot)

        rest_rel_local_transforms = self.make_4x4_transforms(r_mat, pelv_pos).to(self.device)

        global_transforms = self.fk(rest_rel_local_transforms).to(self.device)

        pos = global_transforms[:, :, :3, 3]

        return rot, pelv_pos, pos
</code></pre>
<p>when I train this model I get a runtime error:<br>
<code>RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [32, 4, 4]], which is output 0 of AsStridedBackward0, is at version 26; expected version 25 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!</code></p>
<p>using torch.autograd.set_detect_anomaly(True) prints the following:<br>
<code>File "C:\Users\mayur\AppData\Local\Temp\ipykernel_7976\2772885769.py", line 168, in fk t = global_transforms[:, parent_idx] @ local_transforms[:, bone_idx] (Triggered internally at [C:\cb\pytorch_1000000000000\work\torch\csrc\autograd\python_anomaly_mode.cpp:116](file:///C:/cb/pytorch_1000000000000/work/torch/csrc/autograd/python_anomaly_mode.cpp:116).) return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass.</code><br>
Why is this happening?</p>

kamat_mayur: <p>the error seems to go away if i change <code>global_transforms[:, bone_idx] = t</code> to <code>global_transforms[:, bone_idx] = t.detach()</code> but then the loss never changes during backprop so i’m guessing that is’nt the right thing to do. I am clueless why this is happening.</p>

KFrank: <p>Hi Mayur!</p>
<aside class="quote no-group quote-modified" data-username="kamat_mayur" data-post="1" data-topic="207177" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kamat_mayur/48/71793_2.png" class="avatar"> kamat_mayur:</div>
<blockquote>
<pre><code class="lang-auto">    def fk(self, rest_rel_local_transforms):
        ...
        global_transforms = torch.zeros_like(rest_rel_local_transforms)
        ...
        for bone_idx in range(1, num_bones):
            ...
            t = global_transforms[:, parent_idx] @ local_transforms[:, bone_idx]
            global_transforms[:, bone_idx] = t
</code></pre>
<p>…<br>
<code>RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [32, 4, 4]], which is output 0 of AsStridedBackward0, is at version 26; expected version 25 instead.</code><br>
…<br>
<code>File "C:\Users\mayur\AppData\Local\Temp\ipykernel_7976\2772885769.py", line 168, in fk t = global_transforms[:, parent_idx] @ local_transforms[:, bone_idx]</code></p>
</blockquote>
</aside>
<p>This all points to:</p>
<pre data-code-wrap="python"><code class="lang-python">global_transforms[:, bone_idx] = t
</code></pre>
<p>being the inplace modification that causes your error.  (Writing into a<br>
tensor using indexing counts as an inplace modification.)  You could<br>
check this by printing out the <code>.shape</code> and <code>._version</code> of<br>
<code>global_transforms</code> and see if they match those reported in the error<br>
message.</p>
<p>You could probably fix this by accumulating the updates to<br>
<code>global_transforms</code> in a separate tensors and then constructing<br>
<code>global_transforms</code> after the loop, but doing so would requires<br>
modifying the logic of your loop.</p>
<p>Pytorch’s <a href="https://pytorch.org/docs/2.4/autograd.html#torch.autograd.graph.allow_mutation_on_saved_tensors" rel="noopener nofollow ugc">sweep-inplace-modification-errors-under-the-rug context manager</a><br>
might also fix this for you.</p>
<p>Best.</p>
<p>K. Frank</p>

KFrank: <p>Hi Mayur!</p>
<aside class="quote no-group" data-username="kamat_mayur" data-post="2" data-topic="207177" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kamat_mayur/48/71793_2.png" class="avatar"> kamat_mayur:</div>
<blockquote>
<p><code>global_transforms[:, bone_idx] = t.detach()</code></p>
</blockquote>
</aside>
<p>Yes, this “breaks the computation graph” so you don’t backpropagate<br>
to the point of the inplace modification.</p>
<aside class="quote no-group">
<blockquote>
<p>but then the loss never changes during backprop</p>
</blockquote>
</aside>
<p>Yes, because of the broken computation graph, gradients don’t flow<br>
back to your model parameters, so no actual training occurs.</p>
<p>Best.</p>
<p>K. Frank</p>

kamat_mayur: <p>Ohh I didn’t know about that. I will try it right away. Thanks for your time <a class="mention" href="/u/kfrank">@KFrank</a> . I was literally here the whole time you were typing your response haha. Have a great day!</p>

kamat_mayur: <aside class="quote no-group" data-username="KFrank" data-post="3" data-topic="207177">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Pytorch’s <a href="https://pytorch.org/docs/2.4/autograd.html#torch.autograd.graph.allow_mutation_on_saved_tensors" rel="noopener nofollow ugc">sweep-inplace-modification-errors-under-the-rug context manager </a><br>
might also fix this for you.</p>
</blockquote>
</aside>
<p>So I have used this in my training loop as given below, and it runs fine without any error. The loss decreases as well. I would just like to know if this is the right way to use it…</p>
<pre><code class="lang-auto">for epoch in range(start_epoch, num_epochs):
        model.train()
        running_loss = 0.0
        
        for pos, rot, pelv_pos in dataloader:

            """
            Allowing mutation of tensors since model.fk() uses inplace
            operations being applied to both forward and backward pass 
            as mentioned in the docs
            """
            with torch.autograd.graph.allow_mutation_on_saved_tensors():
                # Send values to device
                pos = pos.to(device)
                rot = rot.to(device)
                pelv_pos = pelv_pos.to(device)

                # Estimate rotations of all bones + pelvis position
                est_rot, est_pelv_pos, est_pos = model(pos)
                est_pos = est_pos[:, 1:, :] # Removed the root position from predictions as it's not contained in ground truth and is (0,0,0) anyway.

                pos = pos.reshape(pos.shape[0], pos.shape[1]//3, 3)
                loss = criterion(est_pos, pos)

                loss.backward()

            optimizer.step()
            scheduler.step()

            # add to the running loss
            running_loss += loss.item()


        epoch_loss = running_loss / len(dataloader)
        print(f'Epoch {epoch}/{num_epochs - 1}, Loss: {epoch_loss:.4f}')
</code></pre>