Neural Network Modules

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_features: int, hidden: int, out_features: int):
        super().__init__()                        # always call first
        self.fc1 = nn.Linear(in_features, hidden)
        self.act = nn.ReLU()
        self.fc2 = nn.Linear(hidden, out_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.act(self.fc1(x))
        return self.fc2(x)

model = MLP(784, 256, 10)
output = model(x_batch)   # calls model.forward(x_batch)

# Parameters — learned, included in optimizer
w = nn.Parameter(torch.randn(3, 3))
self.weight = nn.Parameter(torch.empty(out, inp))

# Buffers — not learned, saved with state_dict, moved with .to()
self.register_buffer('running_mean', torch.zeros(num_features))

# Temporary state — not persisted (just an attribute)
self.cache = None

# Inspecting
list(model.parameters())           # all Parameter objects (recursive)
list(model.named_parameters())     # (name, param) pairs
list(model.buffers())
list(model.named_buffers())

list(model.children())            # direct children only
list(model.modules())             # all modules (recursive, depth-first)
list(model.named_modules())       # (name, module) pairs
list(model.named_children())

model.fc1                         # direct attribute access
getattr(model, 'fc1')             # programmatic access

# Sequential — runs modules in order
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
)

# OrderedDict for named layers
from collections import OrderedDict
model = nn.Sequential(OrderedDict([
    ('fc1', nn.Linear(784, 256)),
    ('relu', nn.ReLU()),
    ('fc2', nn.Linear(256, 10)),
]))

# ModuleList — indexable list, all registered
self.layers = nn.ModuleList([nn.Linear(64, 64) for _ in range(6)])
for layer in self.layers:
    x = layer(x)

# ModuleDict — dict of modules, all registered
self.branches = nn.ModuleDict({'left': nn.Linear(64, 32), 'right': nn.Linear(64, 32)})
x = self.branches['left'](x)

# ParameterList / ParameterDict — same idea for raw Parameters
self.weights = nn.ParameterList([nn.Parameter(torch.randn(64, 64)) for _ in range(4)])
self.named_w = nn.ParameterDict({'a': nn.Parameter(torch.randn(3)), 'b': nn.Parameter(torch.randn(3))})

model.train()        # sets self.training = True on all sub-modules
model.eval()         # sets self.training = False

# Layers affected by mode:
# - nn.Dropout: drops during train, identity during eval
# - nn.BatchNorm*: uses batch stats during train, running stats during eval

# Check current mode
print(model.training)

# Context manager (restores mode on exit)
from contextlib import contextmanager
# no built-in context manager, but common pattern:
was_training = model.training
model.eval()
with torch.no_grad():
    out = model(x)
model.train(was_training)

model.to('cuda')
model.to('cuda:0')
model.to(device)
model.cuda()                      # shortcut
model.cpu()

model.to(torch.float16)           # cast all params/buffers
model.to(dtype=torch.bfloat16)
model.to(device='cuda', dtype=torch.float16)  # combined

model.double()   # float64
model.float()    # float32
model.half()     # float16

# Freeze entire model (common in fine-tuning)
for param in model.parameters():
    param.requires_grad = False

# Freeze specific sub-module
for param in model.backbone.parameters():
    param.requires_grad = False

# Unfreeze later
for param in model.head.parameters():
    param.requires_grad = True

# Only pass trainable parameters to optimizer
optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3
)

# Manual in __init__
def _init_weights(self):
    for m in self.modules():
        if isinstance(m, nn.Linear):
            nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
            nn.init.ones_(m.weight)
            nn.init.zeros_(m.bias)

# Common initializers
nn.init.uniform_(t, a=0.0, b=1.0)
nn.init.normal_(t, mean=0.0, std=1.0)
nn.init.constant_(t, val=0.0)
nn.init.ones_(t)
nn.init.zeros_(t)
nn.init.eye_(t)
nn.init.xavier_uniform_(t, gain=1.0)       # Glorot uniform
nn.init.xavier_normal_(t, gain=1.0)        # Glorot normal
nn.init.kaiming_uniform_(t, mode='fan_in', nonlinearity='leaky_relu')
nn.init.kaiming_normal_(t, mode='fan_in', nonlinearity='relu')
nn.init.orthogonal_(t, gain=1.0)
nn.init.trunc_normal_(t, mean=0.0, std=1.0, a=-2.0, b=2.0)  # truncated normal

# model.apply(fn) — calls fn on every sub-module (post-order DFS)
def init_linear(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        nn.init.zeros_(m.bias)

model.apply(init_linear)

# Forward hook — called after forward(), receives input and output
handle = module.register_forward_hook(fn)

# Forward pre-hook — called before forward(), receives input
handle = module.register_forward_pre_hook(fn)

# Backward hook (full)
handle = module.register_full_backward_hook(fn)

# Always remove hooks when done to avoid memory leaks
handle.remove()

# Context manager pattern
from torch.nn.modules.module import register_module_forward_hook

model.zero_grad()              # zero all .grad on all params
model.parameters(recurse=True) # include sub-modules (default True)

# Get/set state without instantiating new module
state = model.state_dict()        # OrderedDict of all params + buffers
model.load_state_dict(state)      # restore
model.load_state_dict(state, strict=False)  # allow missing/extra keys

# Count parameters
total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)

# torch.compile wraps any nn.Module for TorchDynamo/Inductor acceleration
compiled = torch.compile(model)
compiled = torch.compile(model, mode='reduce-overhead')   # repeat workloads
compiled = torch.compile(model, mode='max-autotune')      # slow compile, fast run
compiled = torch.compile(model, fullgraph=True)           # no graph breaks
compiled = torch.compile(model, backend='aot_eager')      # debug backend