Consider the __init__ method from the Introduction to PyTorch:
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
__init__ contains definitions of the architecture for the various layers. This causes PyTorch to register their weights as learnable parameters. If the parameters aren’t registered, they won’t get trained.
We can manually register them using the nn.Parameter class, but this is not normally done.
Max pooling operations don’t have learnable parameters, so they are defined directly in the forward layer.