TopolearnPicking layers and getting their shapes right is the everyday work of model building.
torch.nn provides ready-made layers. The ones you meet constantly:
nn.Linear(in, out):y = x W^T + bon the last dimension, any leading batch dims.nn.Conv2d(in_ch, out_ch, kernel_size, stride, padding): 2-d convolution overN, C, H, Winput.padding=1with a 3x3 kernel keeps H and W.nn.MaxPool2d(2): halves H and W.nn.Embedding(num, dim): a lookup table turning int64 token ids into vectors.nn.BatchNorm2d(ch)andnn.LayerNorm(dim): normalization; LayerNorm is the norm of transformers.nn.Dropout(p): randomly zeroes activations during training only.- Activations:
nn.ReLU,nn.GELU,nn.SiLU,nn.Tanh,nn.Sigmoid,nn.Softmax(dim).
Most layers also exist as functions in torch.nn.functional (usually imported as F): F.relu(x), F.gelu(x). Use the module form for layers with weights; the functional form is fine for stateless operations inside forward.
If you do not want to compute an input size by hand, nn.LazyLinear(out) infers in_features on the first call.
Print shapes after each layer when building a new model; a mismatch error such as mat1 and mat2 shapes cannot be multiplied (8x512 and 256x10) tells you exactly which dimension to fix.
Example
import torch
from torch import nn
import torch.nn.functional as F
x = torch.randn(2, 3, 28, 28) # N, C, H, W
conv = nn.Conv2d(3, 8, kernel_size=3, padding=1)
print(conv(x).shape) # torch.Size([2, 8, 28, 28])
print(nn.MaxPool2d(2)(conv(x)).shape) # torch.Size([2, 8, 14, 14])
emb = nn.Embedding(num_embeddings=1000, embedding_dim=16)
tokens = torch.tensor([[1, 5, 42]]) # int64 ids
h = emb(tokens)
print(h.shape) # torch.Size([1, 3, 16])
print(nn.LayerNorm(16)(h).shape) # torch.Size([1, 3, 16])
print(F.gelu(torch.tensor([-1.0, 0.0, 1.0]))) # tensor([-0.1587, 0.0000, 0.8413])Remember
- Linear acts on the last dim; Conv2d expects N, C, H, W
- Embedding maps int64 ids to vectors
- Use modules for layers with weights,
F.*for stateless ops
Watch out for
- Hard-coding a Linear input size after conv layers and getting it wrong; print the shape or use
nn.LazyLinear.