Skip to content

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=Trueloss 震荡、收敛慢训练集必须 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备注
00110-01基线 ResNet18lr=1e-3, bs=320.912—
00210-02+ 更强增强同上0.928有效
00310-03换 lr=3e-4bs=320.931略好
00410-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 大模型。

AI 学习路线 · 从入门到专家