How can I get model's gradients in two parts?

I have a neural network divided into two parts. One part computes the first n layers and the other part computes the rest of the layers.

I can obtain the gradients from a single model this way:

grads = torch.autograd.grad(loss, model.parameters())

The thing is I need the gradients of the two parts separately. The main issue I am facing seems to be continuing the first part’s gradients backwards from the second part.

I would appreciate any help as to how I can achieve this. Thanks :slight_smile:


I am not sure to understand the issue. Can you extract the corresponding gradients from grads?
Or provide only the subset of the parameters you are interested in to autograd.grad()?

Hi, thanks for the quick reply.

Say I have a single model and want to obtain gradient values from it. I can do it as I showed in main post general.

Now I have the model as two nn.Module instances, one computing the first part and one computing the second part.

I want to obtain the gradient values from both parts and combine them. Bit of a weird problem but should be possible in some way.

If your full model has model.part1 and model.part2, then you can do:

grad_part1 = torch.autograd.grad(loss, model.part1.parameters(), retain_graph=True)
grad_part2 = torch.autograd.grad(loss, model.part2.parameters())

Does that match what you want?