简介面向具备一定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、调色板 PNGP 模式像素值是调色板索引、彩色 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], dtypenp.uint8) # 默认整张图是背景 for idx, rgb in enumerate(PALETTE[1:], start1): # 精确等值匹配不要用容差容差会把相邻类别吞掉 hit np.all(m np.array(rgb, dtypenp.uint8), axis-1) out[hit] idx # 保存前确认索引范围越界说明调色板没覆盖全 assert out.max() len(PALETTE), f{src_path} 出现未定义的颜色 Image.fromarray(out, modeL).save(dst_path) if __name__ __main__: for split in (train, val): src_dir fraw/masks/{split} dst_dir fdataset/masks/{split} os.makedirs(dst_dir, exist_okTrue) 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 却照常下降——因为模型学会了永远预测背景。modeL那行是保证输出是单通道不写的话某些版本的 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, splittrain, size512, train_modeTrue): 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(p0.5), A.VerticalFlip(p0.2), A.RandomRotate90(p0.5), A.RandomBrightnessContrast(p0.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(imageimg, maskmask) 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要求目标为int64uint8 直接喂进去会报类型错误而且报错信息指向的是 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_interpolationcv2.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 ReLUUNet 的基本单元 def __init__(self, in_ch, out_ch): super().__init__() self.block nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) 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, stride2) 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, sizeskip.shape[-2:], modebilinear, align_cornersFalse) x torch.cat([skip, x], dim1) return self.conv(x) class UNet(nn.Module): def __init__(self, in_ch3, num_classes2, base32): 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) 的 logitsUp里那段尺寸对齐的代码不是凑数的。UNet 要求输入边长是 16 的整数倍四次下采样输入 512 没问题输入 513 就会在下采样过程中丢掉一个像素上采样回来接不上。这份代码用interpolate兜住代价是那一路会有轻微的重采样损失。长期看在 Dataset 里统一Resize到 2 的幂次更干净。outc那行用 1×1 卷积而不是 3×3是因为这一步只做通道映射不做空间聚合。整个网络参数量在base32时大约是 7.8M单张 512×512 图在 8GB 显存上 batch 开到 8 是够的。3.2 输出通道、激活函数与损失的配对这一节是新手翻车最集中的地方三种组合必须记清楚任务形态输出通道模型末尾激活目标形状推荐损失二值前景/背景1不加(B, 1, H, W) floatBCEWithLogitsLoss二值 不平衡1不加(B, 1, H, W) floatBCE Dice多类别互斥num_classes不加(B, H, W) int64CrossEntropyLoss核心原则只有一条损失函数内部已经带了log_softmax/sigmoid模型输出层就不要再加一遍。模型里加了 softmax、外面再接CrossEntropyLoss等于做了两次 softmax梯度会被压得非常小训练速度肉眼可辨地变慢最后还训不出效果。类别不平衡在分割任务里是常态。一片 512×512 的遥感图里待检测的建筑可能只占 2% 的像素纯CrossEntropyLoss下模型十几轮就学会全预测背景loss 降到 0.05IoU 是 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, smooth1.0, ignore_index255): 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, dim1) # 概率化 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(min1) class ComboLoss(nn.Module): def __init__(self, num_classes, ce_weight1.0, dice_weight1.0): super().__init__() self.ce nn.CrossEntropyLoss(ignore_index255) 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-4weight decay 1e-4betas(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(enabledFalse)的上下文里用 fp32 算GradScaler的state_dict在断点续训时必须一起存不存的话恢复训练的第一个 step 大概率梯度溢出被跳过。from torch.cuda.amp import autocast, GradScaler model UNet(num_classes4).cuda() opt torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) sched torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max60) scaler GradScaler() criterion ComboLoss(num_classes4).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_noneTrue) with autocast(): logits model(img) # 损失在 fp32 下算避开 Dice 的除零和溢出 with autocast(enabledFalse): loss criterion(logits.float(), mask) scaler.scale(loss).backward() scaler.step(opt) scaler.update() sched.step()set_to_noneTrue比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, desctrain, leaveFalse): img img.to(device, non_blockingTrue) mask mask.to(device, non_blockingTrue) optimizer.zero_grad(set_to_noneTrue) with torch.cuda.amp.autocast(): logits model(img) with torch.cuda.amp.autocast(enabledFalse): loss criterion(logits.float(), mask) scaler.scale(loss).backward() # 梯度裁剪要在 scaler.step 之前、unscale 之后 scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.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(), lr1e-4, weight_decay1e-4) sched torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_maxepochs) scaler torch.cuda.amp.GradScaler() criterion ComboLoss(num_classesmodel.outc.out_channels).to(device) best 0.0 os.makedirs(ckpt_dir, exist_okTrue) 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(fep {ep:03d} loss {tr_loss:.4f} fmIoU {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_classes4, ignore_index255): model.eval() # 混淆矩阵行是真值列是预测 conf np.zeros((num_classes, num_classes), dtypenp.int64) for img, mask in loader: img img.to(device) logits model(img) pred logits.argmax(dim1).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, minlengthnum_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.benchmarkTrue会让卷积算法选择随机器状态变化追求完全可复现的话要关掉并设deterministicTrue代价是慢 10% 左右。import random import numpy as np import torch def set_seed(seed42, deterministicTrue): 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, optimizerNone, schedulerNone, scalerNone): ckpt torch.load(path, map_locationcpu) 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_lastTrue只在训练集上加验证集不能加否则会漏样本。或者把 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, tile512, stride384, num_classes4): image: (3, H, W) 的 float32 numpy已做归一化 model.eval() _, H, W image.shape # 累积概率和计数图重叠区域最后取平均 prob_sum np.zeros((num_classes, H, W), dtypenp.float32) count np.zeros((H, W), dtypenp.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), modebilinear, align_cornersFalse) prob torch.softmax(logits, dim1)[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(axis0)stride取tile的 3/4重叠 25%这个比例能覆盖大部分拼接缝再高收益递减、耗时翻倍。ys[-1] tile H那个补丁很关键不补的话右边和下边会缺一条缺的宽度正好是(H - tile) % stride很窄但正好落在你最关心的边界上。参数上还有两个地方值得调。tile不要超过训练时的输入尺寸模型没见过更大的感受野硬上会掉点num_classes必须跟训练时完全一致模型输出通道数对不上直接报维度错。交付前我会强制走一遍三项自检缺一项都不敢发出去检查项做法通过标准单图过拟合一张图训 200 步前景 IoU 0.95预处理一致性训练/验证/推理共用同一函数三个脚本输入张量统计量一致边缘样本挑 10 张前景最小的图单独看至少半数能看出轮廓第一项是排查管线断裂第二项是排查「训练好推理差」第三项是排查小目标召回。三项里第二项最容易出问题因为它是靠人手动保证的而人最容易在赶时间的时候复制粘贴。说到赶时间有一次我在一个工业缺陷项目上赶着交付验证集 mIoU 0.86 就直接打包了结果对方拿产线上的图跑缺陷几乎全漏。回头查了两天才发现推理脚本里那行A.Normalize被注释掉了——调试可视化的时候注释的忘了恢复。从那以后我每次交付前都强制走一遍单图过拟合测试和预处理一致性检查宁可多花二十分钟也不愿意再体验一次那种「指标漂亮、实际全废」的滋味。希望帮到你。本文还有配套的精品资源点击获取