Search code examples
pythonpytorch

How does pytorch automatically know what are my model's parameters?


I have defined the custom class as follows:

class MLPClassifier(nn.Module):
    """
    A basic multi-layer perceptron classifier with 3 layers.
    """
  
    def __init__(self, input_size, hidden_size, num_classes):
        """
        The constructor for the MLPClassifier class.
        """
        super(MLPClassifier, self).__init__()

        self.fc1 = nn.Linear(input_size, hidden_size)  # weights & biases for the input-to-hidden layer
        self.ac1 = nn.ReLU()                           # non-linear activation for the input-to-hidden layer
        self.fc2 = nn.Linear(hidden_size, num_classes) # weights & biases for the hidden-to-output layer
        self.ac2 = nn.Softmax(dim=1)                   # non-linear activation for the hidden-to-output layer

When I run the following script I get this:

hyper_param_input_size  = 4
hyper_param_hidden_size = 64
hyper_param_num_classes = 3

model = MLPClassifier(hyper_param_input_size, hyper_param_hidden_size, hyper_param_num_classes)

for p in model.parameters():
    print(p.shape)

>>> torch.Size([64, 4])
>>> torch.Size([64])
>>> torch.Size([3, 64])
>>> torch.Size([3])

How on earth does PyTorch automatically know about my internally defined attributes when I never explicitly told it? Does it loop through everything in the class and check if isinstance(self, nn.Layer) or something?


Solution

  • The nn.Module.parameters function will recursively go through all child modules of the parent and return all its parameters. It does not have to do with the actual structure of your MLPClassifier module. When defining a new submodule attribute inside your __init__ the parent module will register it as a child module, this way their parameters (if any, eg. your nn.ReLU and nn.Softmax don't have any...) can later be accessed via the parameters call.