Post-processing model output while retaining gradients

username: louisbzk
cooked: <p>Hello,</p>
<p>I am trying to train a model in the following fashion:</p>
<pre><code class="lang-auto"># in a training loop, getting a batch x and a target y_true
output = model(x)
output_processed = process_output(output)
loss = criterion(output_processed, y_true)
loss.backward()
optimizer.step()
</code></pre>
<p>I am facing issues implementing the <code>process_output</code> function. Some precisions on the context:</p>
<ul>
<li>my <code>output</code> tensor is of shape <code>(batch_size, output_size)</code>
</li>
<li>my <code>output_processed</code> tensor is of shape <code>(batch_size, some_other_size)</code> where <code>some_other_size &gt; output_size</code>
</li>
<li>the internals of process_output use slicing to assign values, for instance <code>output_processed[:, (0, 1, 2)] = output[:, (0, 1, 2)] + output[:, (3, 4, 5)]</code> (operations may be more complex but they all consist in torch elementary operations + slicing &amp; broadcasting tricks)</li>
</ul>
<p>Now, here is the issue I’m facing. I want to update the model’s weights using the loss computed from <code>processed_output</code>, so <code>process_output</code> begins like this:</p>
<pre><code class="lang-auto">def process_output(output: Tensor) -&gt; Tensor:
    output_processed = (
        torch.empty(size=(output.shape[0], some_other_size))
    ).to(output.device, output.dtype).requires_grad_(True)

    # perform operations: this raises RuntimeError
    output_processed[:, (0, 1, 2)] = output[:, (0, 1, 2)] * 2
    # ... other operations of the sort ...

    return output_processed
</code></pre>
<p>As commented in the function above, the first operation fails and raises <code>RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.</code> I am not sure how to proceed here because I don’t think I understand <code>requires_grad</code> and leaf variables correctly.</p>
<p>From my understanding, setting <code>requires_grad_(False)</code> when instantiating the tensor should “break” my gradient computation, that is calling <code>loss.backward()</code> will result in a wrong update of the model’s weights because I’m not “tracking” the gradients in <code>output_processed</code>. However, a user replied on a similar post that he managed to correctly update his model: <a href="https://discuss.pytorch.org/t/leaf-variable-was-used-in-an-inplace-operation/308/14" class="inline-onebox">Leaf variable was used in an inplace operation - #14 by nima_rafiee</a>.</p>
<p>The more I read on leaf variables, the more I’m confused. The docs for <a href="https://pytorch.org/docs/stable/generated/torch.Tensor.is_leaf.html" rel="noopener nofollow ugc">torch.Tensor.is_leaf</a> say that “Tensors that have <code>requires_grad</code> which is <code>False</code> will be leaf Tensors”, but then say “Only leaf Tensors will have their <code>grad</code> populated during a call to <code>backward()</code>”. So that means that you say “that Tensor does not track gradients”, but then that Tensor is a leaf Tensor so it has its gradient populated ?</p>
<p>So, should I set <code>requires_grad_(True)</code> for <code>output_processed</code> or not ? If yes, would creating different Tensors from the operations in <code>process_output</code> and then concatenating them at the end, rather than slicing an already instantiated Tensor, be a good solution to deal with the error ? Thanks !</p>

username: KFrank
cooked: <p>Hi Louis!</p>
<aside class="quote no-group" data-username="louisbzk" data-post="1" data-topic="181602" 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/louisbzk/48/61099_2.png" class="avatar"> louisbzk:</div>
<blockquote>
<p>I am facing issues implementing the <code>process_output</code> function.</p>
</blockquote>
</aside>
<p>Just implement <code>process_output()</code> using (differentiable) pytorch tensor<br>
functions and everything should work automatically.  (Note, slicing counts<br>
as a differentiable tensor function.)</p>
<p>As far as autograd and updating a model’s weights are concerned, there<br>
is nothing different about tensor functions that occur inside of a model,<br>
and and tensor functions that occur outside of a model such as those in<br>
a post-processing step.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; param = torch.ones (2, 3, requires_grad = True)   # some parameter in your model
&gt;&gt;&gt; output = 2 * param                                # dummy output of model
&gt;&gt;&gt; output_processed = torch.zeros (2, 4)
&gt;&gt;&gt; output_processed                                  # no connection to gradients yet
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.]])
&gt;&gt;&gt; output_processed.requires_grad                    # starts out False
False
&gt;&gt;&gt; output_processed[:, 1] = output[:, 0]             # use slice assignment
&gt;&gt;&gt; output_processed                                  # now carries a grad_fn so you can backpropagate
tensor([[0., 2., 0., 0.],
        [0., 2., 0., 0.]], grad_fn=&lt;CopySlices&gt;)
&gt;&gt;&gt; output_processed.requires_grad                    # "inherited" True from param -- wasn't set explicitly
True
&gt;&gt;&gt; loss = output_processed.sum()                     # some dummy loss
&gt;&gt;&gt; loss.backward()                                   # gradients backpropagate through output_processed
&gt;&gt;&gt; param.grad                                        # to param.grad
tensor([[2., 0., 0.],
        [2., 0., 0.]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: louisbzk
cooked: <p>Hi Frank ! Thank you very much for your answer.</p>
<p>I noticed while debugging that indeed, even when explicitly specifying <code>requires_grad_(False)</code> my tensor inherited the <code>requires_grad</code> property of my model output, but I didn’t see the gradients propagating explicitly like you showed here. You saved me quite a bit of refactoring <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>
<p>Best,<br>
Louis</p>