Model Definition

Code

import torch
import torch.nn as nn

# Sequential for simple stacks
mlp = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(128, 10),
)

# Custom module for flexible architecture
class MLP(nn.Module):
    def __init__(self, in_dim, hidden, out_dim):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden)
        self.fc2 = nn.Linear(hidden, out_dim)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        return self.fc2(x)

model = MLP(784, 128, 10)
print(sum(p.numel() for p in model.parameters()))
print(model)

Explanation

A model subclasses nn.Module, registers layers in __init__, and defines the forward pass. Sequential is a shorthand for pure stacks of layers, while custom modules allow branching, multiple inputs, and conditional logic. The parameters() method exposes learnable tensors for the optimizer.

More PyTorch Snippets