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卷积）时的大致走向，你可以对照着看每一层的空间尺寸和通道变化：

原始论文因为没有用padding，每做一次3×3卷积尺寸就缩2，所以输入输出尺寸不一致，需要靠镜像padding补齐才能拼回原图。现在大家基本都用padding=1的same卷积，尺寸不缩，整个结构就干净很多。这也是我下面代码采用的方式——输入多大，输出就多大，省掉了裁剪对齐的麻烦。

2.3 跳跃连接到底在补什么信息

跳跃连接的作用可以用一句话概括：把编码器在浅层提取到的边缘、纹理、角点这类高频信息，直接送到解码器对应层级，弥补上采样过程中丢失的空间细节。上采样本质上是一种插值，无论是最近邻、双线性还是转置卷积，都不可能凭空造出被池化丢掉的精确边界。

有个很直观的验证方法：把跳跃连接全部断开，其他不变，训练同一个数据集，然后对比预测图的边缘。我在一个小目标分割任务上做过这个对照，断开跳跃连接后，整体IoU只掉了几个点，但边界区域的误差明显变大，细长结构基本粘连成一坨。所以如果你的任务里目标边界很关键（医学影像、遥感、工业缺陷检测都属于这类），跳跃连接不是可选项，是命脉。

另一层理解是梯度通路。深层的梯度回传时，跳跃连接提供了一条更短的路径直达浅层，缓解了深层网络的梯度衰减。这也是为什么UNet即使堆到比较深，训练起来依然相对稳定。

3. 手敲之前先算清尺寸账：跳跃连接对齐的三种翻车方式

3.1 输入尺寸不能被16整除时会发生什么

UNet做了4次下采样，每次减半，所以输入的高和宽必须能被2^4=16整除，否则在某个拼接点上，上采样后的尺寸和编码器保存的特征图尺寸会差1个像素，torch.cat直接报维度不匹配。

这是新手最高频的报错，没有之一。解决办法有三种，我按推荐程度排：

第一种，在数据预处理阶段把图像padding到16的整数倍，推理完成后再裁回原尺寸。这种做法不改变网络，最干净，代价是要记录padding的参数。第二种，把上采样方式改成nn.Upsample(size=skip.shape[2:])，用目标尺寸强制对齐，代码上最省事，但如果尺寸差得多，插值会引入轻微形变。第三种，用padding=1的卷积配合ceil_mode=True的池化，理论上能缓解，但仍然不保证严格对齐，我不推荐。

我自己的习惯是第一种，在Dataset里统一处理。写一个函数，算出需要补多少，用0填充，同时把padding值存下来，推理时反向裁剪。这样训练和推理的路径一致，不会出现"训练时尺寸对、部署时对不上"的尴尬。

import torch.nn.functional as F def pad_to_multiple(x, m=16): _, _, h, w = x.shape ph = (m - h % m) % m pw = (m - w % m) % m # 右、下填充，方便推理时直接切片还原 x = F.pad(x, (0, pw, 0, ph), mode='constant', value=0) return x, (ph, pw)

3.2 池化与上采样的对称性问题

标准的UNet结构是对称的：4次池化对应4次上采样。但如果你为了适配小显存或者小目标，改成5次或者3次下采样，那么解码器也必须严格对应，多一层少一层都会在拼接时出问题。

更隐蔽的一种情况是：你在编码器里用了stride=2的卷积代替池化。这在一些改进版本里很常见，好处是下采样过程可学习。但stride=2的卷积在奇数尺寸输入时，输出是floor((H+2p-k)/2)+1，和池化的floor(H/2)在边界上可能差一位。我的建议是：新手阶段老老实实用MaxPool2d(2)，结构清晰，尺寸行为可预测，等结构吃透了再考虑替换。

3.3 拼接维度写错：cat和add的区别

torch.cat(dim=1)是通道维度拼接，输出通道是两个输入之和；+或者torch.add是逐元素相加，要求通道数相同。UNet用的是前者，这一点从"拼接后通道翻倍再卷积降维"的设计就能看出来。

我见过有人把跳跃连接写成相加，代码不报错，模型也能训，但表达能力弱了一截。为什么？相加要求两路特征的语义已经对齐，而编码器和解码器的特征分布差异其实挺大，强行相加相当于假设它们同分布，融合效果不如拼接后让网络自己学权重。后来有些改进结构用逐元素相加或者门控机制，那是建立在额外的归一化和训练策略之上的，不能简单照搬。

顺带说一句，拼接之后如果立刻接BatchNorm，要注意统计量会受两路特征分布差异的影响。我一般保持"拼接→卷积→BN→激活"的标准顺序，不额外插入归一化。

4. 从DoubleConv到输出头：一份能直接跑起来的PyTorch实现

4.1 双卷积块：UNet的基本积木

UNet里所有卷积都以"两层3×3卷积＋BN＋ReLU"的组合出现，我把这个组合封装成一个类，是整个代码里复用最多的模块。原始论文用的其实是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_ch=None): super().__init__() mid_ch = mid_ch or out_ch self.block = nn.Sequential( nn.Conv2d(in_ch, mid_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(mid_ch), nn.ReLU(inplace=True), nn.Conv2d(mid_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)

这里有个小设计：mid_ch参数允许中间层通道数和输出不同。标准UNet里两层通道一致，但在解码器拼接后，输入通道是两倍，第一层卷积的输入输出就要单独控制，用这个参数写起来更灵活。bias=False是因为后面紧跟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, bilinear=True): super().__init__() if bilinear: # 双线性插值不改变通道数，随后用1x1卷积降一半 self.up = nn.Sequential( nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True), nn.Conv2d(in_ch, in_ch // 2, kernel_size=1) ) else: # 转置卷积：通道数直接减半 self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2) 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, size=skip.shape[-2:], mode='bilinear', align_corners=True) x = torch.cat([skip, x], dim=1) return self.conv(x)

in_ch在这里指的是上采样前的通道数，in_ch // 2是降维后的通道，拼接后正好回到in_ch，所以DoubleConv(in_ch, out_ch)的输入输出就对上了。这个推导建议你自己在纸上写一遍，写通了整个解码器的形状逻辑就通了。

上采样方式的选择我单独说一下。双线性插值加1×1卷积的参数量更少，训练更稳，不容易出现棋盘格伪影；转置卷积是"学出来的上采样"，理论上更灵活，但卷积核尺寸和步长的组合如果不当，输出会出现明显的网格状条纹。我在中小规模数据集上默认用双线性，只有在数据量足够、且对细节恢复要求极高时才试转置卷积。

4.3 主网络组装与输出头

把上面的积木拼起来，就是完整的UNet。通道基数我用base=64，和原始论文一致，显存紧张的话改成32也能跑。

class UNet(nn.Module): def __init__(self, in_channels=3, num_classes=1, base=64, bilinear=True): 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_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)

输出头用1×1卷积，把64通道压到类别数。二分类任务num_classes=1，配合BCEWithLogitsLoss；多分类任务设成类别数，配合交叉熵。这里注意：不要在最后一层加Sigmoid或者Softmax，因为PyTorch的这两个损失函数内部已经包含了对应的激活，重复做会出问题。

用torchsummary或者手动打印一下形状验证：

if __name__ == '__main__': model = UNet(in_channels=1, num_classes=1, base=64) 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，参数量在base=64、双线性模式下大约7.8M。这个规模在单张消费级显卡上跑batch=8是没问题的。

5. 代码跑通只是起点：损失函数、数据增强与显存之间的拉扯

5.1 二分类分割的损失该怎么选

结构搭完，跑通一次前向，很多人就以为万事大吉了，结果训了二十个epoch，模型输出全是背景。原因通常是类别极度不均衡，前景占比可能只有百分之几，交叉熵在这种情况下一开始就被背景主导，梯度几乎全指向"预测背景"。

我常用的组合是BCEWithLogitsLoss加一个Dice损失，按比例加权。BCE提供稳定的逐像素梯度，Dice直接优化重叠度，对前景占比敏感，两者互补。权重我一般设成0.5 × BCE + 0.5 × Dice，样本特别不均衡时把Dice提到0.7。

class DiceLoss(nn.Module): def __init__(self, smooth=1.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(dim=1) union = probs.sum(dim=1) + targets.sum(dim=1) 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(height=256, width=256, scale=(0.7, 1.0), p=1.0), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.2), A.Rotate(limit=45, border_mode=0, p=0.5), A.ElasticTransform(alpha=30, sigma=5, p=0.3), A.RandomBrightnessContrast(p=0.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输入、base=64的情况下，batch=8大概占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会报错或者给出奇怪的loss，Dice也会算错。我踩过一次，图像做了归一化，顺手把掩码也一起处理了，排查了大半天。

第三个是学习率。UNet配上Adam，初始学习率我用1e-3或者3e-4，太大容易在早期就震荡到坏解，太小则前几个epoch看起来毫无进展。可以先用一个小数据集过拟合测试：取10张图，关掉增强，训练到loss接近0。如果这一步都做不到，就别急着上全量数据。

6.2 训练loss降、验证loss涨，问题在正则而不是结构

这是典型的过拟合信号。先加增强，再加weight decay（1e-4量级），还不行就上Dropout。UNet里Dropout一般加在瓶颈层和每次下采样之后，p=0.1~0.3，不要加在输出层附近。

另一种情况是验证集指标卡在一个值不动。这时候要检查验证时的预处理是否和训练一致，特别是归一化的均值和方差、resize的方式。训练用随机裁剪，验证用直接resize，尺度分布不一致，指标自然上不去。

下面这张表是我遇到过的几类典型症状和对应处理，供对照：

6.3 评估指标别只看像素准确率

分割任务里像素准确率是个很有欺骗性的指标。前景只占5%的时候，模型全预测背景也能拿到95%的准确率，看起来很美。真正有意义的是IoU和Dice，它们只在前景和预测的交集上做计算，能真实反映分割质量。

计算IoU时记得加一个极小的平滑项，避免分母为0。另外，二分类的分割结果通常需要先按阈值二值化，阈值取0.5是个默认值，但如果你的任务更看重召回（比如病灶筛查），可以适当调低到0.3~0.4，把更多可疑区域圈进来，后续再人工筛。

7. 几个只有真正敲过一遍才会记住的细节

bias=False配BN这个组合，我在第一次手敲时完全没在意，后来发现参数量比对不上，才回头逐个层核对。BN的减均值操作会把卷积的常数偏置吃掉，所以那个参数是冗余的，删掉不影响结果，还能少一点显存。

上采样用align_corners=True还是False，这个参数在双线性插值时会影响边界像素的对齐方式。做分割时我倾向True，因为输入输出尺寸一致，边界像素能严格对应；如果做的是尺寸变化较大的任务，False可能更合适。两种都能训，但训练和推理要保持一致，中途改会导致结果漂移。

还有一个容易忽略的点：推理时不要忘了model.eval()。BN在训练模式下用的是当前batch的统计量，验证时如果忘了切换，同一个样本在不同batch里的输出会不一样，指标波动大得莫名其妙。我第一次遇到时还以为是数据有问题，查了很久。

输入通道数也要和实际数据对上。灰度图是1通道，RGB是3通道，如果数据集读出来是4通道（比如带alpha），记得在预处理里裁掉，否则第一层卷积直接报维度错误。这种错误信息很直白，但如果不熟悉报错格式，也容易卡住。

最后说一个结构上的扩展思路：如果你要处理的是三维数据（比如体素或者多帧序列），把所有的Conv2d、BatchNorm2d、MaxPool2d换成对应的3d版本，上采样换成Upsample(mode='trilinear')，整个结构可以原样搬过去。代价是显存消耗呈立方增长，通常需要把base降到16或32，并配合patch切块训练。我在一个三维任务上就是这么改的，逻辑一行没动，只换了层类型和输入尺寸，跑通得很快。