nn.Module

import torch
import torch.nn as nn
import torch.nn.functional as F

# 1) Basic Module
class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim,  hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(0.5)

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

model = MLP(784, 256, 10)
print(model)                       # auto-generated summary

# 2) Inspect parameters
for name, param in model.named_parameters():
    print(name, param.shape, param.requires_grad)

total = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'Trainable: {total:,}')

# 3) Move to device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)
x = torch.randn(32, 784).to(device)
logits = model(x)

# 4) Save / load state
torch.save(model.state_dict(), 'model.pt')
model.load_state_dict(torch.load('model.pt', weights_only=True))

# 5) nn.Sequential — quick layer stack
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
)

# 6) ModuleList + ModuleDict — for collections of layers
class DeepNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(784, 256),
            nn.Linear(256, 256),
            nn.Linear(256, 10),
        ])
        # Or by name:
        self.heads = nn.ModuleDict({
            'classifier': nn.Linear(256, 10),
            'regressor':  nn.Linear(256, 1),
        })

    def forward(self, x, head='classifier'):
        for layer in self.layers[:-1]:
            x = F.relu(layer(x))
        return self.heads[head](x)

# 7) Common layers
# Linear / fully connected
nn.Linear(in_features, out_features, bias=True)

# Convolutional
nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
nn.Conv1d(...)
nn.Conv3d(...)
nn.MaxPool2d(kernel_size=2)
nn.AvgPool2d(2)
nn.AdaptiveAvgPool2d((1, 1))

# Normalisation
nn.BatchNorm1d(num_features)
nn.BatchNorm2d(num_channels)
nn.LayerNorm(normalized_shape)
nn.GroupNorm(num_groups, num_channels)

# Activation
nn.ReLU()
nn.LeakyReLU(0.1)
nn.GELU()
nn.SiLU()                              # aka Swish
nn.Tanh()
nn.Sigmoid()
nn.Softmax(dim=-1)

# Regularisation
nn.Dropout(p=0.5)
nn.Dropout2d(p=0.25)                   # entire channels

# Recurrent
nn.LSTM(input_size, hidden_size, num_layers=2, bidirectional=True, batch_first=True)
nn.GRU(input_size, hidden_size)
nn.RNN(input_size, hidden_size)

# Embedding
nn.Embedding(num_embeddings, embedding_dim)
nn.EmbeddingBag(num_embeddings, embedding_dim, mode='mean')

# Transformer
nn.TransformerEncoderLayer(d_model=512, nhead=8)
nn.TransformerEncoder(layer, num_layers=6)
nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)

# 8) Activation in forward (no parameters → F. functional)
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.fc    = nn.Linear(64 * 7 * 7, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.flatten(1)
        return self.fc(x)

# F. vs nn. — when to use which:
#   nn.X — has learnable parameters (BatchNorm, Conv, Linear, Embedding)
#   F.X  — stateless (ReLU, MaxPool, dropout when you control training mode)

# 9) train() / eval() — controls Dropout + BatchNorm behaviour
model.train()              # dropout active, BN uses batch stats
logits = model(x)
loss   = loss_fn(logits, y)
loss.backward()

model.eval()               # dropout off, BN uses running stats
with torch.inference_mode():
    preds = model(x)

# 10) Freezing parameters (transfer learning)
for param in base.parameters():
    param.requires_grad = False

# Or freeze specific layers
for name, param in model.named_parameters():
    if name.startswith('fc1') or name.startswith('fc2'):
        param.requires_grad = False

# Train only unfrozen
optimiser = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3,
)

# 11) Custom layer with parameters
class Linear2(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(in_dim, out_dim) * 0.01)
        self.bias   = nn.Parameter(torch.zeros(out_dim))

    def forward(self, x):
        return x @ self.weight + self.bias

# nn.Parameter() — tensors that get registered as model params automatically.

# 12) Apply — initialise weights
def init_weights(m):
    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, nonlinearity='relu')

model.apply(init_weights)

# 13) Hooks — peek at activations / gradients
activations = {}
def get_activation(name):
    def hook(model, input, output):
        activations[name] = output.detach()
    return hook

model.conv1.register_forward_hook(get_activation('conv1'))
_ = model(x)
print(activations['conv1'].shape)

# Backward hook for gradients
def grad_hook(module, grad_input, grad_output):
    print(module.__class__.__name__, grad_output[0].norm())

for module in model.modules():
    module.register_full_backward_hook(grad_hook)

# 14) Compose — composition is cleaner than mega-classes
class ResBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
        self.bn1   = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
        self.bn2   = nn.BatchNorm2d(channels)

    def forward(self, x):
        residual = x
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.bn2(self.conv2(x))
        return F.relu(x + residual)

class ResNet(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.stem = nn.Sequential(nn.Conv2d(3, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU())
        self.blocks = nn.Sequential(*[ResBlock(32) for _ in range(4)])
        self.head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, num_classes))

    def forward(self, x):
        return self.head(self.blocks(self.stem(x)))

# 15) state_dict — for checkpointing + transfer
state = model.state_dict()
for k, v in state.items():
    print(k, v.shape)

model.load_state_dict(state, strict=False)    # allow shape mismatches

# 16) torch.compile — speed up forward/backward (PyTorch 2.0+)
model = torch.compile(model)    # uses TorchInductor backend; 1.5-2x faster on many models

# 17) Best practices
#   • Keep __init__ thin: just declare submodules + params
#   • Put activation/computation logic in forward()
#   • Use Sequential for linear stacks; subclass for branching
#   • Use nn.ModuleList / ModuleDict, NOT plain Python list/dict — params get registered properly
#   • Set device once (model.to(device)) — children auto-move
#   • Use .train() / .eval() at the right time

# 18) Common bugs
#   • Putting layers in a plain list → parameters not registered
#   • Forgetting .to(device) → tensor-mismatch errors
#   • Mixing F.dropout in forward without checking training mode
#   • Not zero-grad-ing the optimizer → grads accumulate across batches

import torch.nn as nn
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 256), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(256, 10),
        )
    def forward(self, x):
        return self.net(x)