Building a Neural Network in PyTorch

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

# io.thecodeforge: Production-grade MLP implementation with verification
# Demonstrates the correct nn.Module structure and parameter registration pattern
class ForgeClassifier(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, num_classes: int):
        # MANDATORY: initializes _parameters, _modules, _buffers, _hooks dictionaries
        # Without this, no layer you assign to self will be registered with PyTorch
        super(ForgeClassifier, self).__init__()

        # Structure defined here once — layers are created and registered at init time
        # PyTorch intercepts these assignments via __setattr__ and adds them to _modules
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.bn1 = nn.BatchNorm1d(hidden_size)   # running_mean/var stored as buffers
        self.dropout = nn.Dropout(p=0.3)          # disabled in eval mode automatically
        self.fc2 = nn.Linear(hidden_size, num_classes)

        # Verify the model structure at init time — catch shape bugs during development
        self._verify_forward(input_size, num_classes)

    def _verify_forward(self, input_size: int, num_classes: int):
        """Run a dummy forward pass at init to catch dimension mismatches immediately."""
        with torch.no_grad():
            dummy = torch.randn(2, input_size)  # batch_size=2 for BatchNorm1d compatibility
            out = self.forward(dummy)
            assert out.shape == (2, num_classes), (
                f"Output shape mismatch: expected (2, {num_classes}), got {out.shape}"
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Data flow only — no layer construction here
        # fc1 -> BatchNorm -> ReLU -> Dropout -> fc2
        x = self.fc1(x)         # (batch, input_size) -> (batch, hidden_size)
        x = self.bn1(x)         # normalize across batch dimension
        x = F.relu(x)           # element-wise activation
        x = self.dropout(x)     # zeroes 30% of activations during training, no-op in eval
        x = self.fc2(x)         # (batch, hidden_size) -> (batch, num_classes)
        return x                 # raw logits — apply softmax outside, or use CrossEntropyLoss


# Instantiate — _verify_forward runs immediately, catches shape bugs at construction time
model = ForgeClassifier(input_size=784, hidden_size=256, num_classes=10)

print(model)
print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
print(f"Total parameters (incl. buffers): {sum(p.numel() for p in model.parameters()):,}")

# Verify parameter registration is correct
registered_names = [n for n, _ in model.named_parameters()]
print(f"Registered parameter groups: {registered_names}")

# Move entire model to GPU atomically — all registered parameters and buffers move together
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
print(f"Model device: {next(model.parameters()).device}")

# io.thecodeforge: Production model persistence patterns
# Covers inference deployment, training resumption, and safe loading
import torch
import os
from pathlib import Path

MODEL_DIR = Path("io/thecodeforge/models")
MODEL_DIR.mkdir(parents=True, exist_ok=True)


# ─── Pattern 1: Inference deployment — save state_dict only ─────────────────
deployment_path = MODEL_DIR / "classifier_v1.pth"
torch.save(model.state_dict(), deployment_path)
print(f"Saved inference weights: {deployment_path} ({os.path.getsize(deployment_path) / 1e6:.1f} MB)")

# Load for inference — weights_only=True prevents arbitrary pickle execution
inference_model = ForgeClassifier(input_size=784, hidden_size=256, num_classes=10)
inference_model.load_state_dict(
    torch.load(deployment_path, map_location='cpu', weights_only=True)
)
inference_model.eval()  # MANDATORY: disables Dropout, freezes BatchNorm running stats
inference_model = inference_model.to(device)

# Verify loaded weights match the original
for (n1, p1), (n2, p2) in zip(model.named_parameters(), inference_model.named_parameters()):
    assert n1 == n2, f"Parameter name mismatch: {n1} vs {n2}"
    assert torch.equal(p1.cpu(), p2.cpu()), f"Value mismatch for {n1}"
print("Inference model: all parameters loaded and verified.")


# ─── Pattern 2: Training checkpoint — save full state for resumption ─────────
def save_checkpoint(model, optimizer, epoch: int, val_loss: float, path: Path):
    """Save everything needed to resume training exactly where it left off."""
    torch.save({
        'epoch':                epoch,
        'model_state_dict':     model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'val_loss':             val_loss,
    }, path)
    print(f"Checkpoint saved: epoch {epoch}, val_loss {val_loss:.4f}")


def load_checkpoint(model, optimizer, path: Path, device: torch.device):
    """Resume training from a checkpoint — restores model weights and optimizer state."""
    checkpoint = torch.load(path, map_location=device, weights_only=False)  # dict is safe here
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch'] + 1
    best_val_loss = checkpoint['val_loss']
    print(f"Resumed from epoch {checkpoint['epoch']}, val_loss {best_val_loss:.4f}")
    return model, optimizer, start_epoch, best_val_loss


# Example checkpoint save during training loop
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
checkpoint_path = MODEL_DIR / "checkpoint_epoch_10.pth"
save_checkpoint(model, optimizer, epoch=10, val_loss=0.1823, path=checkpoint_path)

# Resume training from checkpoint
fresh_model = ForgeClassifier(input_size=784, hidden_size=256, num_classes=10).to(device)
fresh_optimizer = torch.optim.Adam(fresh_model.parameters(), lr=1e-3)
fresh_model, fresh_optimizer, start_epoch, best_val = load_checkpoint(
    fresh_model, fresh_optimizer, checkpoint_path, device
)
print(f"Training will resume from epoch {start_epoch}")

# io.thecodeforge: Production PyTorch inference container
# Pin the full version triple — never use 'latest' in production
# PyTorch 2.2.0 + CUDA 12.1 + cuDNN 8 is a tested, stable combination for 2026 deployments
FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime

WORKDIR /app

# Install system-level dependencies before pip — this layer caches independently
# libgl1-mesa-glx is required by OpenCV; libgomp1 is required by some PyTorch operations
RUN apt-get update \
    && apt-get install -y --no-install-recommends \
        libgl1-mesa-glx \
        libgomp1 \
        curl \
    && rm -rf /var/lib/apt/lists/*

# Separate requirements from source — requirements layer caches until requirements.txt changes
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy inference source code only
# Model weights are NOT copied here — they are mounted or downloaded at startup
COPY ./src /app/src

# Model path is configurable via environment variable
# In Kubernetes: mount a PVC at /app/models or use an init container to download from S3
ENV MODEL_PATH=/app/models/classifier_v1.pth
ENV MODEL_INPUT_SIZE=784
ENV MODEL_HIDDEN_SIZE=256
ENV MODEL_NUM_CLASSES=10

# Run as non-root user — required by most enterprise security policies
RUN useradd -m -u 1001 forge
USER forge

# Health check verifies the inference service starts and the model loads correctly
HEALTHCHECK --interval=30s --timeout=15s --start-period=30s --retries=3 \
    CMD curl -f http://localhost:8080/health || exit 1

ENTRYPOINT ["python", "src/inference_service.py"]

# io.thecodeforge: Common nn.Module mistake patterns and their correct counterparts
import torch
import torch.nn as nn


# ─── MISTAKE 1: Missing super().__init__() ───────────────────────────────────
class BrokenInit(nn.Module):
    def __init__(self):
        # super().__init__() omitted — _parameters and _modules never created
        self.fc = nn.Linear(10, 2)  # assigned as a plain Python attribute, invisible to PyTorch

    def forward(self, x):
        return self.fc(x)  # AttributeError at runtime: 'BrokenInit' has no attribute 'training'


class CorrectInit(nn.Module):
    def __init__(self):
        super(CorrectInit, self).__init__()  # FIRST LINE — always
        self.fc = nn.Linear(10, 2)           # now registered in _modules

    def forward(self, x):
        return self.fc(x)


# ─── MISTAKE 2: Python list instead of nn.ModuleList ────────────────────────
class BrokenDynamicModel(nn.Module):
    def __init__(self, depth: int):
        super().__init__()
        # Python list: invisible to model.parameters(), model.to(device), state_dict()
        self.layers = [nn.Linear(64, 64) for _ in range(depth)]

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

# Verify the problem:
bad_model = BrokenDynamicModel(depth=3)
print(f"BrokenDynamicModel trainable params: {sum(p.numel() for p in bad_model.parameters())}")
# Output: BrokenDynamicModel trainable params: 0  <-- optimizer has nothing to update


class CorrectDynamicModel(nn.Module):
    def __init__(self, depth: int):
        super().__init__()
        # nn.ModuleList registers all contained modules — parameters are visible and trackable
        self.layers = nn.ModuleList([nn.Linear(64, 64) for _ in range(depth)])

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

good_model = CorrectDynamicModel(depth=3)
print(f"CorrectDynamicModel trainable params: {sum(p.numel() for p in good_model.parameters()):,}")
# Output: CorrectDynamicModel trainable params: 12,480


# ─── MISTAKE 3: Breaking the gradient chain with .numpy() in forward() ──────
class BrokenForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 2)

    def forward(self, x):
        x = self.fc(x)
        # WRONG: breaks the computational graph — gradients cannot flow past this point
        # x = x.detach().cpu().numpy()   # RuntimeError or silent gradient break
        # WRONG: also breaks it
        # x = x.cpu().numpy()            # RuntimeError: can't call numpy() on tensor requiring grad
        return x  # keep everything as PyTorch tensors inside forward()


# ─── MISTAKE 4: Layers defined inside forward() ─────────────────────────────
class BrokenLayerPlacement(nn.Module):
    def __init__(self):
        super().__init__()
        # No layers defined here — they appear in forward() instead

    def forward(self, x):
        # WRONG: creates a new nn.Linear with random weights on every call
        # The optimizer updates weights from the previous call that no longer exist
        fc = nn.Linear(784, 10)  # new random weights every batch — model never learns
        return fc(x)


# ─── Correct: all layers in __init__, only data flow in forward() ────────────
class CorrectLayerPlacement(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(784, 10)  # created once, reused every forward call

    def forward(self, x):
        return self.fc(x)  # same weights every call — optimizer updates persist


# ─── Device placement verification ──────────────────────────────────────────
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CorrectLayerPlacement().to(device)
print(f"All parameters on device: {next(model.parameters()).device}")

# Always call model(x), never model.forward(x)
# model.forward(x) bypasses __call__, which fires hooks, manages training mode, and tracks autograd
test_input = torch.randn(1, 784).to(device)
output = model(test_input)  # correct
print(f"Output shape: {output.shape}")