Checkpoint for a whole subnet

adamce: <p>Hi,<br>
tl;dr: is it possible to make a checkpoint for a whole sub-network and learn its weights with a nice api?</p>
<p>I’ve learned that i can replace</p>
<pre><code class="lang-auto">result = self.net(x, y)[0]
</code></pre>
<p>by</p>
<pre><code class="lang-auto">result = torch.utils.checkpoint.checkpoint(self.net, x, y)[0]
</code></pre>
<p>where in my case self. net is a sub-model with some logic and several layers, and x and y are inputs from some earlier layer. As far as I understand, x and y will be stored during forward, but no other value inside self. net. On the backward pass, the gradient coming from result will see the checkpoint and recalculate everything to get the gradient for x and y. This trades computation for memory.</p>
<p>Previously i didn’t need the gradients in self. net (it was pre-learned), but now I do. So I would like to have something like</p>
<pre><code class="lang-auto">result = torch.utils.checkpoint.checkpoint(self.net, self.net.parameters(), x, y)[0]
</code></pre>
<p>As far as I understand, technically this is very similar to the first case, but the API doesn’t allow it as self. net. parameters() is an iterator.</p>
<p>I’m not very experienced with python and pytorch (only started last summer), but I managed to do the following (and it seems to work):</p>
<pre><code class="lang-auto">class _NetCheckpointWrapper:
    def __init__(self, net, x, y):
        self.net = net
        self.x = x
        self.y = y

    def __call__(self, *args, **kwargs):
        return self.net(self.x, self.bias)


# forward:
wrapper = _NetCheckpointWrapper(self.net, x, y)
net_params = tuple(self.net.parameters())
result = torch.utils.checkpoint.checkpoint(wrapper, *net_params)[0]
</code></pre>
<p>This doesn’t look very good to me as a programmer, but maybe it is idiomatic in python/pytorch? I’m not sure about the usage of net_params – coming from c++ i’m afraid net_params might refer to different objects than self. net. parameters(), or that it might change in the future. It is quite verbose as well.</p>
<p>Is there a better way to achieve that?</p>

albanD: <p>Hi,</p>
<p>The short answer is that you don’t need to worry, and if all the parameters of <code>self.net</code> do require gradients, you can simply do  <code>result = torch.utils.checkpoint.checkpoint(self.net, x, y)[0]</code> and use <code>loss.backward()</code> and all the parameters inside <code>self.net</code> will have their <code>.grad</code> field populated.</p>
<p>The longer answer is that the whole limitation about not supporting <code>autograd.grad</code> with checkpoint is to make this simple construct possible. You don’t need to pass all the things that require gradients as input to the checkpoint. Only the inputs to your network.</p>

adamce: <p>Thanks for the answer!</p>
<aside class="quote no-group" data-username="albanD" data-post="2" data-topic="65295">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>The short answer is that you don’t need to worry, and if all the parameters of <code>self.net</code> do require gradients, you can simply do <code>result = torch.utils.checkpoint.checkpoint(self.net, x, y)[0]</code> and use <code>loss.backward()</code> and all the parameters inside <code>self.net</code> will have their <code>.grad</code> field populated.</p>
</blockquote>
</aside>
<p>I tried that first, and i get</p>
<pre><code class="lang-auto">checkpoint.py:25: UserWarning: None of the inputs have requires_grad=True. Gradients will be None
  warnings.warn("None of the inputs have requires_grad=True. Gradients will be None")
Traceback (most recent call last):
 ..
  File "..", line 249, in train_on
    loss.backward()
  File "/home/madam/bin/anaconda3/lib/python3.7/site-packages/torch/tensor.py", line 166, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/madam/bin/anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>
<p>My thinking was that checkpoint removes all requires_grad by design and returns them for the given parameters and therefore it’s not working.</p>
<aside class="quote no-group" data-username="albanD" data-post="2" data-topic="65295">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>The longer answer is that the whole limitation about not supporting <code>autograd.grad</code> with checkpoint is to make this simple construct possible. You don’t need to pass all the things that require gradients as input to the checkpoint. Only the inputs to your network.</p>
</blockquote>
</aside>
<p>I’m sorry I don’t understand this. <code>autograd.grad</code> is does forward differentiation, doesn’t it? my understanding is that i’m doing backward diff only (but maybe it internally inverts the computation tree?).</p>

albanD: <blockquote>
<p>I’m sorry I don’t understand this.  <code>autograd.grad</code>  is does forward differentiation, doesn’t it?</p>
</blockquote>
<p>No it does backward mode as well. Only that instead of populating the <code>.grad</code> field of all the leaf Tensor, it computes and returns the gradients for a list of Tensor that is given.</p>
<blockquote>
<p>My thinking was that checkpoint removes all requires_grad by design</p>
</blockquote>
<p>No the output of the checkpoint module should require gradients. If you put the whole network into it, you might need to set <code>requires_grad=True</code> on the input though. Does it help?</p>

adamce: <aside class="quote no-group" data-username="albanD" data-post="4" data-topic="65295">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>If you put the whole network into it, you might need to set <code>requires_grad=True</code> on the input though. Does it help?</p>
</blockquote>
</aside>
<p>Thanks. Yes, it does help. Although I’m not totally happy with the result either (setting requires_grad=true to a tensor that doesn’t need it):</p>
<pre><code class="lang-auto">if y.requires_grad is False:
   y.requires_grad_(True)  # require gradient for y even when learning net, otherwise checkpoint doesn't work
   result = torch.utils.checkpoint.checkpoint(self.net, x, y)
</code></pre>

albanD: <p>Unfortunately, it is a limitation of the checkpoint module. At least one input to the module needs to require gradients.</p>
<p>Since usually, the first module of your net has some weights that require gradients, it should not have any impact on the graph that is created.</p>