Does pytorch register modules that are assigned to object properties?

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?

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