Thanks for the answer!



 albanD:

The short answer is that you don’t need to worry, and if all the parameters of self.net do require gradients, you can simply do result = torch.utils.checkpoint.checkpoint(self.net, x, y)[0] and use loss.backward() and all the parameters inside self.net will have their .grad field populated.


I tried that first, and i get
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

My thinking was that checkpoint removes all requires_grad by design and returns them for the given parameters and therefore it’s not working.



 albanD:

The longer answer is that the whole limitation about not supporting autograd.grad 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.


I’m sorry I don’t understand this. autograd.grad 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?).