4.2 用一张卡跑通训练环

4.2 用一张卡跑通训练环

本节摘要：分割训练环是「成对批次 → 前向出分数图 → 损失 → 反向 → 优化器步进 → 定期验证与存盘」。本教程走动态图风格：模块定义清晰、循环写在明处。静态图框架能做同一件事，但不在这里做语法对照。先过拟合一两张图，证明环没有写反，再上全数据集。

先说结论

阅读完本节，你应当能够：

1. 写出数据集接口：返回图像张量与整数掩码，形状约定固定
2. 写出训练一步：清梯度、前向、损失、反向、步进
3. 在单图上把损失打到接近零，作为环的验收
4. 保存与加载权重后，推理预处理与训练一致

一、选定一条路走到底

原文并列介绍两种流行框架：一种偏静态图与高层顺序接口，一种偏动态图与模块类。对照它们的张量与自动求导，是另一门课的主线。本课要的是掩码能训练。我选动态图写法：继承模块类，forward 里摆层，循环里自己调。这样损失换成 Dice、掩码要忽略 255、多尺度输入，改起来都在明处。

硬件：一张消费级卡够小 U-Net。显存不够就减边长、减批次、开混合精度。没有卡就用 128 边长、批次 1、CPU 验环。环的正确性与是否上集群无关。

for images, masks in loader:
    optimizer.zero_grad()
    logits = model(images)
    loss = criterion(logits, masks)
    loss.backward()
    optimizer.step()
​

原文简化 U-Net 示例用二值交叉熵和 Adam、学习率 0.001，循环结构正是上面五步。真实模型只是把两层卷积换成有跳跃的结构，环不变。

二、数据类型与设备是第一坑

图像浮点、掩码长整型（交叉熵）或浮点独热（Dice）。通道维在第二维还是最后一维，必须和卷积约定一致。图像与掩码要在同一设备。忽略索引要在损失构造器里声明。批次维为 0。这些写错，报错有时来得晚：广播出一个能反传的标量，数值却在学垃圾。

加载器的工人进程数在 Windows 上要小心，调试阶段先设 0。随机种子在可复现实验里固定，但完全确定还受算法实现影响，不必迷信到小数点后六位。

验证环：model.eval() 与关掉梯度。漏掉 eval，批归一仍用训练统计，验证曲线会抖。存盘存状态字典，也存类别数、均值方差、输入边长。上线缺这三项，权重等于半份。

⚠️ 常见坑：掩码少一个通道维或类别维，损失在错误的轴上 softmax。打印 logits.shape 与 masks.shape 应成为肌肉记忆。
💡 关键直觉：单图损失不降，是环或标签反了；全数据损失降、验证不降，才是泛化问题。先排除前者。

三、单图过拟合验收

取第 2.4 节试标里最干净的一张。学习率可略大，训几百步。输出叠图应逼近真值。做不到：检查是否上下翻转、是否类别 id 错位、是否把 RGB 当 BGR 读入导致颜色分布怪、是否损失忽略了所有前景。做到了：环通了，再把加载器换成全训练集，学习率调回稳妥值。

学习率 0.001 配 Adam 对从头训的小网常见。接预训练骨干往往要更小，骨干与头用不同组。下一节细说。本节只要求：你能改一处学习率看到曲线斜率变化，证明优化器真的在管这些参数。

静态图框架的「编译再 fit」也能过拟合单图。若团队已有那套管线，不必迁移。不要为了教程同时维护两套循环。选定后，第 5 章所有模型都往这一个环上插。

图 环上的五件事

四、日志里要有掩码可视化

只看损失标量会漏掉「前景全灭」。每隔若干步把预测 argmax 上色存下来。Windows 路径不要写进教程正文给读者当作业；在你的实验目录里固定一个可视化出口即可。指标计算放到下一节，环这一节先保证有图可看。

批次大小与边长对分割显存几乎线性。优先保边长（边界细节），批次不够再用梯度累积模拟。这是训练环的工程旋钮，不是模型论文的一部分，却常决定你能不能复现自己上周的实验。

五、Windows 上的加载细节

调试阶段把数据加载的并行工人数设为 0，排除多进程在 Windows 上的序列化与随机种子问题。环通了再加工人。内存里不要把整个数据集先读成一个巨大数组再切，按需读文件，掩码用无损格式。路径拼装用库的路径对象，不要在教程里写盘符作业——产品代码里集中一个数据根配置即可。

过拟合单图时把增强关掉。增强开着，单图也在变，损失曲线会抖，你会误判环没写对。全量阶段再打开第 6.1 节的菜单。混合精度能省显存，若出现 NaN，先关它定位，再考虑损失缩放。

存盘内容清单：模型状态、优化器若要续训、类别表、均值方差、边长、忽略 id、随机种子、代码版本或配置哈希。缺哈希，三个月后的「同样配置」无法证明。推理只加载模型与契约，不加载优化器。

问题：静态图的 fit 一行能不能用？

能。只要掩码形状、忽略索引、回调里的叠图和存盘你都看懂。本教程用显式循环，是为了改 Dice、改忽略、改多尺度时不跟高层封装打架。不是道德洁癖。选定后不要两套循环并存改到不一致。

打印每步的 logits 最小最大。长期全是很大的正或全是负，可能头初始化或学习率把分类推到饱和，梯度变小。单图过拟合仍应能拉开两类 logits。拉不开，查标签是不是全背景，或损失是不是在错误的轴上平均了两次。

六、单图验收清单

增强关闭、eval 关（训练模式）、学习率可略大、几百步、叠图逼近真值。做不到按顺序查：配对、翻转、BGR 与 RGB、忽略把前景全跳过、损失轴、设备不一致。做到了再开全量与增强。可视化每隔若干步存 argmax 上色，防止损失降但前景灭。

Windows 调试工人数 0。存盘带契约：均值方差、边长、类表、忽略、配置哈希。推理不加载优化器。混合精度 NaN 时先关。logits 最小最大长期饱和要查初始化。静态图 fit 能用，但不要两套循环并存。选定动态图后第 5 章只换模块。

批次与边长争显存时保边长，用累积补有效批次。日志记有效批次。环的单元测试是单图，不是全量碰巧能降。

单图清单逐步排查配对与损失轴。可视化防前景灭。工人数 0。存盘契约含哈希。有效批次入日志。两套循环禁止并存。logits 饱和查初始化。混合精度 NaN 先关。增强在单图阶段必须关，否则曲线抖会被误判为环写反。单图截图进合并请求，没有截图的全量实验视为未就绪，和第 4 章支柱页的纪律一致。

单图清单做成逐步脚本，每步打印是否通过。可视化目录按步数命名。工人数配置默认 0。哈希进文件名。有效批次计算公式写在日志第一行。禁止第二套循环的代码搜索。饱和检测阈值写入。单图关增强由配置强制，打不开全量开关除非单图标记通过。截图附件进合并请求模板。设备不一致用一次 .to 封装，散落的 .cuda 视为气味。掩码 dtype 与损失接口对照表贴在 criterion 构造处。续训才加载优化器，发布加载不含优化器。NaN 时自动关混合精度重跑单步，日志标明。这些是环的工程皮肤，没有皮肤，第 5 章模块插上也会从同样的孔漏掉。

一张卡上的最小环只做三件事：取一对、前向、反传并保存权重。设备、非阻塞拷贝、混合精度可以后加，但第一夜必须看到损失下降和单图叠图变干净。过拟合不了单图，全量数据只会让你更忙。冻结循环后禁止再开一套高层 fit，两套学习率与早停会在周报里各说各话。

温故知新

• 五步循环：清梯度、前向、损失、反向、步进
• 不写框架大战：一条动态图路径插到第 5 章
• 形状与设备：打印出来，别让广播装成训练
• 单图过拟合是环的单元测试
• eval 与配置一起存：推理必须能复现预处理
• 可视化预测：防止损失下降但前景全灭

下一节把循环里的标量换成分割真正该优化和该汇报的量：交叉熵、Dice、交并比、优化器脾气。