Extract gradients of a model like model.state_dict()

autograd
1

Hi, I understand that to extract a model’s parameters’ state can be done through
`state = model.state_dict()’. Is it possible to get the model’s gradients with something like ‘gradients = model.gradient_dict()’ so that I can do ‘model.load_gradient_dict(gradients)’ back to the model? Any help will be much appreciated!

2

The following snippet allows you to get a sort of gradient_dict:

import torch
net = torch.nn.Linear(2, 3)
x = torch.rand(4, 2).requires_grad_(True)
loss = net(x).sum() ** 2
loss.backward()
grad_dict = {k:v.grad for k, v in zip(net.state_dict(), net.parameters())}
print(grad_dict)
1 Like
3

Just want to mention that one should be cautious when using this line

grad_dict = {k:v.grad for k, v in zip(net.state_dict(), net.parameters())}

as net.state_dict() 's keys is not guranteed to be in the same order as in net.parameters(), as I tried this in my SSD model. So I would sugget to use:

grad_dict = {k:v.grad for k,v in net.named_parameters()}

(Maybe in that time(2 years ago), we don’t have this issue or named_parameters is not available, I’m not sure, but I think it’s better to use named_parameters nowadays)

1 Like