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

1

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:

2

Hi,

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()?

4

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.

5

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?