S100 · PyTorch Fundamentals — Tensors, Autograd, nn.Module

"""pytorch_lab.py — the same MLP as session 097, in PyTorch.
 
Runs on CPU by default. If CUDA is available it will use GPU automatically.
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
 
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {DEVICE}")
 
 
# ---------------- model ----------------
class MLP(nn.Module):
    """784 → 128 → 64 → 10, ReLU hidden, softmax done by CrossEntropyLoss."""
 
    def __init__(self, in_dim: int = 784, h1: int = 128, h2: int = 64,
                 out_dim: int = 10) -> None:
        super().__init__()
        self.fc1 = nn.Linear(in_dim, h1)
        self.fc2 = nn.Linear(h1, h2)
        self.fc3 = nn.Linear(h2, out_dim)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.view(x.size(0), -1)          # flatten (B, 1, 28, 28) → (B, 784)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)                 # logits — no softmax, CE will do it
 
 
# ---------------- data ----------------
def make_loaders(batch: int = 128):
    tfm = transforms.Compose([
        transforms.ToTensor(),                       # [0, 1] float
        transforms.Normalize((0.1307,), (0.3081,)),  # MNIST mean/std
    ])
    root = "./data"
    train_ds = datasets.MNIST(root, train=True,  download=True, transform=tfm)
    test_ds  = datasets.MNIST(root, train=False, download=True, transform=tfm)
    train_ld = DataLoader(train_ds, batch_size=batch, shuffle=True,  num_workers=2)
    test_ld  = DataLoader(test_ds,  batch_size=1024, shuffle=False, num_workers=2)
    return train_ld, test_ld
 
 
# ---------------- train / eval ----------------
def train_one_epoch(model, loader, optimizer, criterion) -> float:
    model.train()                                    # dropout/BN in training mode
    running_loss = 0.0
    for x, y in loader:
        x, y = x.to(DEVICE), y.to(DEVICE)
        optimizer.zero_grad()                        # clear stale gradients
        logits = model(x)                            # forward
        loss = criterion(logits, y)                  # CE
        loss.backward()                              # autograd — compute grads
        optimizer.step()                             # apply the update
        running_loss += loss.item() * x.size(0)
    return running_loss / len(loader.dataset)
 
 
@torch.no_grad()                                     # turn off autograd → faster, less memory
def evaluate(model, loader, criterion) -> tuple[float, float]:
    model.eval()                                     # dropout/BN in eval mode
    total_loss, correct, n = 0.0, 0, 0
    for x, y in loader:
        x, y = x.to(DEVICE), y.to(DEVICE)
        logits = model(x)
        total_loss += criterion(logits, y).item() * x.size(0)
        correct += (logits.argmax(1) == y).sum().item()
        n += x.size(0)
    return total_loss / n, correct / n
 
 
def count_params(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
 
if __name__ == "__main__":
    torch.manual_seed(0)
 
    train_ld, test_ld = make_loaders(batch=128)
    model = MLP().to(DEVICE)
    print(f"Model: {model}\n  parameters: {count_params(model):,}")
 
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()                # softmax + NLL fused, numerically stable
 
    for epoch in range(1, 11):
        tr_loss = train_one_epoch(model, train_ld, optimizer, criterion)
        te_loss, te_acc = evaluate(model, test_ld, criterion)
        print(f"epoch {epoch:>2}: train_loss={tr_loss:.4f}  "
              f"test_loss={te_loss:.4f}  test_acc={te_acc:.4f}")
 
    # Save the trained model — deployment artifact
    torch.save(model.state_dict(), "mlp_mnist.pt")
    print("Saved to mlp_mnist.pt")
 
    # Load in a fresh session:
    #   model = MLP(); model.load_state_dict(torch.load("mlp_mnist.pt")); model.eval()