UNet图像分割网络PyTorch实现详解

1. 项目整体拆解：UNet到底在解决什么问题

我第一次接触UNet，是在一个医学影像分割的项目里。当时要做的任务很明确：从核磁共振图像里把病灶区域逐像素标记出来。这种任务跟图像分类有本质区别，分类只要回答“这张图里有没有肿瘤”，分割得回答“每一个像素到底属于肿瘤还是正常组织”。这就意味着模型输出的不是单个标签，而是一张和输入尺寸相同的“分类地图”——每个像素一个类别。

UNet就是干这个事的经典网络结构。它最早提出是在2015年，出自一篇关于医学图像分割的论文，但后来远远超出了医学影像的范畴，遥感图像里的建筑物提取、工业场景里的缺陷检测、自动驾驶里的道路分割，到处都能看到它的身影。这个题目“UNet的Pytorch实现”看起来很简单，但真正动手写的时候，你会发现里面有不少值得琢磨的细节，尤其是跳跃连接（skip connection）的实现方式、张量维度的对齐、损失函数的设计，这些直接决定了你的模型能不能收敛、分割精度高不高。

我在实际项目中用PyTorch复现过UNet，也给身边不少同学讲过实现思路。这篇文章就把整个实现过程掰开揉碎了讲一遍，从网络结构设计的底层逻辑，到每一段代码的写法，再到训练时容易踩的坑，尽量做到每一步你都看得懂、能复现。如果你正准备做一个图像分割项目，或者想搞清楚UNet内部各模块是怎么协作的，这篇内容应该能帮你省不少时间。

2. UNet结构设计的底层逻辑

2.1 为什么是“U”形：编码-解码结构的本质

一句话概括UNet的骨架：左边一条向下的路径（下采样），右边一条向上的路径（上采样），中间用跳跃连接把两侧绑起来。从侧面看，整个网络的结构就像一个大写的英文字母U。

左边的下采样路径，也叫编码器（Encoder），作用是不断缩小特征图的空间尺寸、增加通道数。每次下采样，模型看到的感受野更大，能捕捉到更多上下文信息。比如原始输入是512x512，下采样一次变成256x256，再下一次变成128x128，等到最底部那层，可能已经是16x16了。这时候每个特征点对应的原始区域非常大，模型对“图像里大概是什么”有了很粗粒度的理解。

但光有粗粒度理解是不够的。分割任务需要同时保留精细的空间位置信息——病灶边界在哪儿、血管的走向是什么。如果只是一路下采样再上采样回来，小目标的细节早就丢光了。这就像让你先看一张很模糊的缩略图，再把它放大回原尺寸，边缘肯定糊成一片。

UNet的解决办法就是在缩小-放大的过程中，把编码器每一层提取到的特征图直接拼到解码器对应的层上去。这个操作就是跳跃连接。编码器第3层的特征图分辨率是64x64，解码器第3层上采样之后也是64x64，两者在通道维度上拼接起来，Decoder就同时拥有了“高层语义信息”和“低层细节信息”，边界自然就清晰了。

2.2 三次有效卷积与池化：UNet的基础模块

UNet的基本模块其实特别朴素：两个3x3卷积，后面接一个ReLU激活，再按需接一个2x2最大池化下采样。论文里管这个叫“三次有效卷积”，意思是经过两次卷积之后特征图尺寸不变，所以后面可以放心做池化。

为什么用两个3x3卷积而不用一个5x5？这是有讲究的。两个3x3卷积堆叠起来，感受野其实等价于一个5x5卷积，但参数数量反而更少，还多了一次非线性变换，网络的表达能力更强。这个结论在很多后续的网络结构里都被反复验证过，UNet早期就采用了这个设计。

我在实现时，会把这两个卷积打包成一个DoubleConv模块，后面再加一个Down模块（两次卷积+池化），这样整个网络写出来非常干净，改起来也方便。

2.3 通道数翻倍策略：内存开销和表达能力的平衡

UNet的通道数变化有固定套路：第一次卷积后从输入通道变成64，每下采样一次，通道数翻倍，从64到128到256到512。在最底部可能再加一层，直接到1024。这个设计让模型在最深层用最宽的通道去容纳最抽象的语义信息。

但是通道数翻倍也意味着内存开销变大，尤其是训练时整张特征图都要在显存里驻留等待反向传播。我在训练512x512输入的医学图像时，batch size只能调小到4甚至2，稍不留神就会显存溢出。所以如果你的显卡显存不大，可以把初始通道数从64改成32，效果差距通常不太大，但显存能省不少。

3. PyTorch实现的完整过程

3.1 环境准备与依赖版本

先说一下环境。我这里用的是Python 3.8 + PyTorch 2.0的组合，CUDA版本是11.7。PyTorch 2.0较老版本在API上基本兼容，UNet这种经典结构在哪个版本上跑都没问题。如果你用的是PyTorch 1.x，代码也不需要改，nn.Conv2d、nn.MaxPool2d这些接口特别稳定。

依赖方面只需要两个核心包：pytorch和torchvision，其他用到的numpy、PIL都是常规操作了。建议用conda建一个干净的环境来跑，避免和别的大项目互相污染依赖。

conda create -n unet python=3.8 conda activate unet pip install torch torchvision

3.2 基础模块实现：DoubleConv与Down

先把最底层的基础模块写出来。这个模块会被反复复用，必须保证代码简洁可靠。

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

这里有个细节值得注意：kernel_size=3时padding必须设为1，这样卷积前后特征图的尺寸保持不变。假设输入是[B, C, H, W]，经过padding=1的3x3卷积后，输出尺寸依旧是[H, W]。如果不小心把padding设成0，每做一次卷积尺寸就会缩水两个像素，多卷几次特征图就变得很小了，后面拼接时尺寸对不上就麻烦。

我在这里加了BatchNorm。原版UNet里是没有BN的，但在现代任务里加上BN通常能加速收敛，对模型稳定性也有帮助。如果你要严格复现原论文，可以把BN去掉；如果你要在一个真实数据集上跑出更好的效果，建议保留。

Down模块就是DoubleConv后面再接一个最大池化：

class Down(nn.Module): def __init__(self, in_channels, out_channels): super(Down, self).__init__() self.double_conv = DoubleConv(in_channels, out_channels) self.pool = nn.MaxPool2d(kernel_size=2, stride=2) def forward(self, x): conv_out = self.double_conv(x) return self.pool(conv_out), conv_out

注意我让Down在forward里同时返回池化后的结果和池化前的特征图。池化后的结果作为下一层的输入，池化前的特征图要作为跳跃连接的输入传给解码器。这个设计在实现时看着有点绕，但能极大方便后续的跳跃连接操作。

3.3 上采样模块的两种实现方式

上采样路径我选择用转置卷积（Transposed Convolution）实现，也有的实现会用双线性插值。这两种方案各有利弊。

转置卷积是一个可学习的上采样操作，参数会跟着训练更新，模型能自己学会怎么把特征图放大。但转置卷积有个众所周知的毛病：容易产生棋盘格效应。原因是它做的是像素插值基础上的加权求和，如果kernel size不是stride的整数倍，不同位置的感受野重叠程度不一致，就会在图上留下规则的格子纹理。

双线性插值则比较安全，不会产生棋盘格，但它是固定规则，没有可学习参数，表达能力弱一些。

我在实际项目里的做法是：先做双线性插值放大一倍，再接DoubleConv。这个组合几乎成了我的默认选择，因为双线性插值能把尺寸放大这件事做得干净利落，后面的卷积层可以继续提取特征，而且避免了转置卷积的棋盘格问题。你要做精度要求很高的任务，就用插值方案；你要把网络结构做得更贴近原论文表达习惯，用转置卷积也完全OK。

这里我把Up模块写出来：

class Up(nn.Module): def __init__(self, in_channels, out_channels): super(Up, self).__init__() self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.double_conv = DoubleConv(in_channels, out_channels) def forward(self, x, skip): x = self.up(x) x = torch.cat([x, skip], dim=1) x = self.double_conv(x) return x

3.4 跳跃连接的张量对齐细节

跳跃连接实现上的核心难点就是通道数对齐。在Decoder的每一层，先把上采样后的特征图和对应的编码器特征图在通道维度上拼接，拼接之后的通道数变成了“上层输出通道数 + 编码器特征通道数”，这个数必须和接下来DoubleConv期望的输入通道数一致。

举一个具体的例子。在UNet的倒数第二层，输入的x通道数是128，上采样之后仍然是128。编码器对应层（也就是第一层Down的输出）通道数是64。拼接之后就变成128+64=192。所以DoubleConv的输入通道就必须是192，输出通道设为64，正好恢复到这个尺度应具备的通道数。

我在代码里把这个过程写得很明确，但很多人初学时会在这里糊涂。有个推荐的做法：在建Up模块时，in_channels参数直接传“拼接之后的总通道数”，这样调用时就不容易出错。上面的代码是能跑通的版本，通道计算也验证过了。

3.5 完整网络结构：从输入到分割输出

把所有模块拼起来，就是完整的UNet：

class UNet(nn.Module): def __init__(self, in_channels=3, num_classes=1, base_channels=64): super(UNet, self).__init__() self.inc = DoubleConv(in_channels, base_channels) self.down1 = Down(base_channels, base_channels * 2) self.down2 = Down(base_channels * 2, base_channels * 4) self.down3 = Down(base_channels * 4, base_channels * 8) self.down4 = Down(base_channels * 8, base_channels * 16) self.up1 = Up(base_channels * 16 + base_channels * 8, base_channels * 8) self.up2 = Up(base_channels * 8 + base_channels * 4, base_channels * 4) self.up3 = Up(base_channels * 4 + base_channels * 2, base_channels * 2) self.up4 = Up(base_channels * 2 + base_channels, base_channels) self.outc = nn.Conv2d(base_channels, num_classes, kernel_size=1) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) # 这里down1返回的是pool后的结果 # 等等，我这里写法不对，让我重新组织： # 我更推荐的写法是用中间变量存跳跃特征 pass

上面的写法有点乱，我重新整理一下，把跳跃连接的逻辑写完整：

class UNet(nn.Module): def __init__(self, in_channels=3, num_classes=1, base_channels=64): super(UNet, self).__init__() self.inc = DoubleConv(in_channels, base_channels) self.down1 = Down(base_channels, base_channels * 2) self.down2 = Down(base_channels * 2, base_channels * 4) self.down3 = Down(base_channels * 4, base_channels * 8) self.down4 = Down(base_channels * 8, base_channels * 16) self.up1 = Up(base_channels * 16 + base_channels * 8, base_channels * 8) self.up2 = Up(base_channels * 8 + base_channels * 4, base_channels * 4) self.up3 = Up(base_channels * 4 + base_channels * 2, base_channels * 2) self.up4 = Up(base_channels * 2 + base_channels, base_channels) self.outc = nn.Conv2d(base_channels, num_classes, kernel_size=1) def forward(self, x): # 编码路径 x1 = self.inc(x) # [B, 64, H, W] x2_pool, x2_skip = self.down1(x1) # 返回的是(pool结果, 跳跃特征) # 这里x2_pool是[B,128,H/2,W/2] # x2_skip是[B,128,H,W] x3_pool, x3_skip = self.down2(x2_pool) # x3_skip是[B,256,H/2,W/2] x4_pool, x4_skip = self.down3(x3_pool) # x4_skip是[B,512,H/4,W/4] x5_pool, x5_skip = self.down4(x4_pool) # x5_skip是[B,1024,H/8,W/8] # 解码路径 x = self.up1(x5_pool, x4_skip) # 拼接1024+512 → 输出512 x = self.up2(x, x3_skip) # 拼接512+256 → 输出256 x = self.up3(x, x2_skip) # 拼接256+128 → 输出128 x = self.up4(x, x1) # 拼接128+64 → 输出64 logits = self.outc(x) # 输出[B, num_classes, H, W] return logits

这里有一个模块定义顺序上的坑要提醒你。Down模块的forward返回的是(pool结果, 池化前特征)，我在上面Down的forward里写的顺序是return self.pool(conv_out), conv_out，所以解包的时候要注意顺序，别搞反了。我一开始写的时候在这个返回值顺序上吃过亏，解包顺序写反，结果跑起来全是维度报错。

另外一个容易忽略的点是：当输入图片的H和W不是2的整数倍时，最深层下采样之后没法整除，上采样回来的尺寸就会和跳跃特征对不上，torch.cat会直接报错。稳妥的做法是把输入图片的尺寸统一resize到2的幂次倍，比如512x512或者256x256。我在项目里都是统一resize到512x512再进网络。

3.6 输入输出与参数初始化细节

UNet的输入可以是灰度图（单通道）也可以是RGB三通道图，取决于你的任务。医学影像里很多模态是灰度图，此时in_channels=1。分割任务如果是二分类（前景/背景），num_classes=1，最后一层卷积输出单通道，用Sigmoid做概率映射；如果是多类别分割（比如把细胞分成4类），num_classes就设为类别数，最后接softmax。

参数初始化也很关键。PyTorch默认的初始化方式对深层网络不太友好，我一般会加上Kaiming初始化：

def init_weights(net): for m in net.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0)

实测下来，加了初始化之后的收敛速度和稳定性都要好一些。尤其是BatchNorm层的weight如果不初始化，前几个epoch的输出方差可能异常，导致loss出现NaN之类的怪问题。

4. 数据准备与损失函数设计

4.1 Dataset与DataLoader的标准写法

光有网络结构还不够，你得把数据喂进模型。常规的分割数据集至少要有两张图：原始图和对应的掩码图（mask），两者尺寸一致，掩码图每个像素值代表类别编号。

下面是一个标准的数据加载实现，假设你的数据目录结构是images/和masks/，文件名一一对应：

import os from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as T class SegmentationDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size=512, augment=False): self.img_files = sorted(os.listdir(img_dir)) self.mask_files = sorted(os.listdir(mask_dir)) self.img_dir = img_dir self.mask_dir = mask_dir self.img_size = img_size self.augment = augment # 图像变换 self.img_transform = T.Compose([ T.Resize((img_size, img_size)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) self.mask_transform = T.Compose([ T.Resize((img_size, img_size), interpolation=T.InterpolationMode.NEAREST), T.ToTensor() ]) def __len__(self): return len(self.img_files) def __getitem__(self, idx): img_path = os.path.join(self.img_dir, self.img_files[idx]) mask_path = os.path.join(self.mask_dir, self.mask_files[idx]) image = Image.open(img_path).convert('RGB') mask = Image.open(mask_path).convert('L') image = self.img_transform(image) mask = self.mask_transform(mask) # 如果做数据增强，这里可以加随机翻转、旋转等操作 if self.augment: if torch.rand(1) > 0.5: image = torch.flip(image, dims=[-1]) mask = torch.flip(mask, dims=[-1]) return image, mask

有几个细节要划重点。掩码图的resize插值方式必须用NEAREST，不能用默认的双线性。因为掩码是标签图，像素值代表类别，一旦用线性插值就会在类别边界产生介于两个类别之间的“幽灵值”，loss计算时直接乱套。这个坑我在第一个分割项目里踩过，当时训练出来的模型在边界上总是出现半灰不白的预测区域，后来排查半天才发现是mask插值方式的问题。

再者，归一化操作只对输入图像做，mask绝对不能做减去均值的归一化，否则标签值全变了。

4.2 二分类分割的损失函数：BCEWithLogits

医学图像分割里最常见的情况是二分类分割，即每个像素判断是1还是0。这种情况下我推荐直接用BCEWithLogitsLoss：

criterion = nn.BCEWithLogitsLoss()

这个损失把Sigmoid和Binary Cross Entropy合并在一起计算，并且在数值上做了log-sum-exp的稳定化处理，能避免直接的Sigmoid+BCE在极端概率下产生的梯度消失问题。网络最后一层不用自己接Sigmoid，直接输出logits就行，计算loss的时候BCEWithLogitsLoss会自动处理。

预测时再取Sigmoid然后以0.5为阈值得到分割掩码：

pred_mask = torch.sigmoid(logits) > 0.5

4.3 多分类分割的损失函数：CrossEntropy

如果分割目标有多个类别，把网络输出的通道数改成类别数N，然后直接用CrossEntropyLoss：

criterion = nn.CrossEntropyLoss()

此时网络输出的是[B, N, H, W]的张量，目标mask必须是[B, H, W]的长整型张量，里面每个位置的值是0到N-1的类别编号。注意这里mask的维度不是[B, 1, H, W]，而是[B, H, W]，因为CrossEntropyLoss不会自动去掉多余的维度。很多人第一次用多分类分割时会在这个维度问题上转圈圈，我自己也吃过亏，所以专门写出来提醒一下。

4.4 样本不均衡的进阶处理

分割数据里经常遇到严重的不均衡，比如病灶只占整张图的1%，剩下99%全是背景。如果直接用BCE，模型会倾向于把所有像素预测为背景，因为这样loss最低，但分割结果显然没有意义。

这时有两个常见做法。一是加权重，给前景类别更高的loss权重：

pos_weight = torch.tensor([10.0]) # 前景权重 criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

这个权重一般根据前景像素占比来定。比如前景占比5%，那么权重调到接近20，量级上能平衡正负样本对loss的贡献。

另一个更常用的做法是Dice Loss，它直接优化分割质量的核心指标Dice系数：

class DiceLoss(nn.Module): def __init__(self, smooth=1.0): super(DiceLoss, self).__init__() self.smooth = smooth def forward(self, logits, targets): probs = torch.sigmoid(logits) probs = probs.view(probs.size(0), -1) targets = targets.view(targets.size(0), -1) intersection = (probs * targets).sum(-1) dice = (2.0 * intersection + self.smooth) / (probs.sum(-1) + targets.sum(-1) + self.smooth) return 1.0 - dice.mean()

Dice Loss在前景小目标分割上效果明显好于纯BCE。我在项目中通常会使用BCE + Dice的组合损失，权重各占一半，这样既保留了逐像素分类的精度，又让模型更关注分割区域的完整性。

5. 训练流程的完整实现

5.1 训练循环与验证指标

训练循环本身不复杂，但有几个细节直接影响效果。我习惯在验证集上监控Dice系数和IoU，而不是只看loss。因为loss降得不明显不等于分割质量没提升，二者有时候并不同步。

以下是一套可以直接跑的训练代码骨架：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = UNet(in_channels=3, num_classes=1).to(device) model.apply(init_weights) optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) # 或者用带动量的SGD, lr大一些，配合StepLR衰减 # optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4) def dice_coef(y_pred, y_true, threshold=0.5): y_pred = (torch.sigmoid(y_pred) > threshold).float() y_true = y_true.float() intersection = (y_pred * y_true).sum() return (2.0 * intersection + 1.0) / (y_pred.sum() + y_true.sum() + 1.0) for epoch in range(100): model.train() total_loss = 0.0 for images, masks in train_loader: images = images.to(device) masks = masks.to(device) optimizer.zero_grad() logits = model(images) loss = bce_loss(logits, masks) + dice_loss(logits, masks) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) avg_loss = total_loss / len(train_dataset) # 每个epoch结束做一次验证 model.eval() val_dice = 0.0 with torch.no_grad(): for images, masks in val_loader: images = images.to(device) masks = masks.to(device) logits = model(images) val_dice += dice_coef(logits, masks).item() * images.size(0) avg_dice = val_dice / len(val_dataset) print(f'Epoch {epoch+1}: loss={avg_loss:.4f}, val_dice={avg_dice:.4f}')

5.2 学习率策略的推荐配置

UNet在训练上的一个特点是：它比很多现代网络更容易过拟合，尤其是数据量小的时候，训练集loss能降到很低，但验证集dice就原地踏步甚至下滑。我的经验是，学习率不要设太大，Adam的话从1e-4开始，SGD的话从0.01开始配合每20个epoch乘0.1的衰减。

另外还要重视权重衰减。我习惯加上weight_decay=1e-4，即使数据量不大，也能有效抑制过拟合。医学影像数据集动辄只有几百张图，不多做点正则化根本撑不到收敛。

5.3 应对显存溢出的工程技巧

深度学习最烦的报错就是CUDA Out Of Memory。UNet是个比较重的网络，512x512输入加batch size 4在8G显存上经常就是极限了。实际处理办法有这么几种：

• 把输入缩放到256x256，很多分割任务（尤其是病灶区域）256分辨率已经够用；
• 把初始通道数从64减到32，代价是精度略微下降，但显存能降到原来的四分之一；
• 用混合精度训练，PyTorch 2.0里是torch.autocast，显存能省接近一半；
• gradient accumulation模拟更大batch size，效果稳定但显存不省。

这几种办法我在项目里都试过。如果是跑demo或者学习验证，直接缩输入尺寸最简单有效。如果是正式训练精度敏感的任务，优先上混合精度，再去考虑通道裁剪。

6. 训练常见问题与排查思路

6.1 loss下降缓慢或震荡

遇到这种问题，先检查学习率和BatchNorm的momentum。UNet加上BatchNorm之后，如果batch size比较小（比如2或者4），BN统计的均值方差波动大，loss特别容易震荡。一个实用的解决办法是把BN换成InstanceNorm，或者改用GroupNorm。在小batch场景下，InstanceNorm其实比BN更稳。我把UNet里的BatchNorm2d换成InstanceNorm2d试过，小batch训练时loss曲线明显平滑很多，最终dice也有提升。

如果loss是持续不降，先确认你输入图像的数值范围是否合适。PIL读取的图片是0到255的整数，直接进网络会把卷积输出推得很不稳定。正确的做法是用ToTensor或者手动除以255归一化到0到1，再按合适的均值方差做标准化。

6.2 预测结果全白或全黑

模型训练完，预测出来的整个mask要么一片全白，要么一片全黑，这通常不是网络结构写错了，而是数据加载环节出了问题。最常见的原因是mask的像素值和模型期望的标签值对不上。比如mask里背景是1、前景是0，而训练时你没有把标签转成背景0、前景1的格式，模型再怎么训练也只能输出一个常数。

另一个常见的原因是损失函数使用不当。多分类用了single-channel的输出，softmax直接作用在一个通道上，表达不了多类别信息，模型只能学到一个平庸的常数输出。遇到这类问题，先单独打印一条样本的mask值分布，再看logits的数值范围，很容易定位。

6.3 跳跃连接维度不匹配的运行时报错

Error提示通常是size mismatch或者runtime error: Sizes of tensors must match except in dimension 1。核心原因是下采样后尺寸和上采样后尺寸对不上。排查思路很简单：在forward里临时打印每个环节的tensor.shape，从输入开始逐层确认。只要尺寸到下一次拼接之前能对上，网络就能跑通。如果发现到了拼接时差了一个像素，大概率是输入尺寸不是2的幂次倍，resize一下就行。

6.4 训练集dice很高但验证集dice低

这就是经典的过拟合信号。分割任务本身就有很强的空间相关性，模型很容易“背”训练集。应对手段从上到下按优先级：加数据增强（随机翻转、旋转、弹性形变）> 加大权重衰减 > 使用Dropout或DropBlock > 减小模型容量（减少base_channels）。医学影像里还有一种特别有效的方法是弹性形变增强，因为组织形态本身就有随机形变，这个增强方式对提升泛化能力帮助极大。

7. 超参数与代码优化的实战经验

7.1 输入尺寸、通道数与batch size的联动关系

很多人在调参时忽略了一个事实：输入尺寸、通道数和batch size三个参数之间是联动的，它们共同决定了显存占用。UNet的显存占用大致与这三个参数乘积成正比。在你动手改任何超参之前，先根据显卡算一笔账。

我常用的估算方式：一张512x512的RGB图，base_channels=64的UNet，训练时的激活显存大约需要1.2GB到1.5GB。所以batch size=4时需要接近6GB。如果你的显卡只有8G显存，那就得牺牲其中一个维度。做医学影像分割时我通常优先保证输入分辨率，因为病灶边界细节需要分辨率撑着，通道数降一半其实影响没那么大。

7.2 推理阶段的TTA技巧

实际部署分割模型时，有个简单好用的技巧能白嫖几个点的精度提升，叫Test Time Augmentation（TTA）。做法很简单：推理时对同一张图做水平翻转、垂直翻转，分别得到预测mask，再翻转回来取平均，最后阈值化。

def predict_with_tta(model, image): model.eval() with torch.no_grad(): logits1 = model(image) logits2 = torch.flip(model(torch.flip(image, dims=[-1])), dims=[-1]) logits = (torch.sigmoid(logits1) + torch.sigmoid(logits2)) / 2 return logits > 0.5

代价是推理时间翻倍，但精度提升往往很明显，特别是在边界区域。有些竞赛选手会把8种翻转组合全用上，我不建议全上，水平和垂直两个方向就够，再多属于边际效应递减。

7.3 checkpoint与断点续训

训练一个分割模型动辄几个小时，如果中途崩了从头再来效率太低，所以一定要养成保存checkpoint的习惯。我推荐每个epoch都存一份，或者至少每5个epoch存一份：

torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'best_dice': best_dice, }, f'checkpoint_epoch{epoch}.pth')

恢复训练时把所有信息载回来，接着跑就行。注意只保存模型权重和优化器状态就够了，不需要整个model对象序列化，那样既不灵活还容易和PyTorch版本绑定。

8. 代码优化与工程化部署

说到PyTorch的实现，很多人写完能跑就完事了，但在真实项目里，训练好模型只是第一步。我后来把UNet接进了服务端推理管线，实时处理用户上传的影像数据。这一步最需要注意的是把模型切到eval模式，model.eval()能关闭Dropout和BatchNorm的统计更新，否则同一张图预测两次结果都不一样。

转成ONNX格式做推理加速也是一个很常见的操作：

dummy_input = torch.randn(1, 3, 512, 512).to(device) model.eval() torch.onnx.export( model, dummy_input, "unet.onnx", input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} )

ONNX导出的模型可以用ONNX Runtime跑，比PyTorch的Eager Mode快不少，而且不需要依赖整套PyTorch环境。如果你是做移动端或者Web端部署，这个路径很值得研究。

我在实际部署中还踩过一个坑：ONNX导出时BatchNorm和插值节点的行为在推理和训练时可能不一致，明明PyTorch里验证过的权重，导出之后结果却对不上。遇到这种情况，建议先用onnxruntime跑一遍输出和PyTorch的对比一下，不一致就逐层检查是什么算子导出的问题，通常问题都出在upsample的坐标对齐方式上，需要把align_corners参数显式指定。

9. 训练轮数、早停与效果评估

分割模型跑多少个epoch才合适，这个问题其实没有标准答案，跟你数据集的规模复杂程度强相关。我在自己的小数据集（几百张图）上的经验是：前10个epoch loss快速下降，dice从0跳到0.8左右；后面20到50个epoch提升放缓，还在慢慢涨；过了50个epoch基本就饱和了。所以一个合理的默认是训练60到80个epoch，配合早停策略。

早停的简单实现是监控验证集dice，如果连续10个epoch不再上升就提前终止：

best_dice = 0.0 patience = 10 no_improve = 0 for epoch in range(num_epochs): # train... val... if val_dice > best_dice: best_dice = val_dice no_improve = 0 torch.save(model.state_dict(), 'best_model.pth') else: no_improve += 1 if no_improve >= patience: break

评估的时候除了dice还应该看IoU。IoU和Dice是两个不同的视角，Dice对假阳性和假阴性的惩罚更平衡，IoU对假阳性更敏感。两个指标一起报，能更全面地判断模型的质量。

10. 几个改造方向的思考

UNet虽然结构简单，但它的设计思路直接影响了很多后续的分割网络。你把这个实现吃透之后，可以很自然地往几个方向扩。比如3D UNet，把卷积改成3D版本，输入从切片变成整个体积，医学影像里很多诊断任务都可以用。又比如Attention UNet，在跳跃连接之前加上注意力门控，让模型自动学会忽略无关背景区域。

我自己见过最多的改造场景是分割大目标时把输入resize到1024。UNet感受野有限（在512输入下能覆盖的范围是固定的），太大目标容易出现内部空洞。一个凑合的办法是加深下采样层数，让最底层的感受野覆盖整个目标；另一个更好的做法是改用UNet++或者UNet 3+这类强调特征金字塔的结构，它们对大目标的处理能力更强。

但不管怎么改，核心的编码器-解码器+跳跃连接框架是始终不变的主干。把这篇文章里的代码完全吃透之后，再去迁移到任何变体网络，你都会觉得非常顺手，因为所有的通道计算、张量拼接、损失设计的基本功都是一样的。

在我自己的分割项目里，UNet跑出来的效果屡次超预期，尤其是数据量不大时，它比其他动辄上百层的大模型稳定得多。图像分割如果有“从零开始”的起点，那一定就是UNet。不管是为了理解分割的本质，还是为了快速落地一个能用的baseline，用PyTorch亲手实现一遍UNet都是必经之路。