Does pytorch register modules that are assigned to object properties?

1

Say I have a module like:

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn = CNN(params)

And then I do:

module = MyModule()
module.cnn = CNN(some_other_params)

Is the replacement registered? Will there be any nasty side effects down the line?

2

Yes, the assignment of the new module will properly register the new layer with all parameters and is the recommended way to e.g. replace the last linear layer in a fine-tuning use case (assuming the number of classes changed and you would like to retrain this layer). The parameters are also shown in the state_dict and in the parameters() generator:

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn = nn.Conv2d(3, 1, 3, 1, 1)

model = MyModule()
print(model.state_dict())
print(dict(model.named_parameters()))

model.cnn = nn.Linear(3, 3)
print(model.state_dict())
print(dict(model.named_parameters()))

You could run into “side effects”, if you’ve already created an optimizer and don’t add the newly registered parameters to it, since they wouldn’t be trained.

1 Like