PyTorch 工程习惯
这一篇不讲算法,讲「让实验结果可信」的习惯。
同样是调模型,有人能快速定位问题,有人整天在重跑。差别就在这些工程细节上。
一、项目结构:从第一天就这样组织
my-project/
├── configs/ # YAML 配置,而不是散落的 argparse
│ └── base.yaml
├── data/ # 原始数据(不进 git)
├── src/
│ ├── dataset.py # Dataset / DataLoader
│ ├── model.py # 模型定义
│ ├── train.py # 训练入口
│ ├── evaluate.py # 评估
│ └── utils.py # 种子、日志、指标
├── outputs/ # 权重、日志、图表
├── requirements.txt
└── README.md # 怎么跑、结果多少、用了什么配置为什么不用一个 notebook 走到黑:notebook 的执行顺序依赖是隐形炸弹,且无法复现。约定:探索用 notebook,定型后立刻抽成 .py。
二、随机性:可复现的第一道门槛
python
import os, random, numpy as np, torch
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
# 严格确定性(会略慢)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
set_seed(42)现实
即使全部设种子,GPU 上仍可能有极小的非确定性(原子操作顺序)。目标不是比特级复现,而是「结论可复现」:换 3 个种子跑,看指标波动范围。
三、Dataset 与 DataLoader:数据管道的三个坑
python
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
class TextDataset(Dataset):
def __init__(self, df, tokenizer, max_len=256):
self.texts = df["text"].tolist()
self.labels = df["label"].tolist()
self.tok, self.max_len = tokenizer, max_len
def __len__(self):
return len(self.texts)
def __getitem__(self, i):
ids = self.tok(self.texts[i], truncation=True,
max_length=self.max_len)["input_ids"]
return torch.tensor(ids), torch.tensor(self.labels[i])
def collate(batch):
ids, labels = zip(*batch)
ids = pad_sequence(ids, batch_first=True, padding_value=0)
mask = (ids != 0).long()
return ids, torch.stack(labels), mask
dl = DataLoader(ds, batch_size=32, shuffle=True, num_workers=4,
pin_memory=True, collate_fn=collate,
drop_last=True) # 训练时丢弃最后一个不满 batch(BN 需要)| 坑 | 症状 | 解决 |
|---|---|---|
忘了 shuffle=True | loss 震荡、收敛慢 | 训练集必须 shuffle,验证集不 shuffle |
| padding 未屏蔽 | 指标偏低 | mask + 池化时乘 mask |
num_workers 太大 | 卡住/Debug 困难 | 起始 2–4,Windows 上设 0 |
| 自定义 collate 写错 | 形状不匹配报错 | 固定返回 (x, y, mask) 三元组 |
四、日志:不要只 print
python
# 轻量方案:自己写 CSV,够用且零依赖
import csv, time
log_path = "outputs/log.csv"
with open(log_path, "w", newline="") as f:
w = csv.writer(f); w.writerow(["epoch", "train_loss", "val_loss", "val_acc", "lr", "sec"])进阶用 TensorBoard 或 W&B:
python
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("outputs/runs/exp1")
writer.add_scalar("loss/train", tr_loss, epoch)
writer.add_scalar("loss/val", val_loss, epoch)
writer.add_scalar("lr", optimizer.param_groups[0]["lr"], epoch)
writer.close()必须记录的四件事:train_loss、val_loss、val_metric、lr。 必须记录的第二层:数据版本、代码 commit、完整超参。否则一周后你不知道哪个实验是什么。
五、保存与加载
python
# 保存:只存权重(推荐),附上元信息
torch.save({
"model": model.state_dict(),
"config": vars(args),
"epoch": epoch,
"metric": best_metric,
}, "outputs/best.pt")
# 加载
ckpt = torch.load("outputs/best.pt", map_location=device)
model.load_state_dict(ckpt["model"])
print("best metric:", ckpt["metric"])不要 torch.save(model)
保存整个模型对象会把类的引用路径写死,换目录/重构代码就加载失败。 只存 state_dict()。
六、显存与速度:常查的 5 个地方
python
# ① 当前显存
print(f"{torch.cuda.memory_allocated()/1e9:.2f} GB")
# ② 推理必须关梯度
@torch.no_grad()
def predict(model, x): ...
# ③ 及时释放
del xb, logits
torch.cuda.empty_cache()
# ④ 找瓶颈
with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
model(xb)
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))| 现象 | 常见原因 |
|---|---|
| 显存持续增长 | 把 loss/张量存进了 list;未 no_grad |
| GPU 利用率低 | 数据加载慢(num_workers 太小)、batch 太小 |
| 训练比推理慢很多 | 正常(要存激活值 + 反向) |
| 突然 OOM | 某个 batch 特别长 → 用长度分桶或固定 max_len |
七、实验管理:一张表管住所有实验
维护一个 experiments.md:
| id | 日期 | 改动 | 关键超参 | val_acc | 备注 |
|---|---|---|---|---|---|
| 001 | 10-01 | 基线 ResNet18 | lr=1e-3, bs=32 | 0.912 | — |
| 002 | 10-02 | + 更强增强 | 同上 | 0.928 | 有效 |
| 003 | 10-03 | 换 lr=3e-4 | bs=32 | 0.931 | 略好 |
| 004 | 10-04 | + label smoothing | 同上 | 0.933 | 保留 |
铁律:一次只改一个变量。 同时改三处,结果变好了你也不知道归因给谁。
八、从实验到交付
| 阶段 | 检查项 |
|---|---|
| 训练前 | 种子固定、数据划分固定、基线已记录 |
| 训练中 | 日志完整、checkpoint 定期存(防 Colab 掉线) |
| 训练后 | 最优权重已加载、测试集只跑一次 |
| 交付前 | 推理脚本独立可跑、预处理与训练完全一致、README 写清复现命令 |
最后一条最常被忽略:交付时别人拿到你的模型,第一步就是「跑一下看看」。推理脚本必须能独立运行,不依赖 notebook 的执行顺序。
九、一个可直接复用的推理模板
python
class Predictor:
def __init__(self, ckpt_path, device="cpu"):
self.device = device
ckpt = torch.load(ckpt_path, map_location=device)
self.cfg = ckpt["config"]
self.model = build_model(self.cfg).to(device)
self.model.load_state_dict(ckpt["model"])
self.model.eval() # ← 别忘
self.tf = build_transform(train=False) # ← 与训练一致
@torch.no_grad()
def __call__(self, image):
x = self.tf(image).unsqueeze(0).to(self.device)
prob = self.model(x).softmax(-1)[0]
idx = int(prob.argmax())
return {"label": idx, "confidence": float(prob[idx])}通过标准
下一步
L3 完成。进入 L4 大模型。