Why is requires.grad not saved in model.state_dict()?

Hello!

I found a strange quirk in the following code:

import torchvision
model = torchvision.models.vgg19(pretrained=True)

By default all parameters in model have requires.grad set to True. Here we set features.0.weight and features.0.bias to True and all other parameters to False.

learnable_params = ['features.0.weight', 'features.0.bias']
for name, param in model.named_parameters():
    param.requires_grad = name in learnable_params
    
    if int(name.split('.')[1]) < 3 and name.startswith('features'):
        print(name, param.requires_grad)

Here are the print statements from the first 2 feature layers.

features.0.weight True
features.0.bias True
features.2.weight False
features.2.bias False

As we expected, features.0.weight and features.0.bias are True and all other parameters are False.

But then I investigated the model.state_dict():

for name, v in model.state_dict().items():
    if int(name.split('.')[1]) < 3 and name.startswith('features'):
        print(name, v.requires_grad)

which outputted

features.0.weight False
features.0.bias False
features.2.weight False
features.2.bias False

Now, features.0.weight and features.0.bias are False.

Why?

When I updated param.requires_grad did this note save into the model.state_dict()? Why is this happening?

Any suggestions would be appreciated. Thanks!

Use keep_vars=True as described in the docs and the parameters won’t be detached.

Wow, thank you so much!

This perfectly fixed my problem.

The code

for name, v in model.state_dict(keep_vars=True).items():
    if int(name.split('.')[1]) < 3 and name.startswith('features'):
        print(name, v.requires_grad)

outputs

features.0.weight True
features.0.bias True
features.2.weight False
features.2.bias False

as expected.

ptrblck you are the MVP! Thanks!

1 Like