企业官网建设流程全解析

1. 为什么选择复现 UNet 而不是直接调库

1.1 从一次医学图像分割任务说起

去年接手一个细胞切片分割的小项目，图像分辨率不高，但边界要求特别细。当时第一反应是找现成的分割库，试了几个开箱即用的方案，结果在细胞粘连区域的表现都不太理想——要么把两个挨着的细胞糊成一块，要么边缘锯齿严重。后来翻论文才发现，这类密集预测任务里，UNet 的编码器-解码器加跳跃连接结构，天生就是为这种“既要全局语义、又要像素级定位”的场景设计的。

于是决定自己用 PyTorch 复现一遍。说实话，网上 UNet 的实现代码一抓一大把，但真正跑起来能复现论文效果、且代码结构清晰可扩展的并不多。很多实现要么把通道数写死，要么跳跃连接的处理方式跟原论文有出入，要么上采样用简单的插值糊弄过去。我花了大概两周时间，从论文结构图出发，一行行把网络搭起来，中间踩了不少坑，也积累了一些调参和调试的经验。这篇文章就把整个复现过程拆开来讲，从结构设计到代码实现，再到训练技巧和问题排查，尽量把每个环节的“为什么”说清楚。

1.2 复现的核心目标与适用人群

这次复现的目标很明确：用 PyTorch 从零搭建一个结构上忠实于原论文的 UNet，支持任意输入通道和类别数，能在自定义数据集上直接训练，并且代码要足够干净，方便后续替换编码器或加入注意力机制。适合的人群包括：刚接触医学图像分割的学生、想理解 UNet 内部细节的算法工程师、以及需要把 UNet 作为 baseline 做对比实验的研究者。如果你已经能熟练调用现成库，但对其内部机制一知半解，那这篇内容应该能帮你把那些模糊的地方补上。

2. UNet 结构拆解与设计思路

2.1 编码器-解码器架构的核心逻辑

UNet 的结构图看起来像一个大写的 U，左边是编码器，右边是解码器，中间用跳跃连接把对应层级的特征图拼起来。编码器的作用是逐步下采样，提取从局部纹理到全局语义的多尺度特征；解码器则是逐步上采样，把低分辨率的语义特征恢复到原图尺寸，同时利用跳跃连接引入编码器保留的高分辨率细节。

为什么非要这么设计？因为分割任务本质上是一个像素级分类问题，既需要知道“这个区域是什么”（语义），又需要知道“边界在哪里”（定位）。普通的卷积网络下采样几次后，语义信息够了，但空间细节丢得差不多了；而如果不下采样，感受野又不够，语义信息提取不充分。UNet 的跳跃连接相当于给解码器开了一条“捷径”，把编码器里还没被下采样破坏的细节直接送过来，让解码器在恢复分辨率时有据可依。

具体到通道数，原论文的编码器通道依次是 64、128、256、512、1024，每次下采样通道数翻倍。解码器则是 1024、512、256、128、64，每次上采样通道数减半。这个设计不是随便定的，通道数翻倍是为了在空间分辨率降低时，用更多的通道来补偿信息量的损失，保持网络整体的表达能力。

2.2 跳跃连接的具体实现方式

跳跃连接是 UNet 的灵魂，但实现时有几个细节容易出错。原论文里，编码器每一层的特征图在送入下一层之前，会先被保存下来，然后在解码器对应层级上采样之后，与解码器的特征图在通道维度上拼接。注意是拼接，不是相加。拼接意味着通道数会翻倍，所以解码器后续的卷积层输入通道数要相应调整。

另一个细节是裁剪。原论文因为卷积没有 padding，每次卷积后特征图会缩小，所以跳跃连接时需要把编码器的特征图中心裁剪到和解码器特征图一样的尺寸。但现在大家普遍用 padding=1 的 3x3 卷积，特征图尺寸不变，所以裁剪这一步可以省掉。我在复现时选择了带 padding 的卷积，这样代码更简洁，也不用担心尺寸对不齐的问题。

还有一个容易忽略的点：跳跃连接是在 ReLU 之前还是之后？原论文的结构图里，编码器的输出是先经过两次卷积加 ReLU，然后才保存下来用于跳跃连接。解码器那边，上采样后的特征图先和编码器特征拼接，再经过两次卷积加 ReLU。这个顺序会影响梯度的流动，我实测下来按原论文的顺序效果最稳。

2.3 上采样方式的选择与对比

上采样是解码器的关键操作，常见的有三种方式：最近邻插值、双线性插值和转置卷积。原论文用的是转置卷积，也就是反卷积，通过可学习的参数来放大特征图。但转置卷积有个毛病，容易产生棋盘格伪影，尤其是在通道数较多的时候。

我一开始用的转置卷积，训练到中期发现分割结果里有一些规律的网格状噪声，排查了很久才定位到是转置卷积的步长和卷积核尺寸不匹配导致的。后来换成双线性插值加一个 1x1 卷积，效果反而更干净。双线性插值的优势是参数少、不会引入额外伪影，1x1 卷积负责调整通道数，让解码器保持可学习的上采样能力。

如果你追求跟原论文完全一致，那就用转置卷积，但要注意把 kernel_size 设为 2、stride 设为 2，这样输出尺寸刚好翻倍，不会出现重叠或遗漏。如果更看重实际效果和稳定性，双线性插值加 1x1 卷积是更省心的选择。

3. PyTorch 代码实现与关键模块

3.1 双卷积模块的封装

UNet 里反复出现“两次 3x3 卷积 + ReLU”的结构，把它封装成一个模块能让代码清爽很多。这个模块的输入是通道数，输出也是通道数，中间可以加 BatchNorm 也可以不加。原论文没用 BatchNorm，但实际训练时加上 BatchNorm 能加速收敛，尤其是 batch size 比较小的时候。

import torch import torch.nn as nn 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，偏置项会被归一化抵消，留着反而浪费参数。mid_channels 参数是为了在某些变体里让中间通道数不等于输出通道数，标准 UNet 里用不到，但留着方便扩展。

3.2 下采样与上采样模块

下采样用最大池化，kernel_size=2、stride=2，这样特征图尺寸刚好减半。上采样我选了双线性插值加 1x1 卷积的方案，代码里用 nn.Upsample 加 nn.Conv2d 实现。

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 = nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x = torch.cat([x2, x1], dim=1) return self.conv(x)

注意 forward 里的 padding 操作。虽然理论上用 padding=1 的卷积不会改变尺寸，但实际中如果输入图像的尺寸不是 16 的倍数，经过几次下采样再上采样后，尺寸可能对不上。这个 padding 就是用来兜底的，把上采样后的特征图补齐到和编码器特征图一样的尺寸。这个细节很多实现里没有，但实际跑自定义数据时经常遇到，加上能省不少调试时间。

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

把各个模块拼起来就是完整的 UNet。编码器部分依次是 DoubleConv、Down、Down、Down、Down，解码器部分依次是 Up、Up、Up、Up，最后接一个 1x1 卷积输出类别数。

class UNet(nn.Module): def __init__(self, n_channels, n_classes, bilinear=True): super(UNet, self).__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) logits = self.outc(x) return logits

通道数的配置跟原论文基本一致，唯一的变化是 bilinear 模式下最底层的通道数从 1024 变成了 512，因为双线性插值不消耗通道，而转置卷积会把通道数减半。这个调整是为了让解码器的输入通道数匹配，不然拼接时会报错。

4. 训练流程与数据准备

4.1 数据集组织与预处理

UNet 对输入尺寸没有硬性要求，但为了下采样和上采样能对齐，最好把图像尺寸处理成 16 的倍数。我一般把短边缩放到 256 或 512，长边按比例缩放后中心裁剪或 padding 到相同尺寸。医学图像里常见的是灰度图，单通道输入，n_channels 设为 1；如果是彩色图就设为 3。

数据增强方面，水平翻转、垂直翻转、随机旋转 90 度、随机裁剪都是安全的选择。注意不要用会改变像素值分布的增强，比如色彩抖动，因为分割任务对像素值比较敏感。标签图也要跟着做同样的几何变换，用最近邻插值，避免引入不存在的类别。

from torch.utils.data import Dataset from PIL import Image import torchvision.transforms.functional as TF import random class SegmentationDataset(Dataset): def __init__(self, image_paths, mask_paths, size=256): self.image_paths = image_paths self.mask_paths = mask_paths self.size = size def __len__(self): return len(self.image_paths) def __getitem__(self, idx): image = Image.open(self.image_paths[idx]).convert('L') mask = Image.open(self.mask_paths[idx]).convert('L') image = TF.resize(image, (self.size, self.size)) mask = TF.resize(mask, (self.size, self.size), interpolation=Image.NEAREST) if random.random() > 0.5: image = TF.hflip(image) mask = TF.hflip(mask) image = TF.to_tensor(image) mask = torch.from_numpy(np.array(mask)).long() return image, mask

4.2 损失函数与优化器选择

分割任务最常用的损失是交叉熵，但医学图像里经常有类别不平衡的问题，比如背景像素远多于前景。这时候可以在交叉熵里加权重，或者用 Dice Loss。我一般先用带权重的交叉熵跑一版，如果前景分割效果不好再换成 Dice Loss 或两者的组合。

import torch.nn.functional as F def dice_loss(pred, target, smooth=1.): pred = torch.sigmoid(pred) pred = pred.contiguous().view(pred.size(0), -1) target = target.contiguous().view(target.size(0), -1) intersection = (pred * target).sum(dim=1) dice = (2. * intersection + smooth) / (pred.sum(dim=1) + target.sum(dim=1) + smooth) return 1 - dice.mean()

优化器用 Adam 就行，学习率从 1e-3 开始，配合 ReduceLROnPlateau 在验证损失不下降时减半。batch size 根据显存来，8 或 16 都可以。如果显存不够，可以用梯度累积模拟大 batch。

4.3 训练循环与验证指标

训练循环里除了常规的前向传播、反向传播、参数更新，还要注意两点：一是每个 epoch 结束后在验证集上算一下 IoU 或 Dice 系数，二是保存验证指标最好的模型权重，而不是最后一个 epoch 的权重。

def train_one_epoch(model, loader, optimizer, criterion, device): model.train() epoch_loss = 0 for images, masks in loader: images = images.to(device) masks = masks.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, masks) loss.backward() optimizer.step() epoch_loss += loss.item() return epoch_loss / len(loader) def evaluate(model, loader, device): model.eval() iou_sum = 0 with torch.no_grad(): for images, masks in loader: images = images.to(device) masks = masks.to(device) outputs = model(images) preds = torch.argmax(outputs, dim=1) intersection = ((preds == 1) & (masks == 1)).sum().float() union = ((preds == 1) | (masks == 1)).sum().float() iou_sum += (intersection / (union + 1e-8)).item() return iou_sum / len(loader)

IoU 的计算要注意加一个极小值防止除零。如果有多类别，就对每个类别分别算 IoU 再取平均。

5. 实操中遇到的坑与排查技巧

5.1 尺寸不匹配导致的报错

这是复现 UNet 时最常见的问题。报错信息通常是 “Sizes of tensors must match except in dimension 1”，意思就是跳跃连接时编码器和解码器的特征图尺寸对不上。原因一般是输入图像的尺寸不是 16 的倍数，经过四次下采样后尺寸变成了奇数，再上采样回来就差了 1 个像素。

解决办法有两个：一是预处理时把图像尺寸 pad 到 16 的倍数，二是在 Up 模块的 forward 里加上动态 padding，就像前面代码里写的那样。我推荐两个都做，预处理保证大部分情况不出问题，动态 padding 作为兜底。

5.2 损失不下降或震荡

如果训练了几个 epoch 损失还在原地打转，先检查学习率是不是太大。Adam 的默认学习率 1e-3 对 UNet 来说有时候偏大，可以试试 1e-4。另外检查一下数据归一化，图像像素值要缩放到 0 到 1 之间，用 ToTensor 会自动做这件事，但如果你手动处理数据就要注意。

如果损失震荡得厉害，可能是 batch size 太小，梯度噪声大。可以增大 batch size，或者用梯度累积。还有一个容易被忽略的点是 BatchNorm 在 batch size 很小时表现不稳定，如果 batch size 小于 8，可以考虑把 BatchNorm 换成 GroupNorm。

5.3 分割结果出现棋盘格或网格伪影

这个问题通常跟转置卷积有关。转置卷积的 kernel_size 和 stride 如果不匹配，输出特征图上会出现周期性的重叠或遗漏，表现为棋盘格。解决办法是确保 kernel_size 能被 stride 整除，比如 kernel_size=2、stride=2，或者 kernel_size=4、stride=2。如果换了双线性插值还有伪影，那可能是数据增强里的 resize 操作引入了混叠，把插值方式改成双线性或双三次。

5.4 常见问题速查表

5.5 几个提升效果的小技巧

第一个技巧是深度监督。在解码器的每一层上采样后都接一个 1x1 卷积输出辅助预测，然后把这些辅助损失加权加到总损失里。这样能让梯度更好地传到浅层，加速收敛。第二个技巧是测试时增强，把输入图像翻转、旋转后分别预测，再把结果平均回去，通常能提升 1 到 2 个点的 IoU。第三个技巧是后处理，对预测的 mask 做连通域分析，去掉面积过小的孤立区域，这在细胞分割里特别有用。

6. 模型扩展与后续优化方向

6.1 替换编码器为预训练骨干

原版 UNet 的编码器是从零训练的，如果数据集比较小，收敛会很慢。可以把编码器换成在 ImageNet 上预训练过的 ResNet 或 EfficientNet，只训练解码器部分，或者用较小的学习率微调编码器。这样能在小数据集上快速拿到不错的效果。替换时注意把 ResNet 的最后一层全连接去掉，取最后一个卷积块的输出作为编码器输出，然后接到 UNet 的解码器上。

6.2 加入注意力机制

UNet 的跳跃连接是把编码器特征直接拼到解码器，没有做任何筛选。如果编码器里有一些背景噪声，也会被原封不动地传过去。加入注意力门控可以让解码器在拼接前对编码器特征做一次加权，抑制无关区域，突出目标区域。注意力门控的实现不复杂，就是一个 1x1 卷积加 Sigmoid 得到权重图，然后和编码器特征逐元素相乘。

6.3 多类别与多通道输出的处理

如果分割任务有多个类别，输出通道数设为类别数，损失函数用交叉熵，预测时取 argmax。如果输入是多通道，比如 RGB 加深度图，把 n_channels 设为对应通道数就行，第一层卷积会自动适配。注意如果不同通道的数值范围差异很大，要做归一化，不然训练会不稳定。

6.4 部署时的推理优化

训练好的模型如果要部署到实际系统里，推理速度是个要考虑的问题。可以用 TorchScript 把模型导出成静态图，或者用 ONNX 导出后在其他推理引擎上运行。如果对速度要求高，可以把双线性插值换成最近邻插值，或者减少解码器的通道数。另外，推理时把模型设为 eval 模式，关闭 BatchNorm 的统计更新和 Dropout，能省一点计算量。

7. 我个人在实际操作中的体会

复现 UNet 这件事，代码本身不算复杂，真正花时间的是调试和验证。我最大的体会是：不要一上来就追求跟原论文完全一致，先把最简版本跑通，确认数据管道没问题、损失能下降、输出尺寸对得上，然后再逐步加 BatchNorm、加数据增强、换损失函数。每一步只改一个变量，这样出了问题容易定位。

另外，可视化真的很重要。训练过程中定期把预测结果和真实标签叠在一起看，能直观地发现模型是在学边界还是在学区域，是在过拟合还是在欠拟合。我习惯每 5 个 epoch 存一张对比图，训练结束后翻一遍，比看损失曲线有用得多。

最后再分享一个小技巧：如果显存不够，除了减小 batch size，还可以把输入尺寸从 512 降到 256，UNet 对输入尺寸不敏感，256 的输入在多数任务上已经够用了。等模型结构调好了，再逐步增大输入尺寸做微调，这样能在有限资源下把实验跑起来。