PyTorch 使用方式与典型场景

四、使用方式与典型场景

示例一：一个完整可跑的极简训练循环

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

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

# 造一批假数据：1000 条 32 维输入，回归到 1 维
x = torch.randn(1000, 32)
w_true = torch.randn(32, 1)
y = x @ w_true + 0.1 * torch.randn(1000, 1)

ds = TensorDataset(x, y)
dl = DataLoader(ds, batch_size=64, shuffle=True,
                num_workers=2, pin_memory=(device == "cuda"))

model = nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 1)).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2)
lossf = nn.MSELoss()

model.train()
for epoch in range(20):
    total, n = 0.0, 0
    for xb, yb in dl:
        xb = xb.to(device, non_blocking=True)
        yb = yb.to(device, non_blocking=True)
        opt.zero_grad(set_to_none=True)
        loss = lossf(model(xb), yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        opt.step()
        total += loss.item() * xb.size(0); n += xb.size(0)
    print(f"epoch {epoch:02d} loss {total / n:.4f}")

# 保存与加载：只存参数，不存整个对象
torch.save({"model": model.state_dict()}, "ckpt.pt")
model.load_state_dict(torch.load("ckpt.pt", map_location=device)["model"])

model.eval()
with torch.no_grad():
    print(model(x[:4].to(device)))

要点：zero_grad(set_to_none=True) 比默认写法更省内存；梯度裁剪放在 backward 之后、step 之前；保存用 state_dict 而不是整个模型对象，便于跨版本与跨设备加载。

示例二：多卡训练用 torchrun 启动 DDP

# train_ddp.py
import os, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")

model = build_model().to(local_rank)          # 你的模型
model = DDP(model, device_ids=[local_rank])

# 关键：分布式采样器保证每个进程读到不同分片
sampler = torch.utils.data.DistributedSampler(train_ds, shuffle=True)
dl = DataLoader(train_ds, batch_size=32, sampler=sampler, num_workers=4)

for epoch in range(EPOCHS):
    sampler.set_epoch(epoch)                  # 不加这句，每个 epoch 的分片顺序一样
    for xb, yb in dl:
        xb, yb = xb.to(local_rank), yb.to(local_rank)
        out = model(xb)
        loss = lossf(out, yb)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()

# 启动（单机 4 卡）
# torchrun --standalone --nproc_per_node=4 train_ddp.py

DDP 的核心是「每张卡跑一份模型、梯度做全规约平均」，因此学习率要随全局 batch 增长做大致的等比放大；同时要注意只有 rank 0 才写日志与保存 checkpoint，否则多进程会同时写同一个文件。

适合与不适合的场景

适合：研究与工程混合的团队，需要快速试结构、快速定位数值问题；需要庞大生态（视觉、语音、生成模型、量化、分布式训练库）的场景；从训练到推理一条链路打通的场景。

不适合：模型结构完全固定、只需要极致推理吞吐的服务（这类场景交给专门的推理引擎更合适）；强依赖静态图优化与整图部署到边缘设备的场景；只有 CPU 且需要极致推理性能的场景。