Can I use "torch.autograd.grad" in "training_step" funtion?

[chenshuang-zhang]

Hello!

I hope to calculate the gradient of the loss on the model input every batch when training. The calculated gradients are then processed by some other functions and added into the loss function. The way I do this in pytorch is like this, with the torch.autograd.grad function:

model_input.requires_grad = True
model.zero_grad()
model.eval()
grads = torch.autograd.grad(loss, model_input, grad_outputs=None, only_inputs=True, retain_graph=False)[0]
model.train()

Can I add these codes directly into the training_step function in pytorch-lightning?

My concern is:

1. Will this gradient calculation (torch.autograd.grad function) influence the accuracy of model training since I add it in the training_step function?
2. Do I need to set model.zero_grad(), model.eval() and model.train() when calculating the gradients on the input?

Thank you very much!

[akihironitta]

1. Will this gradient calculation (torch.autograd.grad function) influence the accuracy of model training since I add it in the training_step function?

Yes, I believe. I don't think it'll work with retain_graph=False there because backward pass uses the graph to compute gradients wrt weights using the loss returned from training_step.

2. Do I need to set model.zero_grad(), model.eval() and model.train() when calculating the gradients on the input?

Yes, partially. I think you need zero_grad() after grads = torch.autograd.grad(...) to avoid accumulating gradients wrt weights, but I'm not sure why you need eval() and train().

[chenshuang-zhang]

Hey @YingYingFight,

When using DDP, each process will compute its own gradients based on the local data. Therefore, you can compute the local gradients but you need to be careful when updating the model weights, as this needs to be done using the global gradients = MEAN(all local gradients) through a all_reduce.

Hello @tchaton. Sorry I have another two questions about your comment.

(1) If I use only one GPU when training and no mixed precision training, can I use torch.autograd.grad in automatic optimization?

(2) Can I use model.zero_grad() in the training_step function in automatic optimization?

Thank you very much!

[tchaton]

(1) Yes.

(2) With Automatic Optimization, Lightning takes care of zero_grad. However, with manual, you can use model.zero_grad() too ;)

[chenshuang-zhang]

(1) Yes.

(2) With Automatic Optimization, Lightning takes care of zero_grad. However, with manual, you can use model.zero_grad() too ;)

Hi, @tchaton. It's so kind of you! I tried the manual optimization these days, and got what I want already!

To make sure I use correctly, can you help me with the following question?

With manual optimization, do I need to add some codes about mixed precision training？

Thank you very much!

[tchaton]

For, manual optimization, self.manual_backward takes care of precision :)

[chenshuang-zhang]

For, manual optimization, self.manual_backward takes care of precision :)

Cool! Thank you!