从零手写UNet图像分割：结构拆解、训练调参与工程实践

1. 为什么值得亲手复现一遍 UNet

UNet 这个网络结构在图像分割领域算是绕不开的经典。2015 年那篇论文出来之后，医学影像、遥感解译、工业质检、甚至卫星图像分析，到处都能看到它的影子。但有意思的是，很多人第一次接触 UNet 都是在各种教程里直接调库，用现成的segmentation_models_pytorch一行代码加载预训练权重，跑通一个 Demo 就结束了。真正自己从零搭一遍的人反而不多。

我当初决定手写复现，原因很直接：调库跑出来的结果不对劲时，我根本不知道问题出在哪。是编码器下采样太激进？是跳跃连接的特征对齐有偏差？还是上采样的方式选错了？这些疑问在封装好的 API 面前完全无从下手。所以后来我花了几个晚上，把 UNet 从结构到训练流程完整写了一遍，踩了不少坑，也理清了很多之前模糊的地方。

这篇文章适合两类人：一类是刚入门图像分割，想搞明白 UNet 到底长什么样的新手；另一类是用过现成库，但想深入理解每个模块设计意图的进阶者。我会从结构拆解开始，讲到数据管道、损失函数选择、训练调参，再到实际跑起来之后遇到的各种意外情况。代码基于 PyTorch，尽量写得干净可读，你可以直接拿去改。

需要提前说明的是，我复现的是 UNet 的核心架构思想，不是逐行还原论文里的每一个超参数。原始论文针对的是 ISBI 细胞分割挑战赛的数据集，输入是 572×572 的灰度图，输出是 388×388 的分割掩码。这个尺寸设计在今天看来有些别扭，实际项目中我们通常会根据自己的数据特点做调整。所以我会在保持核心结构不变的前提下，采用更通用的实现方式。

2. UNet 的骨架拆解：每个模块到底在干什么

2.1 编码器-解码器结构的设计逻辑

UNet 最核心的思想就是"先压缩再恢复"。编码器部分通过一系列卷积和池化操作，把输入图像的 spatial 维度逐步缩小，同时通道数逐步增加。这个过程本质上是在提取从低级到高级的语义特征——浅层卷积看到的是边缘、纹理，深层卷积看到的是更抽象的形状和结构。

解码器则反过来，通过上采样操作把特征图逐步恢复到原始分辨率。但这里有个关键问题：如果只靠编码器压缩再解码器恢复，很多空间细节信息会在池化过程中丢失。尤其是医学图像分割这种任务，边界往往只有几个像素宽，一旦丢失就很难恢复准确。

所以 UNet 引入了跳跃连接（skip connection），把编码器每一层的特征图直接拼接到解码器对应层。这样解码器在做上采样时，既有来自深层的语义信息，又有来自浅层的空间细节信息。这个设计后来被无数网络借鉴，包括 Feature Pyramid Network 和各种 Transformer 分割模型。

具体到实现层面，编码器的每一层通常是两个 3×3 卷积加 ReLU 激活，然后接一个 2×2 最大池化。解码器每一层是先做上采样（可以是转置卷积或双线性插值），然后把编码器对应层的特征图拼接过来，再接两个 3×3 卷积加 ReLU。最后一层用 1×1 卷积把通道数映射到类别数。

2.2 跳跃连接的特征对齐问题

这里有一个很容易被忽略的细节：编码器和解码器的特征图尺寸必须严格对齐才能拼接。但卷积操作如果不做 padding，每次 3×3 卷积都会让特征图缩小 2 个像素。原始论文用的是 valid padding，所以编码器每层输出尺寸比输入小 4 个像素，解码器上采样后的尺寸又和编码器对应层差了一点，最后需要做中心裁剪才能拼接。

这种设计在实际写代码时非常容易出错。我的建议是直接用 padding=1 的 same 卷积，这样卷积操作不改变特征图尺寸，尺寸变化只发生在池化和上采样阶段，对齐逻辑会清晰很多。代价是边界像素的卷积结果会受到 padding 的影响，但对于大多数分割任务来说，这个影响可以忽略。

另一个细节是拼接和相加的区别。UNet 用的是拼接（concatenation），把通道维度合并。有些网络用相加（element-wise addition），比如 ResNet 的残差连接。拼接保留了更多信息，但参数量更大；相加更轻量，但可能丢失一些细节。UNet 选择拼接，是因为分割任务对细节要求高，宁可多花一点计算量也要保住信息量。

2.3 上采样的三种实现方式对比

上采样是解码器的关键操作，PyTorch 里常见的有三种做法：

我实测下来，转置卷积在训练初期容易不稳定，尤其是学习率稍大的时候。双线性插值虽然简单，但配合后续的卷积层，效果往往不差，而且训练更稳。最近邻插值一般只在推理速度要求极高的场景下用。

如果你追求最佳效果，可以试试转置卷积加适当的权重初始化，但要做好调参的心理准备。如果只是想快速跑通一个 baseline，双线性插值是最省心的选择。

3. 从零搭建：代码逐段拆解与设计决策

3.1 卷积块的封装思路

先定义一个基础的双卷积块，这是 UNet 里重复出现最多的单元：

import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels, mid_channels=None): super().__init__() if not mid_channels: mid_channels = out_channels self.double_conv = nn.Sequential( nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(mid_channels), nn.ReLU(inplace=True), nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.double_conv(x)

这里有几个设计决策值得说明。第一，用了bias=False配合 BatchNorm，因为 BN 本身有偏置项，卷积再加 bias 是冗余的。第二，padding=1保证卷积不改变尺寸，简化后续对齐逻辑。第三，inplace=True的 ReLU 节省一点显存，但在某些需要保留输入值的场景下要慎用。

原始论文没有用 BatchNorm，因为当时 BN 刚出来不久。但现代实现基本都会加上，它能显著加速收敛并提升稳定性。如果你要复现论文的原始效果，可以把 BN 去掉，但训练会慢很多。

3.2 下采样与上采样模块的实现

下采样就是最大池化，没什么好说的：

class Down(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_channels, out_channels) ) def forward(self, x): return self.maxpool_conv(x)

上采样模块稍微复杂一点，因为要处理拼接后的通道数变化：

class Up(nn.Module): def __init__(self, in_channels, out_channels, bilinear=True): super().__init__() if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv = DoubleConv(in_channels, out_channels, in_channels // 2) else: self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) self.conv = DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 = self.up(x1) # 处理尺寸不匹配的情况 diffY = x2.size()[2] - x1.size()[2] diffX = x2.size()[3] - x1.size()[3] x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x = torch.cat([x2, x1], dim=1) return self.conv(x)

那个F.pad的处理很关键。虽然理论上尺寸应该对齐，但实际中如果输入图像的尺寸不是 2 的整数次幂，经过多次下采样后可能会出现奇数尺寸，上采样回来就差了一个像素。这个 padding 逻辑能自动修正，避免运行时崩溃。

3.3 完整网络组装与通道数配置

把上面的模块串起来：

class UNet(nn.Module): def __init__(self, n_channels, n_classes, bilinear=True): super().__init__() self.n_channels = n_channels self.n_classes = n_classes self.bilinear = bilinear self.inc = DoubleConv(n_channels, 64) self.down1 = Down(64, 128) self.down2 = Down(128, 256) self.down3 = Down(256, 512) factor = 2 if bilinear else 1 self.down4 = Down(512, 1024 // factor) self.up1 = Up(1024, 512 // factor, bilinear) self.up2 = Up(512, 256 // factor, bilinear) self.up3 = Up(256, 128 // factor, bilinear) self.up4 = Up(128, 64, bilinear) self.outc = nn.Conv2d(64, n_classes, kernel_size=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)

通道数的配置是 64-128-256-512-1024 的经典设置。这个配置在显存和效果之间取得了不错的平衡。如果你的显存紧张，可以把所有通道数减半；如果数据量大且追求精度，可以适当增加，但要注意过拟合风险。

factor那个变量是为了处理双线性插值和转置卷积在通道数上的差异。双线性插值不改变通道数，所以上采样后通道数还是 1024，拼接后是 1024+512=1536，需要卷积降到 512。转置卷积会把通道数减半，所以逻辑略有不同。这个细节在写代码时很容易搞混，建议画个表格把每层的输入输出通道数都列清楚。

4. 数据管道与损失函数：比网络结构更容易翻车的地方

4.1 数据集类的编写要点

分割任务的数据管道比分类任务复杂，因为输入是图像，输出也是图像（掩码）。我见过太多人在这里踩坑，最常见的问题是图像和掩码的预处理不一致。

from torch.utils.data import Dataset from PIL import Image import numpy as np class SegmentationDataset(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('RGB')) mask = np.array(Image.open(self.mask_paths[idx]).convert('L')) if self.transform: augmented = self.transform(image=image, mask=mask) image = augmented['image'] mask = augmented['mask'] image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 mask = torch.from_numpy(mask).long() return image, mask

关键点：图像要做归一化，掩码要保持整数类别标签。如果你用 albumentations 做数据增强，一定要确保图像和掩码同步变换，否则掩码和图像就错位了。我建议用 albumentations 的Compose同时传入 image 和 mask，它会自动处理同步问题。

另一个坑是掩码的像素值。有些数据集的掩码是 0 和 255，有些是 0 和 1，还有些是多类别的 0、1、2、3。你需要先确认清楚，然后决定是否要做映射。如果是二分类，把 255 映射到 1；如果是多分类，确保标签从 0 开始连续。

4.2 损失函数的选择与组合策略

二分类分割最常用的是 BCEWithLogitsLoss，多分类用 CrossEntropyLoss。但实际项目中，单纯用这两个往往不够，因为分割任务普遍存在类别不平衡问题——背景像素远多于前景像素。

我常用的组合是 Dice Loss 加 CrossEntropyLoss：

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.view(-1) targets = targets.view(-1) intersection = (probs * targets).sum() dice = (2. * intersection + self.smooth) / (probs.sum() + targets.sum() + self.smooth) return 1 - dice class CombinedLoss(nn.Module): def __init__(self, weight_dice=0.5, weight_bce=0.5): super().__init__() self.dice = DiceLoss() self.bce = nn.BCEWithLogitsLoss() self.weight_dice = weight_dice self.weight_bce = weight_bce def forward(self, logits, targets): return self.weight_dice * self.dice(logits, targets) + \ self.weight_bce * self.bce(logits, targets.float())

Dice Loss 直接优化分割评价指标，对类别不平衡不敏感；BCE 提供稳定的梯度信号。两者结合通常比单用其中一个效果好。权重可以根据任务调整，如果前景特别少，可以加大 Dice 的权重。

多分类的情况可以用多类 Dice Loss，或者用带类别权重的 CrossEntropyLoss。我一般会先统计训练集里每个类别的像素占比，然后按反比设置权重。

4.3 评价指标的实现细节

IoU 和 Dice 是最常用的分割指标。实现本身不难，但要注意边界情况的处理：

def calculate_iou(pred, target, num_classes, ignore_background=True): ious = [] pred = pred.view(-1) target = target.view(-1) start = 1 if ignore_background else 0 for cls in range(start, num_classes): pred_inds = pred == cls target_inds = target == cls intersection = (pred_inds & target_inds).sum().item() union = (pred_inds | target_inds).sum().item() if union == 0: ious.append(float('nan')) else: ious.append(intersection / union) return np.nanmean(ious)

当某个类别在预测和目标中都不存在时，union 为 0，这时 IoU 是未定义的。我选择返回 nan 然后用 nanmean 忽略掉，而不是直接算 0 或 1。这个细节在类别较多的数据集上影响很大，处理不好会导致指标虚高或虚低。

5. 训练过程中那些文档不会告诉你的坑

5.1 学习率与优化器的实际选择

UNet 的训练对学习率比较敏感。我试过 Adam 和 SGD 两种优化器，Adam 收敛更快但最终精度有时不如调好的 SGD。如果时间紧，用 Adam 配 1e-3 或 1e-4 的学习率通常能跑出一个不错的结果。如果追求极致精度，SGD 配 momentum=0.9 和 1e-2 的学习率，再配合余弦退火调度，往往能更上一层楼。

学习率调度方面，ReduceLROnPlateau 是最省心的选择，监控验证集的 loss，连续几个 epoch 不下降就减半。但要注意 patience 的设置，太小会导致学习率过早衰减，太大又起不到作用。我一般设 patience=5，factor=0.5。

还有一个容易被忽略的点：BatchNorm 和 batch size 的关系。如果 batch size 太小（比如 2 或 4），BN 的统计量会很不稳定，训练 loss 会剧烈震荡。这种情况下可以考虑用 GroupNorm 替代 BN，或者用梯度累积来模拟更大的 batch size。

5.2 过拟合的识别与应对

分割任务的数据标注成本高，数据集通常不大，过拟合是常态。判断过拟合的方法很简单：训练 loss 持续下降，但验证 loss 在某个 epoch 后开始上升，同时验证集的 IoU 停滞或下降。

应对手段有几个。最直接的是数据增强，随机旋转、翻转、缩放、弹性变形对分割任务都很有效。其次是加 Dropout，可以在编码器的每个 DoubleConv 后面加nn.Dropout2d(p=0.1)。再就是早停，在验证指标最好的时候保存模型，而不是等训练结束。

我个人的经验是，数据增强的收益远大于正则化手段。如果数据增强做到位了，Dropout 甚至可以不加。弹性变形对医学图像特别有用，因为器官和组织的形状本身就有自然变化。

5.3 显存不足时的取舍策略

UNet 在 512×512 输入下，batch size 设为 8 大概需要 8-10GB 显存。如果显存不够，有几个选择：

• 减小 batch size 到 2 或 4，配合梯度累积
• 把输入裁剪成 256×256 的 patch，但要注意边界效应
• 减少通道数，比如把 64-128-256-512-1024 改成 32-64-128-256-512
• 用混合精度训练，torch.cuda.amp能省将近一半显存

混合精度是我最推荐的方式，改动小，收益大。只需要在训练循环里加几行：

scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

注意用混合精度时，loss 计算最好在 autocast 上下文之外做，或者确保 loss 函数支持 fp16。Dice Loss 里的除法操作在 fp16 下可能不稳定，建议把 probs 转成 fp32 再算。

6. 推理与后处理：让结果真正可用

6.1 滑动窗口推理大图

训练时可能用的是 256×256 的 patch，但推理时面对的是 1024×1024 甚至更大的图像。直接缩放会丢失细节，分块推理又会有边界拼接痕迹。滑动窗口加重叠是一种折中方案：

def sliding_window_inference(model, image, window_size=256, stride=128, num_classes=2): model.eval() _, h, w = image.shape output = torch.zeros(num_classes, h, w).to(image.device) count = torch.zeros(1, h, w).to(image.device) for y in range(0, h, stride): for x in range(0, w, stride): y1 = min(y, h - window_size) x1 = min(x, w - window_size) patch = image[:, y1:y1+window_size, x1:x1+window_size].unsqueeze(0) with torch.no_grad(): pred = model(patch) output[:, y1:y1+window_size, x1:x1+window_size] += pred.squeeze(0) count[:, y1:y1+window_size, x1:x1+window_size] += 1 output /= count return output

stride 设为 window_size 的一半，这样每个像素至少被覆盖两次，边界处的预测会更平滑。重叠区域的预测取平均，能有效消除拼接痕迹。

6.2 后处理中的连通域过滤

模型输出的二值掩码往往会有一些孤立的噪点，或者小的假阳性区域。用连通域分析过滤掉面积过小的区域，通常能提升最终指标：

from scipy import ndimage def remove_small_objects(mask, min_size=50): labeled, num_features = ndimage.label(mask) for i in range(1, num_features + 1): if (labeled == i).sum() < min_size: mask[labeled == i] = 0 return mask

min_size 的设置需要根据实际目标的大小来定。如果目标本身就很小，设太大会把真阳性也滤掉。我一般会先统计验证集上目标面积的分布，取一个偏小的值作为阈值。

另一个后处理技巧是形态学闭运算，填补掩码中的小孔洞。ndimage.binary_closing配合适当的结构元素就能做到。但要注意，如果目标本身就有孔洞结构（比如细胞核），闭运算会破坏这些细节，用之前要确认清楚。

6.3 从概率图到最终掩码的阈值选择

二分类模型输出的是 logits，经过 sigmoid 得到概率图。默认阈值是 0.5，但实际中这个阈值不一定最优。如果模型偏向于预测背景，可以适当降低阈值；反之则提高。

我通常会在验证集上遍历一组阈值（比如 0.3 到 0.7，步长 0.05），选 IoU 最高的那个。这个操作成本很低，但有时能带来 1-2 个百分点的提升。多分类的情况类似，可以对每个类别单独调阈值，但要注意类别之间的竞争关系。

7. 复现之后的一些个人体会

整套代码写下来，最深的感受是：UNet 的结构本身并不复杂，真正花时间的是数据管道、损失函数调参和后处理这些"外围"工作。网络结构可能只占整个项目 20% 的工作量，剩下 80% 都在处理数据和调优。

另一个体会是，不要迷信论文里的超参数。原始论文的输入尺寸、通道数、优化器设置都是针对特定数据集的，直接搬到自己的任务上未必最优。我建议先跑通一个 baseline，然后根据验证集的表现逐步调整。每次只改一个变量，记录结果，这样才能搞清楚每个改动到底有没有用。

最后，复现的价值不在于得到和论文一模一样的数字，而在于理解每个设计决策背后的权衡。当你下次遇到一个新的分割任务时，能快速判断该用什么样的结构、什么样的损失函数、什么样的数据增强策略，这才是复现真正的收获。