A practical reference catalog of the most common nn.Module layer types โ each already covered conceptually in earlier categories of this hub, gathered here in one place with exact PyTorch syntax.
Common Layer Types
| Layer | Syntax | Covered Conceptually In |
|---|---|---|
| Fully connected | nn.Linear(in_features, out_features) | Weights and Bias |
| 2D convolution | nn.Conv2d(in_channels, out_channels, kernel_size) | CNN Fundamentals category |
| Max pooling | nn.MaxPool2d(kernel_size) | Max Pooling |
| Batch normalization | nn.BatchNorm2d(num_features) | Batch Normalization |
| Dropout | nn.Dropout(p=0.5) | Dropout |
| Recurrent (RNN/LSTM/GRU) | nn.LSTM(input_size, hidden_size) | LSTM & GRU category |
| Embedding lookup | nn.Embedding(vocab_size, embed_dim) | Word Embeddings |
| Multi-head attention | nn.MultiheadAttention(embed_dim, num_heads) | Multi-Head Attention |
| Sequential container | nn.Sequential(layer1, layer2, ...) | A convenience wrapper, not a layer itself |
Code โ Composing Layers Into a Model
import torch.nn as nn
model = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64 * 8 * 8, 128), nn.ReLU(), nn.Dropout(0.5),
nn.Linear(128, 10)
)
nn.Sequential vs a Custom nn.Module
nn.Sequential is convenient for simple, strictly-linear architectures โ layers execute in the exact order listed, one feeding directly into the next. For anything with branching, skip connections, or conditional logic (like a ResNet's residual blocks, or an encoder-decoder), a custom nn.Module subclass with an explicit forward() method (see nn.Module) is required instead.
Common Mistakes
- Miscalculating the flattened dimension going into the first
nn.Linearlayer after convolutional/pooling layers โ this exact shape depends on the input image size and every preceding layer's stride/padding, and getting it wrong produces a shape-mismatch error. - Using
nn.Sequentialfor an architecture that genuinely needs branching or skip connections โ this forces awkward workarounds; a customnn.Moduleis the correct tool.
Interview Relevance
Q: "When would you use nn.Sequential versus writing a custom nn.Module subclass?" nn.Sequential works well for architectures that are a strictly linear chain of layers, each feeding directly into the next, with no branching. Any architecture needing skip connections, multiple inputs/outputs, or conditional logic in the forward pass (like a ResNet block or an encoder-decoder) requires a custom nn.Module subclass with an explicit forward() method.
Practice Question
Why can't a ResNet-style residual block (requiring \(x + F(x)\)) be expressed cleanly using nn.Sequential alone?