PyTorch 2.5 Complete Guide

1# CPU only
2pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
3
4# CUDA 12.1 (GPU)
5pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
6
7import torch
8print(torch.__version__)
9print(torch.cuda.is_available())   # True if GPU

1import torch
2
3# Create tensors
4x = torch.randn(3, 4, requires_grad=True)
5y = torch.ones(4, 2)
6z = x @ y + 1   # matrix multiply + bias
7
8# Automatic differentiation
9loss = z.sum()
10loss.backward()
11print(x.grad)   # dL/dx
12
13# Move to GPU
14device = "cuda" if torch.cuda.is_available() else "cpu"
15x = x.to(device)

1import torch.nn as nn
2
3class MLP(nn.Module):
4    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
5        super().__init__()
6        self.net = nn.Sequential(
7            nn.Linear(input_dim, hidden_dim),
8            nn.BatchNorm1d(hidden_dim),
9            nn.ReLU(),
10            nn.Dropout(0.3),
11            nn.Linear(hidden_dim, hidden_dim // 2),
12            nn.ReLU(),
13            nn.Dropout(0.3),
14            nn.Linear(hidden_dim // 2, output_dim),
15        )
16
17    def forward(self, x: torch.Tensor) -> torch.Tensor:
18        return self.net(x)
19
20model = MLP(784, 512, 10).to(device)
21print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

1criterion = nn.CrossEntropyLoss()
2optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
3scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
4
5best_acc = 0.0
6for epoch in range(50):
7    model.train()
8    for batch_x, batch_y in train_loader:
9        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
10        optimizer.zero_grad()
11        logits = model(batch_x)
12        loss   = criterion(logits, batch_y)
13        loss.backward()
14        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
15        optimizer.step()
16
17    # Validation
18    model.eval()
19    with torch.no_grad():
20        correct = sum(
21            (model(x.to(device)).argmax(1) == y.to(device)).sum().item()
22            for x, y in val_loader
23        )
24        acc = correct / len(val_loader.dataset)
25
26    scheduler.step()
27    if acc > best_acc:
28        best_acc = acc
29        torch.save(model.state_dict(), "best_model.pt")

1import torchvision.models as models
2import torchvision.transforms as transforms
3
4# Transfer learning with ResNet-50
5model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
6model.fc = nn.Linear(model.fc.in_features, num_classes)  # replace head
7model = model.to(device)
8
9# Data augmentation
10train_transform = transforms.Compose([
11    transforms.RandomResizedCrop(224),
12    transforms.RandomHorizontalFlip(),
13    transforms.ColorJitter(brightness=0.2, contrast=0.2),
14    transforms.ToTensor(),
15    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
16])

1from torch.cuda.amp import autocast, GradScaler
2
3scaler = GradScaler()
4
5for batch_x, batch_y in train_loader:
6    batch_x, batch_y = batch_x.to(device), batch_y.to(device)
7    optimizer.zero_grad()
8
9    with autocast():   # FP16 forward pass
10        logits = model(batch_x)
11        loss   = criterion(logits, batch_y)
12
13    scaler.scale(loss).backward()
14    scaler.unscale_(optimizer)
15    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
16    scaler.step(optimizer)
17    scaler.update()

1# TorchScript — language-independent serialisation
2model.eval()
3scripted = torch.jit.script(model)
4scripted.save("model_scripted.pt")
5
6# ONNX export — run in any ONNX-compatible runtime
7dummy = torch.randn(1, 3, 224, 224).to(device)
8torch.onnx.export(
9    model, dummy, "model.onnx",
10    input_names=["image"], output_names=["logits"],
11    dynamic_axes={"image": {0: "batch_size"}},
12)
13
14# Serve with TorchServe
15torch-model-archiver --model-name resnet50   --version 1.0   --serialized-file model_scripted.pt   --handler image_classifier   --export-path model_store
16torchserve --start --model-store model_store --models resnet50=resnet50.mar