A neural network needs both operations and a way to keep track of what training can change. In PyTorch, `torch.nn.Module` is the base class for those building blocks. An `nn.Linear` layer is a module; a complete network can be one too. Modules can contain other modules, forming a hierarchy rather than a single flat list of operations. section 6.2 section 8.3
To make a module, subclass `nn.Module`, call `super().__init__()` in the constructor, assign layers to attributes such as `self.conv1`, and define `forward(self, x)` to return the result. Use ordinary PyTorch operations inside `forward`; autograd handles the backward pass. This is useful when data must be reshaped or otherwise handled between layers, beyond a simple `nn.Sequential` chain. section 8.3 section 8.3.1
# In[27]:
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
self.act1 = nn.Tanh()
self.pool1 = nn.MaxPool2d(2)
self.conv2 = nn.Conv2d(16, 8, kernel_size=3, padding=1)
self.act2 = nn.Tanh()
self.pool2 = nn.MaxPool2d(2)
self.fc1 = nn.Linear(8 * 8 * 8, 32)
self.act3 = nn.Tanh()
self.fc2 = nn.Linear(32, 2)
def forward(self, x):
out = self.pool1(self.act1(self.conv1(x)))
out = self.pool2(self.act2(self.conv2(out)))
out = out.view(-1, 8 * 8 * 8) #1
out = self.act3(self.fc1(out))
out = self.fc2(out)
return out
Follow the image classifier from convolution and pooling through `out.view(...)` to its linear layers. The constructor stores the layers; `forward` decides how data moves through them. section 8.3.1
Call an instance as `model(x)`, not `model.forward(x)`. The call invokes `forward` while also running module machinery around it. Calling `forward` directly can skip hooks used during forward or backward passes. section 6.2.1
Assigning a child module to a module attribute registers it. `model.parameters()` then gathers parameters recursively, so an optimizer can reach the weights and biases of nested layers. `named_parameters()` helps identify them; an activation such as `nn.Tanh` does not appear because it has no trainable parameters. Do not hide child modules in an ordinary list or dictionary: use `nn.ModuleList` or `nn.ModuleDict` when those containers are needed. section 8.3.2 section 6.3.2 section 6.2
Check what the optimizer can see
For a custom `Net`, run `model = Net()` and inspect `[p.numel() for p in model.parameters()]`. The resulting counts include parameters from its convolutional and linear submodules. Those same collected parameters can be passed to an optimizer. section 8.3.2 section 8.8
After training, a module can also be passed with an example input to `torch.export.export`, which produces a graph of its computation for deployment within the PyTorch ecosystem. section 17.2.2
