做图像分割的人大多绕不开 UNet：它没有花哨的注意力模块，也没有复杂的多尺度金字塔堆叠，却凭一条“先下采样再上采样、同层特征横向拼接”的主干，在医学影像、工业缺陷检测、遥感地物提取这些场景里稳扎稳打。这篇内容围绕 PyTorch 实现 UNet 分割模型展开，把网络为什么这样设计、代码里每个卷积和拼接在做什么、训练时损失函数怎么配、推理后处理怎么写，以及实战中最容易踩的坑一次讲透。无论你是刚接触图像分割、想找一个能跑通的 PyTorch 分割 Demo，还是已经用过现成库但没细看网络结构，都可以顺着下面的思路自己搭一遍。我会把输入输出尺寸推演、通道数变化、损失函数组合、显存控制这些细节都摊开说，代码尽量给到能直接改成自己数据集的粒度，同时补充一些常规文档里不常写的实操心得。
1. UNet整体设计与思路拆解
1.1 从像素分类到端到端分割：U形结构到底解决什么问题
图像分割和图像分类的差别，不在于输出层多几个类别，而在于任务目标从“整张图是什么”变成了“每个像素属于什么”。这个变化会带来一个很直接的矛盾：网络越往深处走，感受野越大，越能判断“这块区域整体像不像病灶”，但特征图分辨率也越低，边界位置越模糊；网络越浅，分辨率越高，边界越清楚，可它看到的局部信息太少，容易把纹理相似的区域判错。UNet 的 U 形结构就是冲着这个矛盾去的。
它的收缩路径负责把图像一步步压小，让卷积核在更大的感受野里提取语义；扩张路径负责把特征图一步步放大，恢复空间分辨率；跳跃连接则把收缩路径里高分辨率的浅层特征直接送到解码器，让边界、纹理这些细节不至于在下采样过程中丢干净。你可以把它理解成做手工拼图：一边把大图拆成小块分类，一边在拼回去的时候保留原图的边线参考。没有跳跃连接，解码器只能靠深层语义“猜”边界，轮廓往往偏圆、偏毛糙；有了跳跃连接，浅层边缘信息能参与最终判断，小目标和细长结构会稳很多。
现代 PyTorch 实现里，原始论文那套“无 padding 卷积导致输入输出尺寸不一致”的做法已经很少用了，更常见的是 3x3 卷积加 padding=1、stride=1，保证卷积前后宽高不变，再用 2x2 最大池化把分辨率减半。这样一来输入 256x256，输出也是 256x256，调试和部署都省心。下面的推演和代码都按这种“same padding”思路来，避免你一上来就被裁剪参数绕晕。
1.2 编码器、瓶颈层、解码器与跳跃连接的职责划分
把 UNet 拆开看，它由四个部分组成：编码器、瓶颈层、解码器和跳跃连接。编码器通常由多个“双层卷积 + 下采样”组成，每一级通道数翻倍、分辨率减半，负责提取从边缘到语义的层次特征；瓶颈层是分辨率最低、通道数最多的一层，负责把全局上下文压进特征向量；解码器通过上采样逐步恢复分辨率，同时接收来自编码器同层的特征拼接；最后一层 1x1 卷积把通道数压到类别数，输出每个像素的 logits。
这里有个很容易被忽略的细节：编码器每一级输出的特征图，不只是给下一级下采样用，还要缓存起来给解码器拼接。所以写 PyTorch 代码时，常用一个列表features = []把每一级卷积后的特征存下来，解码时按从深到浅的顺序取出来拼接。拼接用torch.cat([upsampled, skip], dim=1)，维度是通道维，不是宽高维。拼接后通道数会变成两者之和，所以后续卷积块的输入通道要提前算好，否则一运行就报通道不匹配。
瓶颈层不需要再做跳跃连接，它只负责把最深层的语义传给第一个解码块。解码器的上采样方式常见两种：转置卷积和“插值 + 卷积”。转置卷积可以学习上采样参数，但容易在输出上产生棋盘格伪影；插值加卷积更稳，尤其是小数据集。实际项目里我会优先用nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)接一个 3x3 卷积，调试成本低，边界也更自然。如果追求极致速度，再考虑转置卷积。
1.3 输入输出尺寸与通道数的完整推演
假设输入是单通道灰度图，尺寸 256x256，也就是张量形状[B, 1, 256, 256]。第一级用 64 通道的双层卷积，卷积核 3x3、padding=1、stride=1，根据公式O = (W - K + 2P) / S + 1，输出宽高仍然是 256，所以形状变成[B, 64, 256, 256]。接着 2x2 最大池化，stride=2，宽高减半到 128，通道不变，形状[B, 64, 128, 128]。
第二级卷积把通道从 64 提到 128，宽高保持 128，再池化到 64；第三级 128 到 256，宽高 64 池化到 32；第四级 256 到 512，宽高 32 池化到 16；瓶颈层通常再走一次 512 到 1024 的卷积块，宽高保持 16。这样最深层特征图是[B, 1024, 16, 16]，感受野已经能覆盖很大区域。
解码第一块：对瓶颈特征做 2 倍上采样，宽高从 16 变 32，通道仍 1024；取编码器第四级缓存的 512 通道特征拼接，通道变成 1536；接一个双层卷积把通道压回 512，宽高保持 32。第二块上采样到 64，拼接编码器第三级 256 通道，通道变 768，卷积压到 256；第三块上采样到 128，拼接 128 通道，通道 384 压到 128；第四块上采样到 256，拼接 64 通道，通道 192 压到 64。最后 1x1 卷积把 64 通道压到类别数，二分类就是 1 通道，多分类就是类别数通道。
下面这张表可以把整个流程钉死，写代码时照着核对通道，比反复打印 shape 更快：
| 阶段 | 操作 | 输出分辨率 | 输出通道 | 是否缓存给解码器 |
|---|---|---|---|---|
| 编码1 | DoubleConv | 256x256 | 64 | 是 |
| 编码2 | MaxPool + DoubleConv | 128x128 | 128 | 是 |
| 编码3 | MaxPool + DoubleConv | 64x64 | 256 | 是 |
| 编码4 | MaxPool + DoubleConv | 32x32 | 512 | 是 |
| 瓶颈 | MaxPool + DoubleConv | 16x16 | 1024 | 否 |
| 解码1 | Upsample + Cat + DoubleConv | 32x32 | 512 | 否 |
| 解码2 | Upsample + Cat + DoubleConv | 64x64 | 256 | 否 |
| 解码3 | Upsample + Cat + DoubleConv | 128x128 | 128 | 否 |
| 解码4 | Upsample + Cat + DoubleConv | 256x256 | 64 | 否 |
| 输出 | 1x1 Conv | 256x256 | 类别数 | 否 |
这张表里最容易被写错的是解码第一块的输入通道：上采样后的 1024 加编码器第四级的 512，等于 1536，所以DoubleConv(1536, 512)才是对的。很多人直接写DoubleConv(1024, 512)，运行到拼接就会报错。另一个常错点是最后一次上采样后没有和编码器第一级的 64 通道拼接，导致输出分辨率虽然对了，但浅层细节没用上，边界效果会明显变差。
2. 核心细节解析与PyTorch实现要点
2.1 卷积块设计：为什么是“卷积 + BN + ReLU”重复两次
UNet 里最基础的单元是双层卷积块，通常写成Conv2d -> BatchNorm2d -> ReLU -> Conv2d -> BatchNorm2d -> ReLU。连续两次 3x3 卷积的感受野等于一次 5x5 卷积，但参数更少、非线性更多，特征表达更细。BatchNorm 的作用是把每个通道的激活分布拉回稳定区间，让训练对学习率没那么敏感；ReLU 负责提供非线性。小数据集上 BN 的 batch 统计量可能不稳，如果 batch size 只能开到 2 或 4，可以把 BN 换成 GroupNorm，分组数取 8 或 16，训练会更稳，代价是速度略慢。
写代码时要注意 padding 的选择。3x3 卷积配 padding=1 可以保持宽高不变，这是后面尺寸对齐的基础。如果用了 padding=0，每经过一次卷积宽高各减 2，跳跃拼接时尺寸对不上，还得手动裁剪，非常麻烦。池化统一用MaxPool2d(2)，不要混用 stride=2 的卷积下采样，否则通道和分辨率的对应关系容易乱。激活函数默认 ReLU，如果发现训练后期大量神经元死亡，可以换 LeakyReLU(0.01)，但大多数分割任务 ReLU 够用。
还有一个经验：不要在卷积块里加 Dropout 作为默认配置。分割任务对空间位置敏感，Dropout 随机丢特征容易让边界抖动。如果确实过拟合，优先用数据增强、权重衰减、早停，实在不行再在瓶颈层加少量 Dropout。下面是一个可以直接复用的卷积块：
import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x)这里把bias=False是因为后面接了 BN，BN 的偏移参数可以替代卷积偏置，省一点参数也避免重复。inplace=True的 ReLU 能省显存，但在某些需要保留原特征的调试场景里要注意，它会把输入原地改掉。正式训练没问题，调试阶段如果发现中间变量被改，可以临时关掉。
2.2 上采样与跳跃拼接：尺寸对齐和通道顺序是两大关卡
解码器的核心动作是“上采样 + 拼接 + 卷积”。上采样把低分辨率特征放大两倍，拼接把编码器同层特征接进来，卷积再融合。这里有两个关卡：尺寸必须完全一致，通道顺序必须清楚。用nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)时，输入 16x16 会变成 32x32，和编码器第四级的 32x32 对得上；但如果输入尺寸不是 16 的倍数，比如 250x250，经过四次池化可能变成 15x15 或 16x16，上采样后和跳跃特征差一个像素，torch.cat就会报错。
处理办法有两个：一是训练和推理前把图像 resize 到 16 的倍数，比如 256、320、512；二是在拼接前用F.interpolate把上采样结果强行对齐到跳跃特征的宽高。后者更灵活，但会引入额外插值，边界可能略软。我一般选前者，数据预处理时统一尺寸，省去很多运行时报错。拼接顺序也有讲究：torch.cat([up, skip], dim=1)和torch.cat([skip, up], dim=1)在数学上都能跑，但通道顺序会影响后续卷积的初始化学习。建议固定为“上采样特征在前、跳跃特征在后”，代码可读性更好，预训练权重迁移时也不容易错位。
下面是一个解码块的写法，把上采样、拼接、卷积封在一起，调用时只传通道数：
class UpBlock(nn.Module): def __init__(self, in_ch, skip_ch, out_ch): super().__init__() self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) self.conv = DoubleConv(in_ch + skip_ch, out_ch) def forward(self, x, skip): x = self.up(x) if x.shape[-2:] != skip.shape[-2:]: x = nn.functional.interpolate( x, size=skip.shape[-2:], mode='bilinear', align_corners=False ) x = torch.cat([x, skip], dim=1) return self.conv(x)这段代码里加了一个尺寸兜底，虽然平时用不上，但在数据尺寸不规整时能救急。注意skip.shape[-2:]取的是高和宽，不是通道。dim=1是通道维，别写成dim=-1，否则拼的是宽高，后面卷积直接崩。
2.3 损失函数与输出激活：为什么模型最后一层不要加 Sigmoid
二值分割最常见的损失是BCEWithLogitsLoss，它内部已经把 Sigmoid 和交叉熵合在一起，数值更稳定。所以模型最后一层 1x1 卷积直接输出 logits，不要加nn.Sigmoid()。如果你在模型里加了 Sigmoid，训练时又用BCEWithLogitsLoss，等于做了两次激活，损失会异常，预测不是全黑就是全白。这个坑非常常见，尤其是从分类模型改分割模型时，容易顺手把 Sigmoid 留在输出端。
类别不平衡时，单用 BCE 会让模型偏向背景。医学影像里病灶可能只占几个像素，背景占 99%，模型全预测背景也能拿到很低 loss。解决办法是 BCE 和 Dice Loss 加权组合。Dice 系数关注预测区域和真实区域的重叠，对前景少样本更敏感。实现时给分子分母加平滑项，避免预测全空时分母为零。常见权重是loss = 0.5 * bce + 0.5 * dice，也可以根据验证集调到 0.3/0.7。多分类则用CrossEntropyLoss加多类 Dice，输出通道等于类别数，不要做 softmax 后再传给 CrossEntropyLoss。
下面是一个二值 Dice Loss 的实现，输入是 logits，内部先 Sigmoid：
class DiceLoss(nn.Module): def __init__(self, smooth=1e-6): super().__init__() self.smooth = smooth def forward(self, logits, targets): probs = torch.sigmoid(logits) probs = probs.contiguous().view(probs.size(0), -1) targets = targets.contiguous().view(targets.size(0), -1) intersection = (probs * targets).sum(dim=1) union = probs.sum(dim=1) + targets.sum(dim=1) dice = (2 * intersection + self.smooth) / (union + self.smooth) return 1 - dice.mean()这里view之前加contiguous()是为了防止非连续张量报错。平滑项不要太大，1e-6 到 1e-5 足够，太大反而让 loss 下降变慢。如果标签里有忽略区域，比如标注边界不确定的像素，要在 loss 里用 mask 排除，不要直接当背景训练，否则模型会学到错误的负样本。
2.4 数据管线与标签处理：同步增强比模型本身更容易出问题
分割任务的数据增强和分类不一样，图像旋转、翻转、缩放时，掩码必须做完全相同的变换。用 torchvision 的transforms单独处理图像和标签很容易错位，推荐用 albumentations 这类支持 image 和 mask 同步变换的库，或者在 Dataset 里用同一个随机参数同时处理两者。常见的增强包括随机水平翻转、垂直翻转、小角度旋转、随机缩放裁剪、亮度对比度扰动。弹性变形对医学影像很有效，但参数不要太大，否则边界被拉得不像真实结构。
标签处理有几个细节：一是确保掩码是单通道或类别索引，不要变成三通道 RGB 图；二是二值化阈值要统一，比如标注里 255 表示前景，读进来后要除以 255 或判断大于 127；三是忽略区域用 255 或特定值标记时，要在 loss 和指标里排除；四是训练集和验证集的归一化参数要一致，通常用训练集均值方差，不要各自算各自的。下面是一个 Dataset 的核心片段，展示图像和掩码如何同步读取：
from torch.utils.data import Dataset from PIL import Image import numpy as np class SegDataset(Dataset): def __init__(self, image_paths, mask_paths, transform=None): self.image_paths = image_paths self.mask_paths = mask_paths self.transform = transform def __len__(self): return len(self.image_paths) def __getitem__(self, idx): image = np.array(Image.open(self.image_paths[idx]).convert('L')) mask = np.array(Image.open(self.mask_paths[idx]).convert('L')) mask = (mask > 127).astype(np.float32) if self.transform: augmented = self.transform(image=image, mask=mask) image, mask = augmented['image'], augmented['mask'] image = image.astype(np.float32) / 255.0 image = np.expand_dims(image, axis=0) mask = np.expand_dims(mask, axis=0) return torch.from_numpy(image), torch.from_numpy(mask)如果图像是三通道，就把convert('L')去掉，并在最后按通道堆叠。掩码不要做归一化到 0 到 1 以外的操作，二值任务保持 0/1 浮点即可。DataLoader 的num_workers在 Windows 上容易出问题，建议先设 0 调试通再调大；pin_memory=True在 GPU 训练时能加快数据搬运。
3. 从零搭建UNet的完整实操流程
3.1 环境准备与目录约定：先把工程骨架搭稳
开始写模型前，先把环境和目录定好。PyTorch 版本不需要追最新，1.10 以上、CUDA 和显卡驱动匹配即可。建议用虚拟环境隔离依赖，核心包包括 torch、torchvision、numpy、Pillow、albumentations、tqdm、scikit-image。如果要做可视化，可以加 matplotlib，但训练环境里不必装。下面是一组常见安装命令，实际版本按你的显卡环境调整：
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install numpy pillow albumentations tqdm scikit-image目录建议按“数据、代码、权重、日志、预测结果”分开，不要把所有文件堆在一个文件夹。一个模拟项目的结构可以是这样：
unet_demo/ ├── data/ │ ├── train/images/ │ ├── train/masks/ │ ├── val/images/ │ └── val/masks/ ├── weights/ ├── logs/ ├── preds/ ├── model.py ├── dataset.py ├── losses.py ├── train.py └── predict.py这样拆分的好处是调试时只改一个文件，不会牵一发动全身。model.py只放网络结构，dataset.py只放数据读取和增强，losses.py放损失函数，train.py负责训练循环，predict.py负责推理。很多人把全部代码写在一个脚本里，前期跑得快，后期改一个损失函数要翻半天，反而更慢。
3.2 完整模型代码：把编码器、瓶颈、解码器串起来
下面给出一个可直接运行的 UNet 实现，输入单通道、输出单通道 logits，基础通道数 64。代码里保留了尺寸兜底，方便处理非 2 的幂尺寸。你可以把in_channels改成 3 处理 RGB 图，把out_channels改成类别数做多分类。
import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): 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 UNet(nn.Module): def __init__(self, in_channels=1, out_channels=1, base_ch=64): super().__init__() self.enc1 = DoubleConv(in_channels, base_ch) self.enc2 = DoubleConv(base_ch, base_ch * 2) self.enc3 = DoubleConv(base_ch * 2, base_ch * 4) self.enc4 = DoubleConv(base_ch * 4, base_ch * 8) self.pool = nn.MaxPool2d(2) self.bottleneck = DoubleConv(base_ch * 8, base_ch * 16) self.up4 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) self.dec4 = DoubleConv(base_ch * 16 + base_ch * 8, base_ch * 8) self.up3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) self.dec3 = DoubleConv(base_ch * 8 + base_ch * 4, base_ch * 4) self.up2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) self.dec2 = DoubleConv(base_ch * 4 + base_ch * 2, base_ch * 2) self.up1 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) self.dec1 = DoubleConv(base_ch * 2 + base_ch, base_ch) self.out_conv = nn.Conv2d(base_ch, out_channels, kernel_size=1) def forward(self, x): e1 = self.enc1(x) e2 = self.enc2(self.pool(e1)) e3 = self.enc3(self.pool(e2)) e4 = self.enc4(self.pool(e3)) b = self.bottleneck(self.pool(e4)) d4 = self.up4(b) if d4.shape[-2:] != e4.shape[-2:]: d4 = F.interpolate(d4, size=e4.shape[-2:], mode='bilinear', align_corners=False) d4 = self.dec4(torch.cat([d4, e4], dim=1)) d3 = self.up3(d4) if d3.shape[-2:] != e3.shape[-2:]: d3 = F.interpolate(d3, size=e3.shape[-2:], mode='bilinear', align_corners=False) d3 = self.dec3(torch.cat([d3, e3], dim=1)) d2 = self.up2(d3) if d2.shape[-2:] != e2.shape[-2:]: d2 = F.interpolate(d2, size=e2.shape[-2:], mode='bilinear', align_corners=False) d2 = self.dec2(torch.cat([d2, e2], dim=1)) d1 = self.up1(d2) if d1.shape[-2:] != e1.shape[-2:]: d1 = F.interpolate(d1, size=e1.shape[-2:], mode='bilinear', align_corners=False) d1 = self.dec1(torch.cat([d1, e1], dim=1)) return self.out_conv(d1)这段代码里没有在输出端加 Sigmoid，原因前面说过，损失函数用BCEWithLogitsLoss。如果你要换成多分类，out_channels改成类别数，损失换CrossEntropyLoss，推理时再对输出做argmax。模型参数量可以通过sum(p.numel() for p in model.parameters())打印，基础通道 64 的版本大约三千万参数，如果显存吃紧，把base_ch降到 32，参数会少很多，分割精度在简单任务上下降有限。
3.3 训练循环与验证指标：别只看 loss，要看 Dice 和 IoU
训练循环的骨架很固定：前向传播、算 loss、反向传播、优化器更新、清零梯度。但分割任务有几个额外细节。优化器常用 Adam 或 AdamW，初始学习率 1e-4 到 3e-4，权重衰减 1e-4。学习率调度器可以用ReduceLROnPlateau，监控验证集 Dice，连续几个 epoch 不涨就减半。混合精度训练用torch.cuda.amp，能明显省显存，但要注意 Dice Loss 里的求和操作在 half 精度下可能溢出，建议 loss 计算放在 float32 里。
验证指标至少看 Dice 和 IoU。Dice 对前景敏感，IoU 更严格。二值任务先把 logits 过 Sigmoid，再用阈值 0.5 转成 0/1，然后计算交集和并集。验证时记得model.eval()和torch.no_grad()，否则 BN 会更新统计量，显存也会爆。下面是一个训练循环的核心片段：
import torch from torch.utils.data import DataLoader def train_one_epoch(model, loader, optimizer, bce_loss, dice_loss, device, scaler): model.train() total_loss = 0.0 for images, masks in loader: images = images.to(device, non_blocking=True) masks = masks.to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) with torch.cuda.amp.autocast(): logits = model(images) loss = 0.5 * bce_loss(logits, masks) + 0.5 * dice_loss(logits, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() total_loss += loss.item() * images.size(0) return total_loss / len(loader.dataset) @torch.no_grad() def validate(model, loader, device, threshold=0.5): model.eval() dice_sum, iou_sum, count = 0.0, 0.0, 0 for images, masks in loader: images = images.to(device) masks = masks.to(device) logits = model(images) probs = torch.sigmoid(logits) preds = (probs > threshold).float() intersection = (preds * masks).sum(dim=(1, 2, 3)) union = preds.sum(dim=(1, 2, 3)) + masks.sum(dim=(1, 2, 3)) dice = (2 * intersection + 1e-6) / (union + 1e-6) iou = (intersection + 1e-6) / (union - intersection + 1e-6) dice_sum += dice.sum().item() iou_sum += iou.sum().item() count += images.size(0) return dice_sum / count, iou_sum / count这里set_to_none=True比zero_grad()更省显存。混合精度下scaler的用法要固定，不要漏掉scaler.update()。验证阶段如果显存不够，可以减小 batch size，或者把图像分块推理。指标计算里加 1e-6 是防止除零，别省。
3.4 推理与后处理：整图推理、滑动窗口和阈值选择
训练完保存权重，推理时加载模型，model.eval()后输入图像，得到 logits，过 Sigmoid，再按阈值二值化。小图像可以直接整图推理，大图像比如 1024x1024 或更高，整图放进显存可能爆，这时候用滑动窗口：把大图裁成 256x256 的小块，重叠 32 或 64 像素，逐块推理后拼回去，重叠区域取平均或取最大值。重叠能减少拼接缝，但速度会慢一些。
阈值不要死守 0.5。医学影像里如果漏检代价高，阈值可以降到 0.3 到 0.4，提高召回；如果误检代价高，阈值提到 0.6。阈值最好在验证集上扫一遍，画一条 Dice-阈值曲线，选拐点。后处理可以加连通域分析，去掉面积过小的孤立区域，或者用形态学闭运算填小孔。下面是一个推理片段：
@torch.no_grad() def predict_image(model, image_tensor, device, threshold=0.5): model.eval() image_tensor = image_tensor.to(device) logits = model(image_tensor) probs = torch.sigmoid(logits) mask = (probs > threshold).float() return mask.cpu()如果要做连通域过滤，可以把 mask 转到 numpy，用skimage.measure.label标记，统计每个区域面积，小于min_area的置零。min_area不要拍脑袋定，根据验证集里最小目标面积来设，通常是几十到几百像素。后处理不是越多越好，过度形态学操作会把小目标抹掉，调参时一定要看可视化结果。
4. 常见问题与排查技巧实录
4.1 维度不匹配：先打印 shape，再核对拼接通道
维度不匹配是 UNet 实现里最高频的报错，常见表现是torch.cat时提示宽高不一致，或者卷积输入通道和权重通道对不上。排查顺序很简单：在forward里每个拼接前打印d4.shape和e4.shape，看宽高是否一致、通道是否等于预期。宽高不一致通常是输入尺寸不是 16 的倍数，或者某次池化多算了一层；通道不一致通常是DoubleConv的输入通道写错，尤其是解码第一块忘了加上采样后的通道。
我习惯在模型里加一个debug=False参数，打开时打印每一级 shape，训练前跑一个随机张量torch.randn(2, 1, 256, 256)过一遍。这样不用等 DataLoader，几秒钟就能定位问题。另一个技巧是把拼接写成torch.cat([x, skip], dim=1)后立刻接一个print(x.shape)，如果报错发生在卷积里，就能确认是拼接后的通道数没对上。不要靠猜，shape 打印是最快的。
4.2 显存溢出：先降 batch，再考虑 AMP 和梯度检查点
显存溢出不一定是模型太大，很多时候是输入尺寸、batch size、数据加载器或验证阶段没关梯度导致的。排查顺序：先把 batch size 降到 1，看能不能跑；如果能跑，说明是 batch 太大，逐步往上加，找到显卡的极限。然后检查验证阶段有没有torch.no_grad()，有没有误把计算图保留下来。数据加载器的num_workers太高也会占内存，但一般不影响显存。
如果 batch 已经降到 1 还爆，考虑混合精度、梯度检查点、减小基础通道数、把输入从 512 降到 256。混合精度能省 30% 到 50% 显存，但要注意 loss 里的求和操作。梯度检查点用时间换显存，适合深层网络，UNet 本身不算特别深，收益有限。另一个容易忽略的点是验证时不要保留所有 batch 的输出，算完指标就释放，否则验证集一大也会爆。
4.3 损失不下降与预测全黑全白：从标签、激活和损失三处查
loss 不下降或者预测全黑全白，先查三件事：标签对不对、模型输出有没有多余激活、损失函数和输出是否匹配。标签常见问题是掩码被归一化到 0 到 1 之外，或者前景背景反了；模型输出常见问题是最后加了 Sigmoid，而损失又用BCEWithLogitsLoss；损失常见问题是类别极不平衡时只用 BCE，模型全预测背景。把这三项排除后，再看学习率是不是太大，优化器是不是没清零梯度。
如果预测全黑，先看 logits 的分布，是不是全部小于 0；如果全白，看是不是全部大于 0。可以在训练初期打印logits.min()和logits.max()，正常应该有正有负。还可以单独拿一张训练集图像过模型，看能不能过拟合到这张图，如果连训练集都学不会，说明网络或损失有结构性问题。学习率从 1e-4 开始试，太大容易震荡，太小收敛慢。
4.4 小目标漏检与边界毛刺：损失加权和后处理要配合
小目标漏检通常不是网络容量不够，而是损失被背景主导。可以给前景像素更高的 BCE 权重，或者在 Dice Loss 里对每张图单独算 Dice 再平均，而不是把所有像素混在一起算。深监督也能帮小目标：在解码器中间层加辅助输出，让浅层也直接受损失约束。但深监督会增加显存和代码复杂度，简单任务不必上。
边界毛刺常见原因是上采样方式太硬，或者训练时边界像素权重不够。把转置卷积换成双线性插值加卷积，通常能缓解棋盘格；在损失里给边界附近像素加权，能让模型更关注轮廓。后处理用形态学开闭运算可以去掉毛刺，但核不能太大，3x3 或 5x5 足够。如果毛刺是标注本身带来的，先检查标注质量，不要指望模型学会不一致的边界。
下面这张速查表可以把常见现象和排查方向对上：
| 现象 | 优先排查 | 常见原因 | 处理办法 |
|---|---|---|---|
| 拼接报错 | 宽高和通道 | 输入尺寸非 16 倍数、通道写错 | resize 到 16 倍数、打印 shape 核对 |
| 显存溢出 | batch 和输入尺寸 | 验证没关梯度、batch 太大 | no_grad、降 batch、AMP |
| loss 不降 | 标签和激活 | 标签反了、输出多加了 Sigmoid | 可视化标签、去掉输出激活 |
| 预测全黑 | logits 分布 | 阈值太高、损失不平衡 | 降阈值、BCE+Dice |
| 小目标漏检 | 损失权重 | 背景主导、下采样丢细节 | 前景加权、深监督、裁块训练 |
| 边界毛刺 | 上采样和后处理 | 转置卷积棋盘格、标注噪声 | 双线性上采样、形态学小核 |
5. 应用场景与工程落地影响
5.1 哪些任务适合UNet，哪些任务要谨慎
UNet 最适合“输入输出同尺寸、目标边界重要、数据量中等”的密集预测任务。医学影像里的器官分割、病灶分割，工业质检里的划痕、孔洞、异物检测，遥感里的建筑、道路、水体提取，甚至内容创作里的抠图、蒙版生成，都能用 UNet 做基线。它的优势是结构简单、训练稳定、对小数据集相对友好，尤其是跳跃连接能保留细节，边界表现比纯编码器-解码器结构好。
但 UNet 不是万能的。如果目标尺度差异极大，比如一张图里既有几十像素的大目标又有两三个像素的小目标，固定感受野的 UNet 会吃力，需要考虑多尺度特征金字塔或注意力机制。如果数据量非常大、类别非常多，UNet 的参数量和计算量可能不如更高效的骨干网络。如果任务需要强全局建模，比如长条形道路的连续性，纯卷积的感受野有限，可以结合 Transformer 或空洞卷积。选型时先跑 UNet 基线，再根据失败案例分析是加数据、改损失还是换结构，不要一上来就堆模块。
5.2 从Demo到落地：数据一致性比模型结构更影响上限
真正落地时，模型结构往往不是瓶颈，数据标注的一致性才是。同一个病灶，不同标注者的边界可能差几个像素，模型学出来的边界就会抖。工程上要做的是统一标注规范、定期抽检、对模糊区域设忽略标签。推理端要考虑图像预处理和训练时是否一致，归一化参数、通道顺序、resize 方式都要对齐，否则线上效果会明显掉。
部署时还要关注推理速度和显存占用。UNet 基础版在单张 GPU 上跑 256x256 图像通常很快，但如果要实时处理高分辨率视频流，就要考虑模型剪枝、量化、TensorRT 或 ONNX 导出。导出时注意上采样算子和interpolate的兼容性，动态宽高可能不被某些推理引擎支持，最好固定输入尺寸。监控方面，除了 Dice 和 IoU，还要看预测掩码的面积分布、连通域数量，这些指标能提前发现数据漂移。
我在一个模拟的工业缺陷分割项目里踩过最深的坑，不是网络结构，而是验证集和训练集的预处理不一致：训练时图像除以 255，验证时忘了除，模型在验证集上 Dice 只有 0.1，查了两天才发现。后来我把预处理写进 Dataset 的__getitem__，训练和推理共用同一段代码，问题再没出现过。另一个心得是，不要等训练完才看预测图，每个 epoch 抽几张验证图保存下来，早期就能发现标签反了、阈值不对、边界偏移这些问题，比盯着 loss 曲线有效得多。
