1. 一张512×512的切片为什么普通卷积网络啃不动第一次把UNet的代码从头敲完是在一个图像分割的小任务上。当时手上只有几百张标注好的切片背景面积占比超过九成前景目标边缘还特别碎。我一开始图省事拿了个常规的分类骨干网络加两层卷积直接出结果训练loss降得挺漂亮但把预测图叠回原图一看目标边界的轮廓糊成一团小目标要么整块消失要么被连成大块。那一刻我才真正理解UNet这套结构为什么要长成U形——它解决的不是能不能分出目标的问题而是分出来的边界够不够准的问题。UNet最早是为生物医学图像分割提出的全卷积网络核心构架由编码器下采样路径、解码器上采样路径和跳跃连接三部分组成。它解决的核心矛盾很具体深层特征语义强但分辨率低浅层特征分辨率高但语义弱。分类网络只关心最后那个向量所以一路把特征图压小没问题分割网络要求输出和输入同尺寸的逐像素标签那就必须把分辨率还原回去。而单纯靠上采样还原丢掉的空间细节是找不回来的UNet的解法是在还原过程中把编码器同层的特征图直接搬过来拼接。这篇内容适合两类人看一类是已经会用现成库调模型、但想搞清楚内部每一层到底在干什么的开发者另一类是准备面试、需要白板手推结构或者现场手敲代码的同学。我会把网络结构、张量形状的推导、手敲代码的完整实现、训练侧的几个关键取舍一步一步拆开讲。代码用PyTorch写尽量做到复制粘贴就能跑同时把每个设计选择背后的理由说清楚而不是只给你一堆层。在动手之前先明确一件事UNet的手敲不是背代码而是能在纸上画出数据流能说清楚任意一层输入输出的通道数和空间尺寸能解释为什么拼接之后通道数要减半。这三点想透了代码自然就写出来了而且换个输入尺寸、换个通道基数也不会翻车。2. 把UNet拆成编码、瓶颈、解码三段通道与分辨率怎么走2.1 编码器负责看懂是什么解码器负责说清在哪里编码器的任务和普通卷积分类网络完全一致通过反复的卷积加下采样把感受野撑大让网络看到更大范围的上下文。UNet原始论文里用了4次下采样每次分辨率减半、通道数翻倍。第一次下采样后特征图尺寸是输入的1/2通道从1或3变成64第二次是1/4、128第三次1/8、256第四次1/16、512。到最底部瓶颈层再做一次卷积通道到1024分辨率保持1/16。解码器则是反过来每一步先做上采样把分辨率放大一倍然后把编码器对应层的特征图拿过来在通道维度上拼接再接两层卷积做特征融合。注意顺序先上采样再拼接再卷积。这个顺序不能乱因为拼接要求两张特征图的空间尺寸完全一致。我见过有同学写成先拼接再上采样结果就是尺寸对不上直接报错或者强行插值之后拼接特征错位训练能收敛但精度明显差一截。这个细节看起来小其实决定了整个信息通路是否正确。2.2 通道数在拼接点为什么要砍一半这是手敲UNet时最容易被问到的点。假设解码器当前层上采样后的通道数是512编码器同层的通道数也是512拼接之后就是1024。接下来那组双卷积的第一个卷积输入通道要写1024输出通道写回512或者按你自己设定的基数算这样通道数就砍半了。为什么不干脆保持1024一方面是参数量和显存吃不消另一方面是拼接进来的两路特征本来就来自同一个源头的不同抽象层级融合之后没有必要维持双倍宽度。实际实现里有一种更省参数的写法上采样后把通道降到原来的一半这样拼接后正好等于编码器层的通道数卷积层输入输出通道一致代码更好读。两种写法效果差别不大我一般按后者写因为形状推导更直观。下面这张表是输入为1×1×572×572原始论文的输入尺寸无padding卷积时的大致走向你可以对照着看每一层的空间尺寸和通道变化阶段操作输出尺寸 (C×H×W)输入-1×572×572编码1双卷积3×3valid64×568×568下采样12×2最大池化64×284×284编码2双卷积3×3128×280×280下采样22×2最大池化128×140×140编码3双卷积3×3256×136×136下采样32×2最大池化256×68×68编码4双卷积3×3512×64×64下采样42×2最大池化512×32×32瓶颈双卷积3×31024×28×28解码4上采样拼接双卷积512×56×56解码3上采样拼接双卷积256×104×104解码2上采样拼接双卷积128×200×200解码1上采样拼接双卷积64×388×388输出1×1卷积类别数×388×388原始论文因为没有用padding每做一次3×3卷积尺寸就缩2所以输入输出尺寸不一致需要靠镜像padding补齐才能拼回原图。现在大家基本都用padding1的same卷积尺寸不缩整个结构就干净很多。这也是我下面代码采用的方式——输入多大输出就多大省掉了裁剪对齐的麻烦。2.3 跳跃连接到底在补什么信息跳跃连接的作用可以用一句话概括把编码器在浅层提取到的边缘、纹理、角点这类高频信息直接送到解码器对应层级弥补上采样过程中丢失的空间细节。上采样本质上是一种插值无论是最近邻、双线性还是转置卷积都不可能凭空造出被池化丢掉的精确边界。有个很直观的验证方法把跳跃连接全部断开其他不变训练同一个数据集然后对比预测图的边缘。我在一个小目标分割任务上做过这个对照断开跳跃连接后整体IoU只掉了几个点但边界区域的误差明显变大细长结构基本粘连成一坨。所以如果你的任务里目标边界很关键医学影像、遥感、工业缺陷检测都属于这类跳跃连接不是可选项是命脉。另一层理解是梯度通路。深层的梯度回传时跳跃连接提供了一条更短的路径直达浅层缓解了深层网络的梯度衰减。这也是为什么UNet即使堆到比较深训练起来依然相对稳定。3. 手敲之前先算清尺寸账跳跃连接对齐的三种翻车方式3.1 输入尺寸不能被16整除时会发生什么UNet做了4次下采样每次减半所以输入的高和宽必须能被2^416整除否则在某个拼接点上上采样后的尺寸和编码器保存的特征图尺寸会差1个像素torch.cat直接报维度不匹配。这是新手最高频的报错没有之一。解决办法有三种我按推荐程度排第一种在数据预处理阶段把图像padding到16的整数倍推理完成后再裁回原尺寸。这种做法不改变网络最干净代价是要记录padding的参数。第二种把上采样方式改成nn.Upsample(sizeskip.shape[2:])用目标尺寸强制对齐代码上最省事但如果尺寸差得多插值会引入轻微形变。第三种用padding1的卷积配合ceil_modeTrue的池化理论上能缓解但仍然不保证严格对齐我不推荐。我自己的习惯是第一种在Dataset里统一处理。写一个函数算出需要补多少用0填充同时把padding值存下来推理时反向裁剪。这样训练和推理的路径一致不会出现训练时尺寸对、部署时对不上的尴尬。import torch.nn.functional as F def pad_to_multiple(x, m16): _, _, h, w x.shape ph (m - h % m) % m pw (m - w % m) % m # 右、下填充方便推理时直接切片还原 x F.pad(x, (0, pw, 0, ph), modeconstant, value0) return x, (ph, pw)3.2 池化与上采样的对称性问题标准的UNet结构是对称的4次池化对应4次上采样。但如果你为了适配小显存或者小目标改成5次或者3次下采样那么解码器也必须严格对应多一层少一层都会在拼接时出问题。更隐蔽的一种情况是你在编码器里用了stride2的卷积代替池化。这在一些改进版本里很常见好处是下采样过程可学习。但stride2的卷积在奇数尺寸输入时输出是floor((H2p-k)/2)1和池化的floor(H/2)在边界上可能差一位。我的建议是新手阶段老老实实用MaxPool2d(2)结构清晰尺寸行为可预测等结构吃透了再考虑替换。3.3 拼接维度写错cat和add的区别torch.cat(dim1)是通道维度拼接输出通道是两个输入之和或者torch.add是逐元素相加要求通道数相同。UNet用的是前者这一点从拼接后通道翻倍再卷积降维的设计就能看出来。我见过有人把跳跃连接写成相加代码不报错模型也能训但表达能力弱了一截。为什么相加要求两路特征的语义已经对齐而编码器和解码器的特征分布差异其实挺大强行相加相当于假设它们同分布融合效果不如拼接后让网络自己学权重。后来有些改进结构用逐元素相加或者门控机制那是建立在额外的归一化和训练策略之上的不能简单照搬。顺带说一句拼接之后如果立刻接BatchNorm要注意统计量会受两路特征分布差异的影响。我一般保持拼接→卷积→BN→激活的标准顺序不额外插入归一化。4. 从DoubleConv到输出头一份能直接跑起来的PyTorch实现4.1 双卷积块UNet的基本积木UNet里所有卷积都以两层3×3卷积BNReLU的组合出现我把这个组合封装成一个类是整个代码里复用最多的模块。原始论文用的其实是valid卷积且没有BN我在实现里加了padding和BN前者保证尺寸不变后者让训练更稳、对学习率不那么敏感。import torch import torch.nn as nn class DoubleConv(nn.Module): (conv - BN - ReLU) x 2 def __init__(self, in_ch, out_ch, mid_chNone): super().__init__() mid_ch mid_ch or out_ch self.block nn.Sequential( nn.Conv2d(in_ch, mid_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(mid_ch), nn.ReLU(inplaceTrue), nn.Conv2d(mid_ch, out_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x)这里有个小设计mid_ch参数允许中间层通道数和输出不同。标准UNet里两层通道一致但在解码器拼接后输入通道是两倍第一层卷积的输入输出就要单独控制用这个参数写起来更灵活。biasFalse是因为后面紧跟BN偏置项会被BN的减均值操作抵消加了纯属浪费参数。4.2 下采样块与上采样块的分工下采样块就是最大池化加双卷积上采样块则是上采样、拼接、双卷积三步。我把它们各自封装主网络里循环调用代码可读性会好很多。class Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.pool_conv nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.pool_conv(x) class Up(nn.Module): def __init__(self, in_ch, out_ch, bilinearTrue): super().__init__() if bilinear: # 双线性插值不改变通道数随后用1x1卷积降一半 self.up nn.Sequential( nn.Upsample(scale_factor2, modebilinear, align_cornersTrue), nn.Conv2d(in_ch, in_ch // 2, kernel_size1) ) else: # 转置卷积通道数直接减半 self.up nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size2, stride2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x, skip): x self.up(x) # 兜底对齐防止奇数尺寸导致差1个像素 if x.shape[-2:] ! skip.shape[-2:]: x nn.functional.interpolate( x, sizeskip.shape[-2:], modebilinear, align_cornersTrue) x torch.cat([skip, x], dim1) return self.conv(x)in_ch在这里指的是上采样前的通道数in_ch // 2是降维后的通道拼接后正好回到in_ch所以DoubleConv(in_ch, out_ch)的输入输出就对上了。这个推导建议你自己在纸上写一遍写通了整个解码器的形状逻辑就通了。上采样方式的选择我单独说一下。双线性插值加1×1卷积的参数量更少训练更稳不容易出现棋盘格伪影转置卷积是学出来的上采样理论上更灵活但卷积核尺寸和步长的组合如果不当输出会出现明显的网格状条纹。我在中小规模数据集上默认用双线性只有在数据量足够、且对细节恢复要求极高时才试转置卷积。4.3 主网络组装与输出头把上面的积木拼起来就是完整的UNet。通道基数我用base64和原始论文一致显存紧张的话改成32也能跑。class UNet(nn.Module): def __init__(self, in_channels3, num_classes1, base64, bilinearTrue): super().__init__() self.in_channels in_channels self.num_classes num_classes # 编码器 self.inc DoubleConv(in_channels, base) # 64 self.down1 Down(base, base * 2) # 128 self.down2 Down(base * 2, base * 4) # 256 self.down3 Down(base * 4, base * 8) # 512 factor 2 if bilinear else 1 self.down4 Down(base * 8, base * 16 // factor) # 1024 或 512 # 解码器 up_in base * 16 // factor self.up1 Up(up_in, base * 8 // factor, bilinear) self.up2 Up(base * 8 // factor, base * 4 // factor, bilinear) self.up3 Up(base * 4 // factor, base * 2 // factor, bilinear) self.up4 Up(base * 2 // factor, base, bilinear) # 输出头 self.outc nn.Conv2d(base, num_classes, kernel_size1) 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)输出头用1×1卷积把64通道压到类别数。二分类任务num_classes1配合BCEWithLogitsLoss多分类任务设成类别数配合交叉熵。这里注意不要在最后一层加Sigmoid或者Softmax因为PyTorch的这两个损失函数内部已经包含了对应的激活重复做会出问题。用torchsummary或者手动打印一下形状验证if __name__ __main__: model UNet(in_channels1, num_classes1, base64) x torch.randn(2, 1, 256, 256) y model(x) print(y.shape) # torch.Size([2, 1, 256, 256]) print(sum(p.numel() for p in model.parameters()) / 1e6, M)输入256×256输出还是256×256参数量在base64、双线性模式下大约7.8M。这个规模在单张消费级显卡上跑batch8是没问题的。5. 代码跑通只是起点损失函数、数据增强与显存之间的拉扯5.1 二分类分割的损失该怎么选结构搭完跑通一次前向很多人就以为万事大吉了结果训了二十个epoch模型输出全是背景。原因通常是类别极度不均衡前景占比可能只有百分之几交叉熵在这种情况下一开始就被背景主导梯度几乎全指向预测背景。我常用的组合是BCEWithLogitsLoss加一个Dice损失按比例加权。BCE提供稳定的逐像素梯度Dice直接优化重叠度对前景占比敏感两者互补。权重我一般设成0.5 × BCE 0.5 × Dice样本特别不均衡时把Dice提到0.7。class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__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) inter (probs * targets).sum(dim1) union probs.sum(dim1) targets.sum(dim1) dice (2 * inter self.smooth) / (union self.smooth) return 1 - dice.mean()有个细节Dice用的是sigmoid后的概率但梯度仍然能传回logits不需要额外处理。计算时先展平成(N, -1)再逐样本求比全局求平均更稳不会因为某个大样本主导整个loss。5.2 小数据集上数据增强比换模型更有效医学影像、工业缺陷这类场景标注成本高几百张图是常态。这种规模下UNet的参数量其实是过剩的不加正则几乎必然过拟合。我的增强清单通常包括随机旋转任意角度比90度倍数更有效、随机缩放0.8~1.2、水平垂直翻转、弹性变形、亮度对比度扰动、随机裁剪。其中弹性变形对分割任务特别重要它能模拟组织形变、镜头畸变这类真实变化比几何翻转更接近数据分布。代价是计算开销不小我一般把它放在裁剪之后做减小计算量。用albumentations组合起来大概是这个写法import albumentations as A from albumentations.pytorch import ToTensorV2 train_tf A.Compose([ A.RandomResizedCrop(height256, width256, scale(0.7, 1.0), p1.0), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.2), A.Rotate(limit45, border_mode0, p0.5), A.ElasticTransform(alpha30, sigma5, p0.3), A.RandomBrightnessContrast(p0.3), A.Normalize(mean[0.5], std[0.5]), ToTensorV2(), ]) val_tf A.Compose([ A.Resize(256, 256), A.Normalize(mean[0.5], std[0.5]), ToTensorV2(), ])关键点图像和掩码必须用同一套几何变换参数albumentations的Compose传image和mask两个关键字就能自动同步不要自己手写随机数否则两者错位模型学到的就是错的对应关系。5.3 显存不够时按这个顺序妥协UNet在256×256输入、base64的情况下batch8大概占6~8GB显存。显存吃紧的时候我按以下顺序调整从影响最小的开始第一开混合精度训练。torch.cuda.amp基本能省掉将近一半显存速度也有提升对精度影响很小。第二梯度累积。把batch拆成多次前向再统一更新等效于大batch代价是训练时间变长。第三减小base通道数从64降到32参数量降到四分之一。第四缩小输入尺寸但这会影响小目标的可见性放最后考虑。我不太推荐一上来就用梯度检查点虽然省显存但训练速度会明显下降而且实现上容易出错。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for imgs, masks in loader: optimizer.zero_grad() with autocast(): logits model(imgs) loss criterion(logits, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()用AMP的时候注意Dice损失内部有sigmoid和除法建议在float32下算避免数值不稳定可以在autocast上下文外面单独计算损失。6. 训练曲线不好看时我通常按这个顺序排查6.1 先确认模型是不是什么都没学到判断方法很简单拿一张训练集图片喂给模型看输出是不是接近全0或者全1。如果两者之间几乎没有变化说明梯度没传对问题往往出在损失函数和激活函数的搭配上——比如输出层加了Sigmoid又用了BCEWithLogitsLoss等于做了两次压缩梯度被严重削弱。第二个常见原因是标签的数值范围。掩码必须是0/1或者0~C-1的整数如果读进来的掩码被归一化成了0~1的浮点数BCE会报错或者给出奇怪的lossDice也会算错。我踩过一次图像做了归一化顺手把掩码也一起处理了排查了大半天。第三个是学习率。UNet配上Adam初始学习率我用1e-3或者3e-4太大容易在早期就震荡到坏解太小则前几个epoch看起来毫无进展。可以先用一个小数据集过拟合测试取10张图关掉增强训练到loss接近0。如果这一步都做不到就别急着上全量数据。6.2 训练loss降、验证loss涨问题在正则而不是结构这是典型的过拟合信号。先加增强再加weight decay1e-4量级还不行就上Dropout。UNet里Dropout一般加在瓶颈层和每次下采样之后p0.1~0.3不要加在输出层附近。另一种情况是验证集指标卡在一个值不动。这时候要检查验证时的预处理是否和训练一致特别是归一化的均值和方差、resize的方式。训练用随机裁剪验证用直接resize尺度分布不一致指标自然上不去。下面这张表是我遇到过的几类典型症状和对应处理供对照症状可能原因处理方式输出全为背景类别不均衡、损失被背景主导加Dice、提高前景采样比例边界糊成一团跳跃连接未生效或特征错位检查cat的维度与尺寸对齐loss变成NaN学习率过大、AMP下除零降学习率、Dice加smooth项验证指标停滞训练/验证预处理不一致统一归一化与resize参数训练极慢未开AMP、DataLoader单进程开混合精度、调num_workers6.3 评估指标别只看像素准确率分割任务里像素准确率是个很有欺骗性的指标。前景只占5%的时候模型全预测背景也能拿到95%的准确率看起来很美。真正有意义的是IoU和Dice它们只在前景和预测的交集上做计算能真实反映分割质量。计算IoU时记得加一个极小的平滑项避免分母为0。另外二分类的分割结果通常需要先按阈值二值化阈值取0.5是个默认值但如果你的任务更看重召回比如病灶筛查可以适当调低到0.3~0.4把更多可疑区域圈进来后续再人工筛。7. 几个只有真正敲过一遍才会记住的细节biasFalse配BN这个组合我在第一次手敲时完全没在意后来发现参数量比对不上才回头逐个层核对。BN的减均值操作会把卷积的常数偏置吃掉所以那个参数是冗余的删掉不影响结果还能少一点显存。上采样用align_cornersTrue还是False这个参数在双线性插值时会影响边界像素的对齐方式。做分割时我倾向True因为输入输出尺寸一致边界像素能严格对应如果做的是尺寸变化较大的任务False可能更合适。两种都能训但训练和推理要保持一致中途改会导致结果漂移。还有一个容易忽略的点推理时不要忘了model.eval()。BN在训练模式下用的是当前batch的统计量验证时如果忘了切换同一个样本在不同batch里的输出会不一样指标波动大得莫名其妙。我第一次遇到时还以为是数据有问题查了很久。输入通道数也要和实际数据对上。灰度图是1通道RGB是3通道如果数据集读出来是4通道比如带alpha记得在预处理里裁掉否则第一层卷积直接报维度错误。这种错误信息很直白但如果不熟悉报错格式也容易卡住。最后说一个结构上的扩展思路如果你要处理的是三维数据比如体素或者多帧序列把所有的Conv2d、BatchNorm2d、MaxPool2d换成对应的3d版本上采样换成Upsample(modetrilinear)整个结构可以原样搬过去。代价是显存消耗呈立方增长通常需要把base降到16或32并配合patch切块训练。我在一个三维任务上就是这么改的逻辑一行没动只换了层类型和输入尺寸跑通得很快。 
长河君 · 2026/10/27 3:23:38 · 阅读 1.2k 手敲UNet图像分割：结构拆解与PyTorch训练调优
1. 一张512×512的切片为什么普通卷积网络啃不动第一次把UNet的代码从头敲完是在一个图像分割的小任务上。当时手上只有几百张标注好的切片背景面积占比超过九成前景目标边缘还特别碎。我一开始图省事拿了个常规的分类骨干网络加两层卷积直接出结果训练loss降得挺漂亮但把预测图叠回原图一看目标边界的轮廓糊成一团小目标要么整块消失要么被连成大块。那一刻我才真正理解UNet这套结构为什么要长成U形——它解决的不是能不能分出目标的问题而是分出来的边界够不够准的问题。UNet最早是为生物医学图像分割提出的全卷积网络核心构架由编码器下采样路径、解码器上采样路径和跳跃连接三部分组成。它解决的核心矛盾很具体深层特征语义强但分辨率低浅层特征分辨率高但语义弱。分类网络只关心最后那个向量所以一路把特征图压小没问题分割网络要求输出和输入同尺寸的逐像素标签那就必须把分辨率还原回去。而单纯靠上采样还原丢掉的空间细节是找不回来的UNet的解法是在还原过程中把编码器同层的特征图直接搬过来拼接。这篇内容适合两类人看一类是已经会用现成库调模型、但想搞清楚内部每一层到底在干什么的开发者另一类是准备面试、需要白板手推结构或者现场手敲代码的同学。我会把网络结构、张量形状的推导、手敲代码的完整实现、训练侧的几个关键取舍一步一步拆开讲。代码用PyTorch写尽量做到复制粘贴就能跑同时把每个设计选择背后的理由说清楚而不是只给你一堆层。在动手之前先明确一件事UNet的手敲不是背代码而是能在纸上画出数据流能说清楚任意一层输入输出的通道数和空间尺寸能解释为什么拼接之后通道数要减半。这三点想透了代码自然就写出来了而且换个输入尺寸、换个通道基数也不会翻车。2. 把UNet拆成编码、瓶颈、解码三段通道与分辨率怎么走2.1 编码器负责看懂是什么解码器负责说清在哪里编码器的任务和普通卷积分类网络完全一致通过反复的卷积加下采样把感受野撑大让网络看到更大范围的上下文。UNet原始论文里用了4次下采样每次分辨率减半、通道数翻倍。第一次下采样后特征图尺寸是输入的1/2通道从1或3变成64第二次是1/4、128第三次1/8、256第四次1/16、512。到最底部瓶颈层再做一次卷积通道到1024分辨率保持1/16。解码器则是反过来每一步先做上采样把分辨率放大一倍然后把编码器对应层的特征图拿过来在通道维度上拼接再接两层卷积做特征融合。注意顺序先上采样再拼接再卷积。这个顺序不能乱因为拼接要求两张特征图的空间尺寸完全一致。我见过有同学写成先拼接再上采样结果就是尺寸对不上直接报错或者强行插值之后拼接特征错位训练能收敛但精度明显差一截。这个细节看起来小其实决定了整个信息通路是否正确。2.2 通道数在拼接点为什么要砍一半这是手敲UNet时最容易被问到的点。假设解码器当前层上采样后的通道数是512编码器同层的通道数也是512拼接之后就是1024。接下来那组双卷积的第一个卷积输入通道要写1024输出通道写回512或者按你自己设定的基数算这样通道数就砍半了。为什么不干脆保持1024一方面是参数量和显存吃不消另一方面是拼接进来的两路特征本来就来自同一个源头的不同抽象层级融合之后没有必要维持双倍宽度。实际实现里有一种更省参数的写法上采样后把通道降到原来的一半这样拼接后正好等于编码器层的通道数卷积层输入输出通道一致代码更好读。两种写法效果差别不大我一般按后者写因为形状推导更直观。下面这张表是输入为1×1×572×572原始论文的输入尺寸无padding卷积时的大致走向你可以对照着看每一层的空间尺寸和通道变化阶段操作输出尺寸 (C×H×W)输入-1×572×572编码1双卷积3×3valid64×568×568下采样12×2最大池化64×284×284编码2双卷积3×3128×280×280下采样22×2最大池化128×140×140编码3双卷积3×3256×136×136下采样32×2最大池化256×68×68编码4双卷积3×3512×64×64下采样42×2最大池化512×32×32瓶颈双卷积3×31024×28×28解码4上采样拼接双卷积512×56×56解码3上采样拼接双卷积256×104×104解码2上采样拼接双卷积128×200×200解码1上采样拼接双卷积64×388×388输出1×1卷积类别数×388×388原始论文因为没有用padding每做一次3×3卷积尺寸就缩2所以输入输出尺寸不一致需要靠镜像padding补齐才能拼回原图。现在大家基本都用padding1的same卷积尺寸不缩整个结构就干净很多。这也是我下面代码采用的方式——输入多大输出就多大省掉了裁剪对齐的麻烦。2.3 跳跃连接到底在补什么信息跳跃连接的作用可以用一句话概括把编码器在浅层提取到的边缘、纹理、角点这类高频信息直接送到解码器对应层级弥补上采样过程中丢失的空间细节。上采样本质上是一种插值无论是最近邻、双线性还是转置卷积都不可能凭空造出被池化丢掉的精确边界。有个很直观的验证方法把跳跃连接全部断开其他不变训练同一个数据集然后对比预测图的边缘。我在一个小目标分割任务上做过这个对照断开跳跃连接后整体IoU只掉了几个点但边界区域的误差明显变大细长结构基本粘连成一坨。所以如果你的任务里目标边界很关键医学影像、遥感、工业缺陷检测都属于这类跳跃连接不是可选项是命脉。另一层理解是梯度通路。深层的梯度回传时跳跃连接提供了一条更短的路径直达浅层缓解了深层网络的梯度衰减。这也是为什么UNet即使堆到比较深训练起来依然相对稳定。3. 手敲之前先算清尺寸账跳跃连接对齐的三种翻车方式3.1 输入尺寸不能被16整除时会发生什么UNet做了4次下采样每次减半所以输入的高和宽必须能被2^416整除否则在某个拼接点上上采样后的尺寸和编码器保存的特征图尺寸会差1个像素torch.cat直接报维度不匹配。这是新手最高频的报错没有之一。解决办法有三种我按推荐程度排第一种在数据预处理阶段把图像padding到16的整数倍推理完成后再裁回原尺寸。这种做法不改变网络最干净代价是要记录padding的参数。第二种把上采样方式改成nn.Upsample(sizeskip.shape[2:])用目标尺寸强制对齐代码上最省事但如果尺寸差得多插值会引入轻微形变。第三种用padding1的卷积配合ceil_modeTrue的池化理论上能缓解但仍然不保证严格对齐我不推荐。我自己的习惯是第一种在Dataset里统一处理。写一个函数算出需要补多少用0填充同时把padding值存下来推理时反向裁剪。这样训练和推理的路径一致不会出现训练时尺寸对、部署时对不上的尴尬。import torch.nn.functional as F def pad_to_multiple(x, m16): _, _, h, w x.shape ph (m - h % m) % m pw (m - w % m) % m # 右、下填充方便推理时直接切片还原 x F.pad(x, (0, pw, 0, ph), modeconstant, value0) return x, (ph, pw)3.2 池化与上采样的对称性问题标准的UNet结构是对称的4次池化对应4次上采样。但如果你为了适配小显存或者小目标改成5次或者3次下采样那么解码器也必须严格对应多一层少一层都会在拼接时出问题。更隐蔽的一种情况是你在编码器里用了stride2的卷积代替池化。这在一些改进版本里很常见好处是下采样过程可学习。但stride2的卷积在奇数尺寸输入时输出是floor((H2p-k)/2)1和池化的floor(H/2)在边界上可能差一位。我的建议是新手阶段老老实实用MaxPool2d(2)结构清晰尺寸行为可预测等结构吃透了再考虑替换。3.3 拼接维度写错cat和add的区别torch.cat(dim1)是通道维度拼接输出通道是两个输入之和或者torch.add是逐元素相加要求通道数相同。UNet用的是前者这一点从拼接后通道翻倍再卷积降维的设计就能看出来。我见过有人把跳跃连接写成相加代码不报错模型也能训但表达能力弱了一截。为什么相加要求两路特征的语义已经对齐而编码器和解码器的特征分布差异其实挺大强行相加相当于假设它们同分布融合效果不如拼接后让网络自己学权重。后来有些改进结构用逐元素相加或者门控机制那是建立在额外的归一化和训练策略之上的不能简单照搬。顺带说一句拼接之后如果立刻接BatchNorm要注意统计量会受两路特征分布差异的影响。我一般保持拼接→卷积→BN→激活的标准顺序不额外插入归一化。4. 从DoubleConv到输出头一份能直接跑起来的PyTorch实现4.1 双卷积块UNet的基本积木UNet里所有卷积都以两层3×3卷积BNReLU的组合出现我把这个组合封装成一个类是整个代码里复用最多的模块。原始论文用的其实是valid卷积且没有BN我在实现里加了padding和BN前者保证尺寸不变后者让训练更稳、对学习率不那么敏感。import torch import torch.nn as nn class DoubleConv(nn.Module): (conv - BN - ReLU) x 2 def __init__(self, in_ch, out_ch, mid_chNone): super().__init__() mid_ch mid_ch or out_ch self.block nn.Sequential( nn.Conv2d(in_ch, mid_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(mid_ch), nn.ReLU(inplaceTrue), nn.Conv2d(mid_ch, out_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x)这里有个小设计mid_ch参数允许中间层通道数和输出不同。标准UNet里两层通道一致但在解码器拼接后输入通道是两倍第一层卷积的输入输出就要单独控制用这个参数写起来更灵活。biasFalse是因为后面紧跟BN偏置项会被BN的减均值操作抵消加了纯属浪费参数。4.2 下采样块与上采样块的分工下采样块就是最大池化加双卷积上采样块则是上采样、拼接、双卷积三步。我把它们各自封装主网络里循环调用代码可读性会好很多。class Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.pool_conv nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.pool_conv(x) class Up(nn.Module): def __init__(self, in_ch, out_ch, bilinearTrue): super().__init__() if bilinear: # 双线性插值不改变通道数随后用1x1卷积降一半 self.up nn.Sequential( nn.Upsample(scale_factor2, modebilinear, align_cornersTrue), nn.Conv2d(in_ch, in_ch // 2, kernel_size1) ) else: # 转置卷积通道数直接减半 self.up nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size2, stride2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x, skip): x self.up(x) # 兜底对齐防止奇数尺寸导致差1个像素 if x.shape[-2:] ! skip.shape[-2:]: x nn.functional.interpolate( x, sizeskip.shape[-2:], modebilinear, align_cornersTrue) x torch.cat([skip, x], dim1) return self.conv(x)in_ch在这里指的是上采样前的通道数in_ch // 2是降维后的通道拼接后正好回到in_ch所以DoubleConv(in_ch, out_ch)的输入输出就对上了。这个推导建议你自己在纸上写一遍写通了整个解码器的形状逻辑就通了。上采样方式的选择我单独说一下。双线性插值加1×1卷积的参数量更少训练更稳不容易出现棋盘格伪影转置卷积是学出来的上采样理论上更灵活但卷积核尺寸和步长的组合如果不当输出会出现明显的网格状条纹。我在中小规模数据集上默认用双线性只有在数据量足够、且对细节恢复要求极高时才试转置卷积。4.3 主网络组装与输出头把上面的积木拼起来就是完整的UNet。通道基数我用base64和原始论文一致显存紧张的话改成32也能跑。class UNet(nn.Module): def __init__(self, in_channels3, num_classes1, base64, bilinearTrue): super().__init__() self.in_channels in_channels self.num_classes num_classes # 编码器 self.inc DoubleConv(in_channels, base) # 64 self.down1 Down(base, base * 2) # 128 self.down2 Down(base * 2, base * 4) # 256 self.down3 Down(base * 4, base * 8) # 512 factor 2 if bilinear else 1 self.down4 Down(base * 8, base * 16 // factor) # 1024 或 512 # 解码器 up_in base * 16 // factor self.up1 Up(up_in, base * 8 // factor, bilinear) self.up2 Up(base * 8 // factor, base * 4 // factor, bilinear) self.up3 Up(base * 4 // factor, base * 2 // factor, bilinear) self.up4 Up(base * 2 // factor, base, bilinear) # 输出头 self.outc nn.Conv2d(base, num_classes, kernel_size1) 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)输出头用1×1卷积把64通道压到类别数。二分类任务num_classes1配合BCEWithLogitsLoss多分类任务设成类别数配合交叉熵。这里注意不要在最后一层加Sigmoid或者Softmax因为PyTorch的这两个损失函数内部已经包含了对应的激活重复做会出问题。用torchsummary或者手动打印一下形状验证if __name__ __main__: model UNet(in_channels1, num_classes1, base64) x torch.randn(2, 1, 256, 256) y model(x) print(y.shape) # torch.Size([2, 1, 256, 256]) print(sum(p.numel() for p in model.parameters()) / 1e6, M)输入256×256输出还是256×256参数量在base64、双线性模式下大约7.8M。这个规模在单张消费级显卡上跑batch8是没问题的。5. 代码跑通只是起点损失函数、数据增强与显存之间的拉扯5.1 二分类分割的损失该怎么选结构搭完跑通一次前向很多人就以为万事大吉了结果训了二十个epoch模型输出全是背景。原因通常是类别极度不均衡前景占比可能只有百分之几交叉熵在这种情况下一开始就被背景主导梯度几乎全指向预测背景。我常用的组合是BCEWithLogitsLoss加一个Dice损失按比例加权。BCE提供稳定的逐像素梯度Dice直接优化重叠度对前景占比敏感两者互补。权重我一般设成0.5 × BCE 0.5 × Dice样本特别不均衡时把Dice提到0.7。class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__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) inter (probs * targets).sum(dim1) union probs.sum(dim1) targets.sum(dim1) dice (2 * inter self.smooth) / (union self.smooth) return 1 - dice.mean()有个细节Dice用的是sigmoid后的概率但梯度仍然能传回logits不需要额外处理。计算时先展平成(N, -1)再逐样本求比全局求平均更稳不会因为某个大样本主导整个loss。5.2 小数据集上数据增强比换模型更有效医学影像、工业缺陷这类场景标注成本高几百张图是常态。这种规模下UNet的参数量其实是过剩的不加正则几乎必然过拟合。我的增强清单通常包括随机旋转任意角度比90度倍数更有效、随机缩放0.8~1.2、水平垂直翻转、弹性变形、亮度对比度扰动、随机裁剪。其中弹性变形对分割任务特别重要它能模拟组织形变、镜头畸变这类真实变化比几何翻转更接近数据分布。代价是计算开销不小我一般把它放在裁剪之后做减小计算量。用albumentations组合起来大概是这个写法import albumentations as A from albumentations.pytorch import ToTensorV2 train_tf A.Compose([ A.RandomResizedCrop(height256, width256, scale(0.7, 1.0), p1.0), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.2), A.Rotate(limit45, border_mode0, p0.5), A.ElasticTransform(alpha30, sigma5, p0.3), A.RandomBrightnessContrast(p0.3), A.Normalize(mean[0.5], std[0.5]), ToTensorV2(), ]) val_tf A.Compose([ A.Resize(256, 256), A.Normalize(mean[0.5], std[0.5]), ToTensorV2(), ])关键点图像和掩码必须用同一套几何变换参数albumentations的Compose传image和mask两个关键字就能自动同步不要自己手写随机数否则两者错位模型学到的就是错的对应关系。5.3 显存不够时按这个顺序妥协UNet在256×256输入、base64的情况下batch8大概占6~8GB显存。显存吃紧的时候我按以下顺序调整从影响最小的开始第一开混合精度训练。torch.cuda.amp基本能省掉将近一半显存速度也有提升对精度影响很小。第二梯度累积。把batch拆成多次前向再统一更新等效于大batch代价是训练时间变长。第三减小base通道数从64降到32参数量降到四分之一。第四缩小输入尺寸但这会影响小目标的可见性放最后考虑。我不太推荐一上来就用梯度检查点虽然省显存但训练速度会明显下降而且实现上容易出错。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for imgs, masks in loader: optimizer.zero_grad() with autocast(): logits model(imgs) loss criterion(logits, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()用AMP的时候注意Dice损失内部有sigmoid和除法建议在float32下算避免数值不稳定可以在autocast上下文外面单独计算损失。6. 训练曲线不好看时我通常按这个顺序排查6.1 先确认模型是不是什么都没学到判断方法很简单拿一张训练集图片喂给模型看输出是不是接近全0或者全1。如果两者之间几乎没有变化说明梯度没传对问题往往出在损失函数和激活函数的搭配上——比如输出层加了Sigmoid又用了BCEWithLogitsLoss等于做了两次压缩梯度被严重削弱。第二个常见原因是标签的数值范围。掩码必须是0/1或者0~C-1的整数如果读进来的掩码被归一化成了0~1的浮点数BCE会报错或者给出奇怪的lossDice也会算错。我踩过一次图像做了归一化顺手把掩码也一起处理了排查了大半天。第三个是学习率。UNet配上Adam初始学习率我用1e-3或者3e-4太大容易在早期就震荡到坏解太小则前几个epoch看起来毫无进展。可以先用一个小数据集过拟合测试取10张图关掉增强训练到loss接近0。如果这一步都做不到就别急着上全量数据。6.2 训练loss降、验证loss涨问题在正则而不是结构这是典型的过拟合信号。先加增强再加weight decay1e-4量级还不行就上Dropout。UNet里Dropout一般加在瓶颈层和每次下采样之后p0.1~0.3不要加在输出层附近。另一种情况是验证集指标卡在一个值不动。这时候要检查验证时的预处理是否和训练一致特别是归一化的均值和方差、resize的方式。训练用随机裁剪验证用直接resize尺度分布不一致指标自然上不去。下面这张表是我遇到过的几类典型症状和对应处理供对照症状可能原因处理方式输出全为背景类别不均衡、损失被背景主导加Dice、提高前景采样比例边界糊成一团跳跃连接未生效或特征错位检查cat的维度与尺寸对齐loss变成NaN学习率过大、AMP下除零降学习率、Dice加smooth项验证指标停滞训练/验证预处理不一致统一归一化与resize参数训练极慢未开AMP、DataLoader单进程开混合精度、调num_workers6.3 评估指标别只看像素准确率分割任务里像素准确率是个很有欺骗性的指标。前景只占5%的时候模型全预测背景也能拿到95%的准确率看起来很美。真正有意义的是IoU和Dice它们只在前景和预测的交集上做计算能真实反映分割质量。计算IoU时记得加一个极小的平滑项避免分母为0。另外二分类的分割结果通常需要先按阈值二值化阈值取0.5是个默认值但如果你的任务更看重召回比如病灶筛查可以适当调低到0.3~0.4把更多可疑区域圈进来后续再人工筛。7. 几个只有真正敲过一遍才会记住的细节biasFalse配BN这个组合我在第一次手敲时完全没在意后来发现参数量比对不上才回头逐个层核对。BN的减均值操作会把卷积的常数偏置吃掉所以那个参数是冗余的删掉不影响结果还能少一点显存。上采样用align_cornersTrue还是False这个参数在双线性插值时会影响边界像素的对齐方式。做分割时我倾向True因为输入输出尺寸一致边界像素能严格对应如果做的是尺寸变化较大的任务False可能更合适。两种都能训但训练和推理要保持一致中途改会导致结果漂移。还有一个容易忽略的点推理时不要忘了model.eval()。BN在训练模式下用的是当前batch的统计量验证时如果忘了切换同一个样本在不同batch里的输出会不一样指标波动大得莫名其妙。我第一次遇到时还以为是数据有问题查了很久。输入通道数也要和实际数据对上。灰度图是1通道RGB是3通道如果数据集读出来是4通道比如带alpha记得在预处理里裁掉否则第一层卷积直接报维度错误。这种错误信息很直白但如果不熟悉报错格式也容易卡住。最后说一个结构上的扩展思路如果你要处理的是三维数据比如体素或者多帧序列把所有的Conv2d、BatchNorm2d、MaxPool2d换成对应的3d版本上采样换成Upsample(modetrilinear)整个结构可以原样搬过去。代价是显存消耗呈立方增长通常需要把base降到16或32并配合patch切块训练。我在一个三维任务上就是这么改的逻辑一行没动只换了层类型和输入尺寸跑通得很快。