Python+UNet图像分割实战：从数据掩膜到IoU推理排错

简介：面向具备一定Python与深度学习基础、希望用U-Net完成图像分割实践的开发者与学习者，这份压缩包围绕分割任务的数据准备、模型搭建、损失函数选择、优化器配置、训练验证与推理预测展开，并聚焦大尺寸图像分块预测后的平滑融合处理，可用于卫星图像等场景的效果验证。包内共21个文件，以jpg、png示例图与前后对比gif为主，配有核心Python脚本、README说明、requirements依赖清单及许可证等，压缩包约5.6MB，结构轻量，便于快速查阅与复现实验。已有10164人学习下载，说明其流程与示例具有一定参考价值。读者可从中获得U-Net分割的实践脉络、分块预测与平滑拼接的实现思路，以及借助示例图对比分割前后的直观参考，适合用来搭建基线、理解关键环节或迁移到自有数据集。

1. 从一张标注图说起：Python + UNet 图像分割真正的卡点在哪

拿到一份 Python 图像分割的 UNet 代码包，很多人的第一反应是装完依赖直接python train.py，十几分钟后盯着纹丝不动的 loss 曲线怀疑显卡。这种场景我见过太多次：代码本身没毛病，卡住的地方几乎都在数据和指标这两头。UNet 那套编码器—解码器加跳跃连接的结构十年来没怎么变过，PyTorch 里两百行就写得完，真正决定你能否跑出可用掩膜的是掩膜本身怎么存、几何增强怎么同步、IoU 的分子分母怎么数。这份资源面向手里已经有标注图、想把分割模型训到能出图的人，医学影像、遥感地块、工业缺陷检测都吃得下，也适合当作第一个完整的语义分割工程来练手。

2. 数据先过关：掩膜格式、目录结构与 Dataset 的写法

模型训不动的时候，八成不是网络写错了，而是喂进去的掩膜根本不是你以为的那张图。这一章把数据侧的三个高频问题拆开：标注导出的掩膜有多少种形态、目录和文件名怎么对齐、增强为什么必须让图像和掩膜走同一条随机种子。

2.1 标注工具的导出五花八门，先统一成单通道索引图

同一个分割任务，不同人标注出来的掩膜能凑出四五种格式。常见的有：二值图（背景 0、前景 255）、调色板 PNG（P 模式，像素值是调色板索引）、彩色 PNG（每个类别一个 RGB 三元组）、以及 LabelMe 这类工具导出的 JSON 多边形。训练时网络要的是形状为(H, W)的整型张量，每个像素值直接就是类别编号，背景为 0、类别从 1 开始递增。

这里有两个必须提前定死的前提。第一，如果用CrossEntropyLoss，目标里的每个值都必须落在[0, num_classes-1]区间内，ignore_index一般设 255，所以有效区域里不能出现 255。第二，类别编号必须连续，从 0 开始一个个往上排，中间空一个编号不会报错，但那个类别永远学不出来，因为采样时根本抽不到。

我一般会先跑一个转换脚本，把所有掩膜压成统一的索引图，再进训练流程。转换脚本本身不复杂，坑都在细节里：

import os import numpy as np from PIL import Image # 类别名 -> RGB 的映射，顺序决定索引编号，索引 0 留给背景 PALETTE = [ (0, 0, 0), # 0 背景 (255, 0, 0), # 1 类别 A (0, 255, 0), # 2 类别 B (0, 0, 255), # 3 类别 C ] def to_index_mask(src_path, dst_path): """把彩色 / 调色板 / 二值掩膜统一成单通道 uint8 索引图""" # 关键：先 convert("RGB")，P 模式直接 np.array 读到的是调色板索引而非颜色 m = np.array(Image.open(src_path).convert("RGB")) out = np.zeros(m.shape[:2], dtype=np.uint8) # 默认整张图是背景 for idx, rgb in enumerate(PALETTE[1:], start=1): # 精确等值匹配，不要用容差，容差会把相邻类别吞掉 hit = np.all(m == np.array(rgb, dtype=np.uint8), axis=-1) out[hit] = idx # 保存前确认索引范围，越界说明调色板没覆盖全 assert out.max() < len(PALETTE), f"{src_path} 出现未定义的颜色" Image.fromarray(out, mode="L").save(dst_path) if __name__ == "__main__": for split in ("train", "val"): src_dir = f"raw/masks/{split}" dst_dir = f"dataset/masks/{split}" os.makedirs(dst_dir, exist_ok=True) for name in os.listdir(src_dir): if name.lower().endswith((".png", ".jpg", ".bmp")): to_index_mask(os.path.join(src_dir, name), os.path.join(dst_dir, name))

这段脚本有两个地方值得单独说。convert("RGB")那行看着多余，实际上是最容易翻车的一步：P 模式的 PNG 用np.array读出来是 0~N 的索引值，跟你在图片查看器里看到的颜色完全是两码事，不做转换直接匹配 RGB 会一个像素都对不上，最后得到一张全黑的掩膜，训练时 loss 却照常下降——因为模型学会了永远预测背景。mode="L"那行是保证输出是单通道，不写的话某些版本的 PIL 会按输入模式回写。

assert那句是给自己留的后悔药。真实标注里经常出现边缘抗锯齿产生的中间色，比如红绿边界上出现(128, 128, 0)，精确匹配匹配不上，这些像素就落到背景里去了。跑完脚本统计一下每类的像素占比，如果某个类别突然掉到 0.01% 以下，多半就是颜色没对上。

最后一句提醒：不要用mask = mask / 255这种写法处理二值图。前景 255 除以 255 变成 1.0，看着没错，但背景 0 除以 255 还是 0，模型输出只有一个通道，BCEWithLogitsLoss能吃，换成CrossEntropyLoss就直接报维度错误。二值任务和单类别多类任务在张量形状上不是一回事，这一点在第 3 章会再展开。

2.2 目录结构、文件名对齐与 Dataset 实现

数据目录我习惯固定成这样，训练脚本只认这一个约定，换数据集只改根路径：

dataset/ ├── images/ │ ├── train/ a001.png a002.png ... │ └── val/ b001.png b002.png ... └── masks/ ├── train/ a001.png a002.png ... └── val/ b001.png b002.png ...

图像和掩膜靠「去扩展名后的文件名」一一对应。这里有个特别隐蔽的问题：很多人用sorted(os.listdir(img_dir))和sorted(os.listdir(mask_dir))分别取两个列表，靠下标对齐。一旦某个目录里混进一个多余文件，或者两个目录的文件名前缀不一致导致排序结果不同（img_2排在img_10后面这种），后面所有样本的图像和掩膜就整体错位了。模型照样能训，loss 照样降，只是学的是「把 A 图的掩膜往 B 图上贴」，最后验证集指标难看得莫名其妙。

稳妥做法是求交集再排序：

import os import torch import numpy as np from PIL import Image from torch.utils.data import Dataset import albumentations as A from albumentations.pytorch import ToTensorV2 class SegDataset(Dataset): def __init__(self, root, split="train", size=512, train_mode=True): self.img_dir = os.path.join(root, "images", split) self.mask_dir = os.path.join(root, "masks", split) self.size = size self.train_mode = train_mode # 用集合求交集，两个目录对不上的文件直接丢弃而不是静默错位 img_stems = {os.path.splitext(f)[0] for f in os.listdir(self.img_dir)} mask_stems = {os.path.splitext(f)[0] for f in os.listdir(self.mask_dir)} self.stems = sorted(img_stems & mask_stems) if len(self.stems) == 0: raise RuntimeError("图像与掩膜文件名没有交集，检查命名规则") # 训练和验证用不同的增强管线，但归一化参数必须一致 common = [A.Resize(size, size)] if train_mode: self.tf = A.Compose( common + [ A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.2), A.RandomRotate90(p=0.5), A.RandomBrightnessContrast(p=0.3), A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ] ) else: self.tf = A.Compose(common + [ A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ]) def __len__(self): return len(self.stems) def __getitem__(self, i): stem = self.stems[i] img = np.array(Image.open( os.path.join(self.img_dir, stem + ".png")).convert("RGB")) mask = np.array(Image.open( os.path.join(self.mask_dir, stem + ".png"))) # albumentations 的 mask 参数只接受几何变换，颜色变换不会作用到 mask out = self.tf(image=img, mask=mask) img_t = out["image"] # (3, H, W) float32 mask_t = out["mask"].long() # (H, W) int64 assert mask_t.dim() == 2, f"掩膜维度异常: {mask_t.shape}" return img_t, mask_t

几个参数和实现细节。A.Normalize用了 ImageNet 的均值和方差，是因为后面大概率要加载预训练编码器，归一化方式必须和预训练时保持一致，否则特征分布对不上，微调效果会明显变差。mask_t.long()这一步不能省，CrossEntropyLoss要求目标为int64，uint8 直接喂进去会报类型错误，而且报错信息指向的是 loss 那一行，容易误判成模型输出维度的问题。

img_stems & mask_stems求交集看起来保守，实际上是主动暴露问题的做法。训练正式开始前打印一下len(img_stems) - len(self.stems)，如果不是 0，说明有文件命名不一致，先修数据再谈模型。

2.3 几何增强必须同步，颜色增强绝不能碰掩膜

数据增强在分割任务里比分类任务危险得多。分类任务里你把图翻转、裁剪、调色，标签还是那个标签；分割任务里图像做了水平翻转，掩膜必须跟着翻同样的角度，差一个像素都不行。albumentations 的设计就是所有空间变换同时作用于image和mask两个参数，颜色类变换只作用在image上，这是它比手工写 torchvision transform 省事的地方。

但即便用了正确的库，仍有三个地方要盯住。

第一是插值方式。A.Resize默认对掩膜使用最近邻插值，这是对的；如果某天为了「让掩膜看起来平滑」手动把插值改成双线性，类别索引会被插出 0.5、1.7 这种小数值，再取整就凭空制造出一堆不存在的类别边界。我一般会在 Compose 里显式写明mask_interpolation=cv2.INTER_NEAREST，把意图固定下来。

第二是随机裁剪。A.RandomCrop对稠密预测任务的默认行为可能裁出一块完全不含前景的补丁，如果这类样本占比过高，模型会偏向全预测背景。按 0.1~0.2 的比例使用、或者干脆不用裁剪只在Resize上做文章，对小目标任务更友好。

第三是训练和验证的归一化必须完全相同。我见过有人在验证集上忘了写A.Normalize，模型在训练时吃的输入是零均值单位方差的，验证时吃的是 0~255 的原始像素，指标直接掉二十个点，还以为是过拟合。

提示：A.Normalize之后必须紧跟ToTensorV2()，顺序反了会把归一化作用到已经变成(3, H, W)的张量上，维度对不上直接抛异常。

3. UNet 结构、损失函数与训练配置的配对关系

网络本身几乎是整个工程里最不需要动脑子的部分，原版 UNet 的结构十年没变。真正需要想清楚的是三件配套的事：输出通道数跟损失函数怎么配、类别不平衡怎么处理、batch size 跟归一化层怎么搭。这三件事配错，网络写得再漂亮也训不出东西。

3.1 编码器—解码器与跳跃连接的最小实现

UNet 的结构一句话能说完：左边四次下采样把感受野做大，右边四次上采样把分辨率还原，每一级上采样之前把左边同分辨率的特征图按通道拼过来。拼接这个动作就是所谓的跳跃连接，作用是让解码器在还原边界时能直接拿到浅层的高频细节，而不必从低分辨率特征里硬猜。

原版通道数是 64-128-256-512-1024，显存吃紧的时候按比例砍一半到 32-64-128-256-512，精度损失通常在 1~2 个 IoU 点以内，对验证流程来说完全可以接受。

import torch import torch.nn as nn class DoubleConv(nn.Module): """两层 3x3 卷积 + BN + ReLU，UNet 的基本单元""" def __init__(self, in_ch, out_ch): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x) class Down(nn.Module): """最大池化下采样 + 双层卷积""" def __init__(self, in_ch, out_ch): super().__init__() self.block = nn.Sequential(nn.MaxPool2d(2), DoubleConv(in_ch, out_ch)) def forward(self, x): return self.block(x) class Up(nn.Module): """转置卷积上采样 + 与浅层特征拼接 + 双层卷积""" def __init__(self, in_ch, out_ch): super().__init__() self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride=2) self.conv = DoubleConv(in_ch, out_ch) def forward(self, x, skip): x = self.up(x) # 输入尺寸不是 16 的整数倍时，上采样后与 skip 可能差 1 个像素 if x.shape[-2:] != skip.shape[-2:]: x = nn.functional.interpolate( x, size=skip.shape[-2:], mode="bilinear", align_corners=False) x = torch.cat([skip, x], dim=1) return self.conv(x) class UNet(nn.Module): def __init__(self, in_ch=3, num_classes=2, base=32): super().__init__() c = [base, base * 2, base * 4, base * 8, base * 16] self.inc = DoubleConv(in_ch, c[0]) self.down1 = Down(c[0], c[1]) self.down2 = Down(c[1], c[2]) self.down3 = Down(c[2], c[3]) self.down4 = Down(c[3], c[4]) self.up1 = Up(c[4], c[3]) self.up2 = Up(c[3], c[2]) self.up3 = Up(c[2], c[1]) self.up4 = Up(c[1], c[0]) # 输出层不加 softmax/sigmoid，交给损失函数内部处理 self.outc = nn.Conv2d(c[0], num_classes, 1) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) return self.outc(x) # (B, num_classes, H, W) 的 logits

Up里那段尺寸对齐的代码不是凑数的。UNet 要求输入边长是 16 的整数倍（四次下采样），输入 512 没问题，输入 513 就会在下采样过程中丢掉一个像素，上采样回来接不上。这份代码用interpolate兜住，代价是那一路会有轻微的重采样损失。长期看，在 Dataset 里统一Resize到 2 的幂次更干净。

outc那行用 1×1 卷积而不是 3×3，是因为这一步只做通道映射不做空间聚合。整个网络参数量在base=32时大约是 7.8M，单张 512×512 图在 8GB 显存上 batch 开到 8 是够的。

3.2 输出通道、激活函数与损失的配对

这一节是新手翻车最集中的地方，三种组合必须记清楚：

核心原则只有一条：损失函数内部已经带了log_softmax/sigmoid，模型输出层就不要再加一遍。模型里加了 softmax、外面再接CrossEntropyLoss，等于做了两次 softmax，梯度会被压得非常小，训练速度肉眼可辨地变慢，最后还训不出效果。

类别不平衡在分割任务里是常态。一片 512×512 的遥感图里，待检测的建筑可能只占 2% 的像素，纯CrossEntropyLoss下模型十几轮就学会全预测背景，loss 降到 0.05，IoU 是 0。常见做法是把 Dice 和 CE 按 1:1 加权：

import torch import torch.nn as nn import torch.nn.functional as F class DiceLoss(nn.Module): def __init__(self, num_classes, smooth=1.0, ignore_index=255): super().__init__() self.num_classes = num_classes self.smooth = smooth self.ignore_index = ignore_index def forward(self, logits, target): # logits: (B, C, H, W)；target: (B, H, W) prob = torch.softmax(logits, dim=1) # 概率化 valid = (target != self.ignore_index) target_clean = target.clone() target_clean[~valid] = 0 # one-hot 后把无效像素也置零，避免它们参与分子的累加 onehot = F.one_hot(target_clean, self.num_classes).permute(0, 3, 1, 2).float() onehot = onehot * valid.unsqueeze(1).float() prob = prob * valid.unsqueeze(1).float() dims = (0, 2, 3) # 在 batch 和空间维求和 inter = (prob * onehot).sum(dims) union = prob.sum(dims) + onehot.sum(dims) dice = (2 * inter + self.smooth) / (union + self.smooth) # 只对当前 batch 里真实出现过的类别取平均，避免小类别被稀释 present = onehot.sum(dims) > 0 return 1 - (dice * present).sum() / present.sum().clamp(min=1) class ComboLoss(nn.Module): def __init__(self, num_classes, ce_weight=1.0, dice_weight=1.0): super().__init__() self.ce = nn.CrossEntropyLoss(ignore_index=255) self.dice = DiceLoss(num_classes) self.ce_w = ce_weight self.dice_w = dice_weight def forward(self, logits, target): return (self.ce_w * self.ce(logits, target) + self.dice_w * self.dice(logits, target))

smooth取 1.0 是为了防止某类别当轮完全没出现时除零。present那一行是这份实现跟多数公开代码不同的地方：公开实现通常对所有类别一视同仁地求平均，包括背景。背景像素占比动辄 90% 以上，它的 Dice 轻松到 0.99，把平均值拉得很高，看起来指标漂亮，前景实际上没学会。只统计当前 batch 出现过的类别，指标会更低但更真实。

3.3 优化器、学习率与混合精度的搭配

优化器用AdamW，学习率 1e-4，weight decay 1e-4，betas=(0.9, 0.999)默认即可。从零训练时 1e-3 也能跑，但配合 BN 和小 batch 容易在前期震荡；加载预训练编码器时必须降到 1e-4 甚至 1e-5，编码器部分的学习率再乘 0.1，防止预训练特征被一开始就冲掉。

调度器优先CosineAnnealingLR配个warmup。前 3 个 epoch 线性爬升到目标学习率，之后余弦下降到底。warmup 在 batch 小的时候作用明显，因为 BN 统计量在前几十步还不准，大学习率会让参数往奇怪的方向跑。

混合精度训练能省将近一半显存，torch.cuda.amp三行就接进去了。但有两个细节：Dice Loss 里的除法和clamp在 fp16 下容易出nan，我一般把损失计算放进autocast(enabled=False)的上下文里用 fp32 算；GradScaler的state_dict在断点续训时必须一起存，不存的话恢复训练的第一个 step 大概率梯度溢出被跳过。

from torch.cuda.amp import autocast, GradScaler model = UNet(num_classes=4).cuda() opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=60) scaler = GradScaler() criterion = ComboLoss(num_classes=4).cuda() for epoch in range(60): model.train() for img, mask in train_loader: img, mask = img.cuda(), mask.cuda() opt.zero_grad(set_to_none=True) with autocast(): logits = model(img) # 损失在 fp32 下算，避开 Dice 的除零和溢出 with autocast(enabled=False): loss = criterion(logits.float(), mask) scaler.scale(loss).backward() scaler.step(opt) scaler.update() sched.step()

set_to_none=True比zero_grad()省一次显存写入，在显存吃紧时有用。autocast只包住前向，损失单独拿出来，这个写法比较啰嗦但对分割任务最实用。

4. 训练循环、指标计算与断点续训

网络和数据都摆好了，接下来是把它们串成一条能跑、能看、能接着跑的训练管线。这一章的重点不在「能跑起来」，而在「跑出来的数字可信」——指标算错比模型训坏更麻烦，因为你看不出来。

4.1 一轮完整训练循环的骨架

循环本身不长，值得说的是几个位置放什么。我一般会在每个 epoch 结束时只在验证集上算指标、只在指标提升时保存权重，训练集的 loss 只用于早期判断是否收敛。

import os import torch from tqdm import tqdm def train_one_epoch(model, loader, optimizer, criterion, scaler, device): model.train() total_loss = 0.0 for img, mask in tqdm(loader, desc="train", leave=False): img = img.to(device, non_blocking=True) mask = mask.to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) with torch.cuda.amp.autocast(): logits = model(img) with torch.cuda.amp.autocast(enabled=False): loss = criterion(logits.float(), mask) scaler.scale(loss).backward() # 梯度裁剪要在 scaler.step 之前、unscale 之后 scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() total_loss += loss.item() * img.size(0) return total_loss / len(loader.dataset) def fit(model, train_loader, val_loader, epochs, ckpt_dir, device): opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs) scaler = torch.cuda.amp.GradScaler() criterion = ComboLoss(num_classes=model.outc.out_channels).to(device) best = 0.0 os.makedirs(ckpt_dir, exist_ok=True) for ep in range(epochs): tr_loss = train_one_epoch(model, train_loader, opt, criterion, scaler, device) val_iou, val_dice = evaluate(model, val_loader, device) sched.step() print(f"ep {ep:03d} loss {tr_loss:.4f} " f"mIoU {val_iou:.4f} Dice {val_dice:.4f}") state = { "model": model.state_dict(), "optimizer": opt.state_dict(), "scheduler": sched.state_dict(), "scaler": scaler.state_dict(), # 断点续训必须带上 "epoch": ep, "best": best, } torch.save(state, os.path.join(ckpt_dir, "last.pth")) if val_iou > best: best = val_iou torch.save(state, os.path.join(ckpt_dir, "best.pth"))

clip_grad_norm_的位置有讲究。在scaler.scale(loss).backward()之后、scaler.step()之前调用，并且要先scaler.unscale_()，否则裁剪的是被放大了 65536 倍的梯度，阈值 1.0 等于没裁。这个顺序错了不会报错，只是梯度裁剪完全失效，某几个 batch 梯度爆炸导致 loss 变nan的时候你会一头雾水。

loss.item() * img.size(0)按样本数加权再除以总样本数，这样最后一个不满的 batch 不会把平均值拉偏。直接对 batch loss 求算术平均是常见写法，样本数能被 batch size 整除时没区别，不能整除时会有几个千分点的误差。

4.2 IoU 与 Dice 的分子分母，别在 batch 层面求平均

分割指标最容易犯的错是「每个 batch 算一次 IoU 然后取平均」。实际上每个 batch 的前景像素数量差异巨大，一张图里前景占 30%、下一张占 0.5%，两者的 IoU 直接算术平均等于给小目标过高的权重，最终数字跟真实性能对不上。

正确做法是累加整个验证集的混淆矩阵，最后一次性算出来：

import torch import numpy as np @torch.no_grad() def evaluate(model, loader, device, num_classes=4, ignore_index=255): model.eval() # 混淆矩阵：行是真值，列是预测 conf = np.zeros((num_classes, num_classes), dtype=np.int64) for img, mask in loader: img = img.to(device) logits = model(img) pred = logits.argmax(dim=1).cpu().numpy() gt = mask.numpy() valid = gt != ignore_index # bincount 一次性统计 (gt * C + pred) 的频次，比双层 for 快很多 idx = (gt[valid].astype(np.int64) * num_classes + pred[valid].astype(np.int64)) cnt = np.bincount(idx, minlength=num_classes ** 2) conf += cnt.reshape(num_classes, num_classes) # 每类 IoU = TP / (TP + FP + FN)，背景类索引 0 通常不计入 mIoU inter = np.diag(conf) union = conf.sum(1) + conf.sum(0) - inter iou = inter / np.maximum(union, 1) fg = iou[1:] # 去掉背景 # Dice = 2TP / (2TP + FP + FN) dice = 2 * inter / np.maximum(conf.sum(1) + conf.sum(0), 1) return float(fg.mean()), float(dice[1:].mean())

np.bincount那个技巧是为了避免对每个像素做 Python 循环，整个验证集几百张图跑下来差别在秒级和分钟级。ignore_index的像素在累加前就被过滤掉了，这一点很重要——如果掩膜里有 255 的「不确定区域」，把它们当背景算进去会把 IoU 压得很低，而且这个偏差会随着忽略区域的比例变化。

指标打印出来后，如果某一类的 IoU 明显低于其他类，先去看那一类的像素占比和标注质量，不要急着调超参。我遇到过的多数「某类学不出来」，最后查出来是那一类的标注只有几十个样本，或者标注标准本身有分歧。

4.3 日志、随机种子与断点续训

可复现这件事在调试阶段价值极高。torch.manual_seed、np.random.seed、random.seed三个都要设，cudnn.benchmark=True会让卷积算法选择随机器状态变化，追求完全可复现的话要关掉并设deterministic=True，代价是慢 10% 左右。

import random import numpy as np import torch def set_seed(seed=42, deterministic=True): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) if deterministic: torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def resume(path, model, optimizer=None, scheduler=None, scaler=None): ckpt = torch.load(path, map_location="cpu") model.load_state_dict(ckpt["model"]) if optimizer is not None: optimizer.load_state_dict(ckpt["optimizer"]) if scheduler is not None: scheduler.load_state_dict(ckpt["scheduler"]) if scaler is not None and "scaler" in ckpt: scaler.load_state_dict(ckpt["scaler"]) # 返回时把 epoch + 1，避免恢复后重跑同一个 epoch return ckpt["epoch"] + 1, ckpt.get("best", 0.0)

日志建议按 JSON Lines 追加写，一个 epoch 一行，包含 epoch、train_loss、val_iou、val_dice、lr。用 pandas 读进来画曲线，比从终端往回翻要省事得多。断点续训时resume返回epoch + 1这个细节，忘了加一的话恢复后会把已经训过的那个 epoch 再跑一遍，学习率调度会整体错位一格。

5. 常见问题排查：五个把训练拖死的现场

下面这五条都是我在别人的代码包和自己写崩的代码里反复见过的。共同特点是：不报错，或者报的错跟真实原因隔了一层，光看 traceback 找不着北。

5.1 loss 降到 0.05、掩膜全黑

现象：训练日志里 loss 一路往下走得很顺，二十轮之后稳定在 0.05 附近，但拿验证图做推理，输出整张图全是背景。

原因：前景像素占比太低，纯CrossEntropyLoss的全局最优解就是全预测背景。这时候 loss 曲线看起来极其健康，指标却是 0。

解决：先统计训练集里前景像素占比，低于 5% 就把损失换成CE + Dice的组合，Dice 直接优化重叠度，对不平衡不敏感。另外看evaluate输出的 mIoU 是不是真的在动，不要只盯 loss。如果训了一百轮 mIoU 还是 0，检查掩膜转换脚本，前面 2.1 提到的convert("RGB")漏写就正好是这个症状。

5.2 CUDA out of memory，减小 batch 也没用

现象：把 batch size 从 8 降到 4 再降到 1，仍然 OOM，显存占用不随 batch 线性下降。

原因：分辨率是显存占用的平方项。512×512 换成 1024×1024，特征图面积翻四倍，batch 从 8 降到 1 只省了八分之一，抵不过分辨率带来的四倍增长。另一个可能是有张量被意外地挂在了计算图上，比如在with torch.no_grad()外面做验证。

解决：先降分辨率再降 batch，两个一起调。验证阶段一定要包@torch.no_grad()。如果显存还是不够，把base从 64 砍到 32，再不行把up里的转置卷积换成interpolate + conv，能再省一点峰值显存。混合精度开了之后还没缓解，检查一下损失计算是不是被漏在autocast外面用 fp32 跑了一整套。

5.3 报错 Expected more than 1 value per channel

现象：训练跑到一半突然抛ValueError: Expected more than 1 value per channel when training, got input size [1, 64, 1, 1]，前面的 epoch 都正常。

原因：数据集样本数不能被 batch size 整除，最后一个 batch 只剩 1 张图，BatchNorm2d在训练模式下需要每个通道至少两个样本才能算方差。偶数样本数的时候碰不到，换了个数据集就冒出来。

解决：DataLoader加drop_last=True，只在训练集上加，验证集不能加否则会漏样本。或者把 BN 换成GroupNorm，对小 batch 更友好，代价是收敛速度略慢一点。

5.4 验证指标很高，可视化一看全糊

现象：mIoU 报出 0.88，但把预测结果叠到原图上看，边界一塌糊涂，小目标整片丢掉。

原因：几乎可以确定是类别像素占比极不均衡，背景占了 90% 以上。mIoU 对背景类不敏感，前景学得再差，背景的高分也能把平均值撑起来。

解决：把背景类从 mIoU 里排除掉（4.2 的代码已经这么做了），单独看各类的 IoU。再对比一下训练集和验证集的分布，如果验证集里小目标占比高很多，指标低是正常的，说明划分本身不均衡。真正的修法是分层的划分策略，按前景像素占比排序后再切分。

5.5 推理结果和验证结果对不上

现象：evaluate里跑出来的 mIoU 是 0.85，把同一批图单独写个推理脚本读进来跑，只有 0.4。

原因：两个脚本的预处理不一致。最常见的是验证时用了Net.eval()但推理脚本忘了，BN 用 batch 统计量而不是运行统计量；或者是推理脚本里没做Normalize，输入还是 0~255；还有一种是Resize的尺寸跟训练时不一样，模型能吃但特征分布已经偏了。

解决：把预处理和模型构建抽成一个build_model()/preprocess()函数，训练、验证、推理三处都调同一个，不要复制粘贴。写完之后拿一张训练集里的图做单图过拟合测试：只用这一张图训 200 步，模型应该能把它预测得几乎完美，做不到就说明管线某处有问题，先别继续往下跑。

注意：单图过拟合测试的收敛时间通常在半分钟以内，如果跑了五分钟还没贴合，一定是数据到模型之间断了，不用怀疑模型容量。

6. 推理与验证：大图滑窗和交付前的强制自检

训练完之后最容易被跳过的是推理环节，很多人直接把验证集的evaluate结果当成最终交付质量。但真实场景里的图往往比训练图大得多，直接整张塞进网络会 OOM，缩放到 512 又会把细节抹掉。工业上常见的做法是滑窗推理：把大图切成带重叠的块，逐块预测，再把结果拼回去，重叠区域取平均。

import numpy as np import torch import torch.nn.functional as F @torch.no_grad() def sliding_window_infer(model, image, device, tile=512, stride=384, num_classes=4): """image: (3, H, W) 的 float32 numpy，已做归一化""" model.eval() _, H, W = image.shape # 累积概率和计数图，重叠区域最后取平均 prob_sum = np.zeros((num_classes, H, W), dtype=np.float32) count = np.zeros((H, W), dtype=np.float32) ys = list(range(0, max(H - tile, 0) + 1, stride)) xs = list(range(0, max(W - tile, 0) + 1, stride)) # 补上右下角最后一块，避免边缘留白 if ys[-1] + tile < H: ys.append(H - tile) if xs[-1] + tile < W: xs.append(W - tile) for y in ys: for x in xs: patch = torch.from_numpy( image[:, y:y + tile, x:x + tile]).unsqueeze(0).to(device) logits = model(patch) # 边缘块尺寸不足时插值回 tile 大小再累加 logits = F.interpolate(logits, size=(tile, tile), mode="bilinear", align_corners=False) prob = torch.softmax(logits, dim=1)[0].cpu().numpy() prob_sum[:, y:y + tile, x:x + tile] += prob count[y:y + tile, x:x + tile] += 1 return (prob_sum / np.maximum(count, 1e-6)).argmax(axis=0)

stride取tile的 3/4，重叠 25%，这个比例能覆盖大部分拼接缝，再高收益递减、耗时翻倍。ys[-1] + tile < H那个补丁很关键，不补的话右边和下边会缺一条，缺的宽度正好是(H - tile) % stride，很窄但正好落在你最关心的边界上。

参数上还有两个地方值得调。tile不要超过训练时的输入尺寸，模型没见过更大的感受野，硬上会掉点；num_classes必须跟训练时完全一致，模型输出通道数对不上直接报维度错。

交付前我会强制走一遍三项自检，缺一项都不敢发出去：

第一项是排查管线断裂，第二项是排查「训练好推理差」，第三项是排查小目标召回。三项里第二项最容易出问题，因为它是靠人手动保证的，而人最容易在赶时间的时候复制粘贴。

说到赶时间，有一次我在一个工业缺陷项目上赶着交付，验证集 mIoU 0.86 就直接打包了，结果对方拿产线上的图跑，缺陷几乎全漏。回头查了两天才发现推理脚本里那行A.Normalize被注释掉了——调试可视化的时候注释的，忘了恢复。从那以后我每次交付前都强制走一遍单图过拟合测试和预处理一致性检查，宁可多花二十分钟，也不愿意再体验一次那种「指标漂亮、实际全废」的滋味。希望帮到你。

本文还有配套的精品资源，点击获取