What needs to be a class member for models?


I am not clear about what needs to be a class member for a model and what not.
Let’s take a simple model:

class CNN(torch.nn.Module):
def init(self):
super(CNN, self).init()

    self.conv1 = torch.nn.Conv2d(1, 8, 3, 1)
    self.pool1 = torch.nn.MaxPool2d(4)
    self.relu1 = torch.nn.ReLU()

    self.conv2 = torch.nn.Conv2d(8, 8, 3, 1)
    self.pool2 = torch.nn.MaxPool2d(4)
    self.relu2 = torch.nn.ReLU()

1- Do we need to include the pooling and activation functions/objects as class members.
2- Let’s say we use the same pooling and/or activation functions in our model in different layers. Do we need to include them as separate class members or can we use the same object?

As far as I understand optimizer will need to access model parameters to update them. So as long as an object does not have a parameter to be updated we don’t need to include it as a class member. What are the examples of pooling/activation/etc. objects that have parameters?


  1. No. I tend to use the functional variants, F.max_pool2d and friends.
  2. No, as long as both uses require the same options for the pooling you can reuse the same object. You could also use a conv object twice, but then it would use the same weights both times, which is generally not what you want.

Correct. Note that even when an object has learnable parameters it doesn’t necessarily have to be a class member, you can register it manually using self.add_module("name", object).

The pooling layers listed in the docs don’t have parameters, and as far as I can tell, the only activation listed in the docs that has a learnable parameter is PReLU. The various normalisation layers have learnable params if they are initialised with the affine=True option.

1 Like