Adding new nn.Parameter in Conv2d (that gets saved in state_dict)

1

Conv2d contains two parameters that get saved in the state_dict: weight and bias. I want to add a new parameter of the same shape as bias (purpose is irrelevant here) that also gets saved in the state_dict. Currently, I am saving the parameter in the module itself in a nn.ParameterList of the same length as the number of Conv2d layers I have. However, I would prefer to have the parameter stored inside the Conv2d layer it belongs to.

How and where do I have to modify Conv2d to add my custom parameter?

2

You can directly assign new parameters to an nn.Module:

conv = nn.Conv2d(3, 3, 3)
print(dict(conv.named_parameters()))

conv.my_param = nn.Parameter(torch.randn(1))
print(dict(conv.named_parameters()))

which will then show up in .parameters() as well as the state_dict.
Note however, that my_param won’t be used as the forward method wasn’t overridden.

3

Thanks, this worked perfectly! I didn’t expect it to be this easy.