PyTorch foundations

🎯 Goal

After reading this chapter:

• You will know PyTorch’s tensor, autograd, nn.Module, and DataLoader APIs
• You will be able to build your own model via nn.Module
• You will be able to write a complete training loop (on CPU or GPU)
• You will know how to save, load, and run inference on a model
• You will get familiar with production deployment (torch.jit, ONNX)

What to learn

• Tensor — NumPy ndarray + GPU support + autograd
• Autograd — automatic differentiation
• nn.Module — building models
• nn.Linear, nn.Conv2d, nn.RNN — layers
• Loss functions — nn.MSELoss, nn.CrossEntropyLoss, etc.
• Optimizers — optim.SGD, optim.Adam
• Dataset and DataLoader — batch loading
• Device management — CPU/GPU/MPS
• Saving/loading — state_dict
• TorchScript — production export

Libraries

# CPU
pip install torch torchvision torchaudio

# CUDA 12.1 (NVIDIA GPU)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# Mac M1/M2/M3 — default install works with MPS

Verify:

import torch
print(torch.__version__)
print(torch.cuda.is_available())     # NVIDIA GPU
print(torch.backends.mps.is_available())  # Mac

Important topics

Tensor — the heart of PyTorch

import torch

# Creation
a = torch.tensor([1, 2, 3])              # int64
b = torch.tensor([1.0, 2.0, 3.0])        # float32
c = torch.zeros(3, 4)                    # 3x4 zeros
d = torch.randn(2, 3)                    # normal random
e = torch.arange(10)                     # [0..9]

# From NumPy
import numpy as np
arr = np.array([1, 2, 3])
t = torch.from_numpy(arr)                # share memory!
arr_back = t.numpy()                     # share memory!

# Attributes
print(a.shape, a.dtype, a.device)        # torch.Size([3]) torch.int64 cpu

Device management

# Best device automatically
device = "cuda" if torch.cuda.is_available() else (
    "mps" if torch.backends.mps.is_available() else "cpu"
)

# Move tensor to device
x = torch.randn(1000, 1000).to(device)
model = MyModel().to(device)

# Note: both tensors must be on the same device # y = x @ y # ❌ if y is on CPU # y = x @ y.to(device) # ✅

Autograd — magic

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1     # y = x² + 3x + 1
y.backward()                # dy/dx = 2x + 3
print(x.grad)               # tensor(7.0) ← at x=2: 2*2+3=7

# Main mechanism — a computational graph is built, and when backward is called # gradients are computed for each x

nn.Module pattern

import torch.nn as nn

class MyModel(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.3)
        self.activation = nn.ReLU()
    
    def forward(self, x):
        x = self.activation(self.fc1(x))
        x = self.dropout(x)
        x = self.activation(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x

model = MyModel(input_dim=784, hidden_dim=256, output_dim=10)
print(model)
print(f"Parameters:{sum(p.numel() for p in model.parameters()):,}")

Dataset and DataLoader

from torch.utils.data import Dataset, DataLoader

class CSVDataset(Dataset):
    def __init__(self, csv_path):
        df = pd.read_csv(csv_path)
        self.X = torch.tensor(df.drop("target", axis=1).values, dtype=torch.float32)
        self.y = torch.tensor(df["target"].values, dtype=torch.long)
    
    def __len__(self):
        return len(self.y)
    
    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

train_dataset = CSVDataset("train.csv")
train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,       # parallel data loading
    pin_memory=True,     # fast for GPU
)

Code examples

Complete training loop

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

device = "cuda" if torch.cuda.is_available() else "cpu"

# Model
model = MyModel(784, 256, 10).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

# Training loop
def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss, total_correct, total = 0, 0, 0
    
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        
        optimizer.zero_grad()
        logits = model(X)
        loss = criterion(logits, y)
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item() * X.size(0)
        total_correct += (logits.argmax(dim=1) == y).sum().item()
        total += X.size(0)
    
    return total_loss / total, total_correct / total

@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss, total_correct, total = 0, 0, 0
    
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        logits = model(X)
        loss = criterion(logits, y)
        
        total_loss += loss.item() * X.size(0)
        total_correct += (logits.argmax(dim=1) == y).sum().item()
        total += X.size(0)
    
    return total_loss / total, total_correct / total

# Training
EPOCHS = 20
for epoch in range(EPOCHS):
    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)
    
    print(f"Epoch{epoch+1}/{EPOCHS}"
          f"Train: loss={train_loss:.4f}, acc={train_acc:.4f}"
          f"Val: loss={val_loss:.4f}, acc={val_acc:.4f}")

Saving and loading a model

# Only weights (RECOMMENDED)
torch.save(model.state_dict(), "model.pt")

# Load
model = MyModel(784, 256, 10)
model.load_state_dict(torch.load("model.pt", map_location="cpu"))
model.eval()

# Full checkpoint (for resuming training)
torch.save({
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": loss,
}, "checkpoint.pt")

checkpoint = torch.load("checkpoint.pt")
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

TorchScript — production export

# Tracing (needs input examples)
example_input = torch.randn(1, 784).to(device)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_traced.pt")

# Scripting (full Python control flow)
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

# Loading (no Python needed!)
loaded = torch.jit.load("model_traced.pt")
output = loaded(torch.randn(1, 784))

ONNX export

torch.onnx.export(
    model,
    example_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    opset_version=17,
)

# Loading in ONNX Runtime
import onnxruntime as ort
sess = ort.InferenceSession("model.onnx")
output = sess.run(None, {"input": input_array})[0]

Backend integration

PyTorch model in FastAPI

from fastapi import FastAPI
from pydantic import BaseModel
import torch
from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app):
    app.state.device = "cuda" if torch.cuda.is_available() else "cpu"
    app.state.model = MyModel(784, 256, 10).to(app.state.device)
    app.state.model.load_state_dict(torch.load("model.pt", map_location=app.state.device))
    app.state.model.eval()
    yield

app = FastAPI(lifespan=lifespan)

class Input(BaseModel):
    features: list[float]  # 784 elements

@app.post("/predict")
@torch.no_grad()
def predict(data: Input):
    X = torch.tensor([data.features], dtype=torch.float32).to(app.state.device)
    logits = app.state.model(X)
    probs = torch.softmax(logits, dim=1)
    pred_class = probs.argmax(dim=1).item()
    confidence = probs[0, pred_class].item()
    return {"class": pred_class, "confidence": confidence}

Batch prediction (efficient)

@app.post("/predict/batch")
@torch.no_grad()
def predict_batch(items: list[Input]):
    X = torch.tensor([item.features for item in items], dtype=torch.float32).to(device)
    logits = app.state.model(X)
    probs = torch.softmax(logits, dim=1)
    return [
        {"class": int(p.argmax().item()), "confidence": float(p.max().item())}
        for p in probs
    ]

Production tips

1. model.eval() — Dropout and BatchNorm behave differently in production
2. torch.no_grad() — disables gradient tracking (faster + memory)
3. torch.inference_mode() — no_grad + bonus optimization
4. Batching — 64 inputs per request — better GPU utilization
5. TorchServe — batching, versioning, A/B test in production (Month 6)
6. Async serving — asyncio + to_thread (CPU bound) or Triton/BentoML

Resources

• PyTorch tutorials — pytorch.org/tutorials
• “Deep Learning with PyTorch” — Eli Stevens (free PDF: pytorch.org/deep-learning-with-pytorch)
• PyTorch Lightning — wrapper for less boilerplate
• Karpathy — “Let’s build GPT”(YouTube) — deep PyTorch
• Hugging Face Course — for PyTorch transformers

🏋️ Exercises

🟢 Easy

1. Create a torch.randn(3, 4) tensor, do transpose, sum, mean.
2. Find the gradient of f(x) = x³ at x=3 with requires_grad=True.
3. Create nn.Linear(10, 1), do a forward pass, output the number of parameters.

🟡 Medium

1. MNIST MLP: Get 95%+ accuracy on MNIST with a 2-layer MLP.
2. Custom dataset: Create your own Dataset class with CSV.
3. GPU check: Train the model on CPU and GPU, measure the time difference.

🔴 Hard

1. FastAPI + PyTorch service: MNIST classifier, image upload, returns prediction. With Docker.
2. TorchScript benchmark: Compare a plain model and its TorchScript version by latency (timeit).
3. Multi-GPU: Train on 2+ GPUs with nn.DataParallel or DistributedDataParallel (using Colab Pro or Kaggle).

Capstone

notebooks/month-03/02_pytorch_mnist.ipynb:

• Load the MNIST dataset via torchvision.datasets
• Write a 3-layer MLP
• Train + Validation loop
• 97%+ accuracy on the test set
• Confusion matrix
• Visualize the worst examples
• Export the model to TorchScript
• Create a FastAPI endpoint

✅ Checklist

• Tensor creation, operations, device transfer
• Autograd basics (requires_grad, backward, grad)
• nn.Module subclassing
• Batch loading with DataLoader
• Writing training loops (train + eval mode, zero_grad, optimizer.step)
• Saving and loading a model (state_dict)
• TorchScript or ONNX export
• PyTorch serving with FastAPI

Moving on to TensorFlow and Keras.