PyTorch 入门（工程向）

0. 张量与自动求导速记

import torch

x = torch.randn(4, 5, device='cpu')
w = torch.randn(5, 3, requires_grad=True)
y = (x @ w).sum()
y.backward()
print(w.grad.shape)  # (5, 3)

• 模型与数据必须在同一 device。
• backward() 后梯度累加到 .grad，每步先 optimizer.zero_grad()。
• 推理用 model.eval() + torch.no_grad()。

1. 推荐工程目录

project/
├── configs/
├── datasets/
├── models/
├── trainers/
├── tools/
└── train.py

2. 标准化训练模板

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

class MLP(nn.Module):
    def __init__(self, in_dim=784, num_classes=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, 256), nn.ReLU(),
            nn.Linear(256, 128), nn.ReLU(),
            nn.Linear(128, num_classes)
        )

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


def train_one_epoch(model, loader, optimizer, criterion, device="cuda"):
    model.train()
    total_loss = 0.0
    total = 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        logits = model(xb)
        loss = criterion(logits, yb)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * xb.size(0)
        total += xb.size(0)
    return total_loss / total

3. 常见技巧

• torch.cuda.amp：混合精度节省显存。
• grad_clip：防止梯度爆炸。
• model.eval() + torch.no_grad()：验证/推理必须加。
• pin_memory=True + num_workers>0：提升数据加载性能。

4. 检查点保存

torch.save({
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "epoch": epoch,
}, "ckpt.pt")

恢复：

ckpt = torch.load("ckpt.pt", map_location="cpu")
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
start_epoch = ckpt["epoch"] + 1