Topolearnnn.Module is how every PyTorch model, layer and loss is structured, from a two-layer MLP to an LLM.
A model is a class that inherits from torch.nn.Module. You create sub-layers in __init__ and assign them to attributes; you describe how input flows through them in forward.
Assigning a module or nn.Parameter to self registers it. That registration is what makes these work automatically:
model.parameters()andmodel.named_parameters()iterate all trainable weights (pass them to the optimizer).model.to(device)moves every parameter and buffer.model.state_dict()collects weights for saving.model.train()/model.eval()switch the mode of every sub-module.
Always call super().__init__() first, and call the model as model(x), not model.forward(x), so that hooks run.
Layers held in a plain Python list are not registered; use nn.ModuleList or nn.ModuleDict instead. For a simple stack of layers, nn.Sequential(layer1, layer2, ...) saves writing a class at all, and its items can be indexed like a list.
Because forward is ordinary Python, you can use loops, conditionals and print statements inside it, which makes custom architectures straightforward to write and debug.
Example
import torch
from torch import nn
class MLP(nn.Module):
def __init__(self, n_in, n_hidden, n_out):
super().__init__()
self.fc1 = nn.Linear(n_in, n_hidden) # registered automatically
self.fc2 = nn.Linear(n_hidden, n_out)
def forward(self, x):
return self.fc2(torch.relu(self.fc1(x)))
model = MLP(4, 16, 3)
out = model(torch.randn(8, 4)) # call the model, not .forward
print(out.shape) # torch.Size([8, 3])
print(sum(p.numel() for p in model.parameters())) # 131
for name, p in model.named_parameters():
print(name, tuple(p.shape)) # fc1.weight (16, 4), fc1.bias (16,), ...
seq = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
print(seq(torch.randn(8, 4)).shape) # torch.Size([8, 3])Remember
- Create layers in
__init__, compute inforward, callmodel(x) - Attribute assignment registers parameters; plain lists do not
nn.Sequentialcovers simple layer stacks
Watch out for
- Storing layers in a Python list, so the optimizer never sees their parameters; use
nn.ModuleList. - Forgetting
super().__init__(), which fails as soon as you assign a layer.