这几年和不少同学、同行聊U-Net大家普遍都经历过同一个阶段代码从 GitHub 上拷下来能跑验证集的 Dice 分数还挺好看但不会被提问“编码器里为什么每层都套两个卷积”“跳跃连接真的拿掉会怎样”“一张 512×512 的图过完最底层分辨率变成多少”十有八九答不透。这篇文章就是想帮你把这件事彻底补上用 PyTorch 从零手写一遍 U-Net边写边算把每个卷积、每次池化、每次上采样和每条跳跃连接背后的尺寸变化、通道变化、参数来源全部摊开讲清楚。适不适合你两拨人最对口一是准备认真学语义分割但还没动手写过完整网络的新手二是已经能跑通现成代码、但想真正读懂模型内部为什么这样设计的进阶选手。读完你可以只靠torch.nn的常用模块把 U-Net 复现出来不用抄任何第三方封装库同时能根据实际报错自己定位形状不匹配、损失不降、显存爆炸这三类最常见问题。从题目看“手算 U-Net”这几个字重点不在“写”而在“算”。真正把每一步的输入输出尺寸推一遍你对整个网络的理解程度会完全不一样。下面我按设计思路、环境准备、逐模块实现、训练要点、踩坑记录这五块展开。1. 先理解设计图U-Net 为什么长这样1.1 编码器-解码器到底在解决什么问题图像分割的本质是给每一个像素做类别判断所以网络的输出不能像图像分类那样最后压成一个向量而要保留空间上的高和宽。这里就有个真实矛盾像素级判断需要“看清楚细节”而判断一个像素属于“猫”还是“背景”又必须知道它周围一大片的上下文否则光看一个点根本没信息量。感受野这个概念就是关键。一个卷积层看局部堆叠多个卷积后每个输出像素对应输入区域越来越大。想让感受野扩大最直接的办法是下采样比如最大池化把特征图分辨率减半。分辨率小了后面卷积的感受野增长速度就更快同时计算量也大幅下降。这本质和图像压缩一个道理把有效信息凝聚到更少的像素里但代价是空间细节损失边缘、纹理这类高频信息最容易丢。顺着这个逻辑网络自然的形态就出来了前半段不断减小空间尺寸、增大通道数把图像“编码”成高度抽象的特征图后半段再逐步把尺寸放大、通道数恢复从抽象特征中“解码”出逐像素的预测结果。这就是编码器-解码器结构。U-Net 只是把这种思想做成了一套非常干净的模板。1.2 跳跃连接的真实作用如果单纯编码再解码会出现一个明显问题解码器越往高层走越难恢复出精细的边缘和位置信息。为什么因为编码器下采样过程中已经把这些细节抹掉了后面的上采样只能“无中生有”。拿生活里的事情打比方你把一张高清大图压成缩略图再强行放大回来画面总是模糊的因为丢失的像素不可能凭空还原。U-Net 的答案是一条条跳跃连接把编码器每一层的特征图直接拼接到解码器对应层的输入上。底层特征图分辨率高、空间细节全但语义弱高层特征图语义强但分辨率低两者拼接起来正好互补。这个设计带来的好处不仅仅是信息保留还有一个往往被忽略的间接作用跳跃连接让梯度可以绕过中间的瓶颈层直接向解码器传播训练时更稳定深层网络也不容易因为梯度问题而收敛过慢。很多文章喜欢把 U-Net 比作字母“U”的形状其实真正让它好用的核心不是形状而是这些横向拼接的“捷径”。1.3 顺着张量走一遍先建立手感动手写代码之前先把一个 256×256 的输入在标准 U-Net 里各阶段输出张量和通道数走一遍这步搞清楚了后面写代码基本不会错。模块输入形状输出形状初始卷积块1×3×256×2561×64×256×256下采样第1层1×64×256×2561×128×128×128下采样第2层1×128×128×1281×256×64×64下采样第3层1×256×64×641×512×32×32下采样第4层1×512×32×321×1024×16×16上采样第1层1×1024×16×161×512×32×32上采样第2层1×512×32×321×256×64×64上采样第3层1×256×64×641×128×128×128上采样第4层1×128×128×1281×64×256×256输出卷积1×64×256×2561×类别数×256×256可以看到一个非常整齐的规律每经过一次下采样空间尺寸减半、通道数翻倍每经过一次上采样空间尺寸翻倍、通道数减半。通道数翻倍的意义在于分辨率降低后每个像素需要承载的信息变多了用更多通道来表达不同维度的语义特征正好平衡信息总量。这也是为什么这条曲线看起来像字母 U两边的计算量和显存开销大致对称。2. 动手前的准备PyTorch 环境与依赖2.1 PyTorch 安装与版本选择实现 U-Net 其实不需要太新的版本PyTorch 1.13 以上的常用 API 就足够不过建议直接装当前稳定的 2.x 系列性能、编译支持和算子优化都更完善。安装时最关键的是 CUDA 版本匹配装错了会出现装了 GPU 版但torch.cuda.is_available()返回 False 的尴尬局面。# 先确认本地显卡驱动支持的 CUDA 版本 nvidia-smi # 再按照 PyTorch 官方给出的对应命令安装例如 CUDA 11.8 或 12.1 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118装完之后一定要做两件事验证环境导入 PyTorch 并检查 GPU 是否可用顺便看一眼版本号。import torch print(torch.__version__) print(torch.cuda.is_available()) if torch.cuda.is_available(): print(torch.cuda.get_device_name(0))这里想提醒一句如果是在云端 GPU 环境或某些受限网络环境下下载torch这个大体积包可能比较慢推荐用镜像源加速或者直接用云环境预装好的基础镜像再自行安装依赖。2.2 需要用到的库除了 PyTorch 本身下面这些库会贯穿整个实现和数据流程numpy图像数据转数组、mask 处理、指标计算的基本工具。torchvision提供常用的图像变换也有部分预训练模型可以辅助迁移学习。matplotlib可视化原图、mask 和预测结果训练时看效果非常直观。tqdm训练循环进度条每个 epoch 跑了多少一眼就能看到。albumentations做数据增强的利器接口设计比 torchvision 的 transform 更适合分割任务因为它能同步处理 image 和 mask不用担心几何变换导致两者错位。如果不想额外装太多依赖最精简的一套就是torch、numpy、torchvision、matplotlib、tqdm后面用到 albumentations 时会单独说明。环境准备好之后就可以开始写网络结构了。3. 从零手写 U-Net完整 PyTorch 实现3.1 最小单元DoubleConv 模块U-Net 里最基本的模块是“双卷积块”也就是连续两个 3×3 卷积每个卷积后面跟着 BatchNorm 和 ReLU。为什么是两个而不是一个两个 3×3 卷积叠加等价于一个 5×5 卷积的感受野但参数量只有后者的 18/25而且中间多了一次非线性变换模型表达能力更强。所以在 U-Net 的各种实现里几乎每一个层级内部都固定用双卷积。激活函数这里用 ReLU 和 BatchNorm 的组合算是 U-Net 经典实现中最省心、最稳定的配置。BatchNorm 的作用是让每层输入分布尽量稳定对深度网络的收敛速度提升非常明显。要注意的是 Batchnorm 在训练和推理时有行为差异训练时统计当前 batch 的均值和方差并更新全局统计量推理时直接用全局统计量所以代码里必须正确切换model.train()和model.eval()。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().__init__() self.conv nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Conv2d(out_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x)这里padding1是为了保持卷积前后特征图宽高不变。如果去掉 padding每过一个卷积高宽会减 2几个层叠下来后续拼接时对尺寸的处理会变得非常繁琐所以主流实现都统一用padding1。3.2 下采样路径编码器下采样路径做的事情就是“压缩特征”先经过一个 DoubleConv 提取特征再用最大池化把空间尺寸减半然后进入下一层继续提取。最大池化选择 2×2、步长为 2纯操作没有可学习参数计算上很便宜这是它相比步长为 2 的卷积做下采样的明显优势。class Down(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.pool nn.MaxPool2d(kernel_size2, stride2) self.conv DoubleConv(in_channels, out_channels) def forward(self, x): x self.pool(x) x self.conv(x) return x按最标准的 U-Net 配置通道数会按64 - 128 - 256 - 512 - 1024逐层增加。如果你想要更轻量级的版本可以把第一层的 64 改成 32后面继续翻倍参数总量会缩小四倍左右。是否调整需要结合数据集规模和显存大小来判断后面训练章节会具体说。3.3 上采样路径解码器与跳跃连接融合解码器的核心是上采样和跳跃连接融合。上采样我选择转置卷积尺寸 2×2、步长为 2刚好把特征图宽高各放大一倍。转置卷积是有可学习参数的它相比双线性插值更能从数据中学到如何恢复细节。不过它也有副作用步长为 2 时容易产生棋盘格伪影实践中很多项目会用双线性上采样配合 3×3 卷积替代各有取舍。这里先保持经典结构方便和原始论文对照。上采样时必须注意一个关键细节编码器路径的经过最大池化得到的特征图尺寸与解码器上采样后的尺寸可能不完全一致。比如输入尺寸不是 2 的整数次幂时经过交替的偶数次下采样和上采样两个待拼接的特征图可以差一两个像素。所以我要在拼接前对尺寸做裁剪或填充让它们严格对齐否则torch.cat会直接报维度不匹配的错误。class Up(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.up nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size2, stride2) self.conv DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 self.up(x1) diff_y x2.size()[2] - x1.size()[2] diff_x x2.size()[3] - x1.size()[3] x1 F.pad(x1, [ diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2, ]) x torch.cat([x2, x1], dim1) return self.conv(x)这段代码里x1是上采样后的特征图x2是来自编码器对应层的跳跃连接特征图。先计算两边高和宽各自差多少然后对x1做非对称的 padding最后在通道维上拼接起来。拼接之后通道数变成x2的通道数 上采样结果的通道数再送入 DoubleConv 去融合这两部分信息。从通道数的角度看Up的in_channels是来自解码器上一层的通道数out_channels是这一层输出后的通道数。以上采样第一层为例上一层的通道数是 1024转置卷积后变成 512和跳跃连接的 512 通道拼接后重新变成 1024再经过 DoubleConv 输出 512。这个“先减半、再翻倍、再减半”的通道节奏是有意设计的目的是让融合后 DoubleConv 的输入输出通道一致保持信息量的均衡。3.4 组装完整 U-Net有了 Down、Up、DoubleConv 这几个积木组装一个完整 U-Net 就非常简单了。我按经典参数写一遍输入通道数设为in_channels语义分割的类别数设为num_classes。二分类任务时num_classes1多分类任务时就等于实际类别总数。class UNet(nn.Module): def __init__(self, in_channels3, num_classes1): super().__init__() self.inc DoubleConv(in_channels, 64) self.down1 Down(64, 128) self.down2 Down(128, 256) self.down3 Down(256, 512) self.down4 Down(512, 1024) self.up1 Up(1024, 512) self.up2 Up(512, 256) self.up3 Up(256, 128) self.up4 Up(128, 64) self.outc nn.Conv2d(64, 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) out self.outc(x) return outforward里留下的x1到x5就是被保存下来的跳跃连接特征图解码器按顺序把每一层的结果和对应的跳跃连接特征图融合。需要注意最后输出层的1×1卷积没有跟任何激活函数原因稍后训练部分会解释二分类时通常配合 BCEWithLogitsLoss 使用多分类时配合 CrossEntropyLoss 使用。如果想直观检查网络结构是否拼装正确最简单的办法是随机生成一个小张量跑一次前向传播model UNet(in_channels3, num_classes1) x torch.randn(1, 3, 256, 256) out model(x) print(out.shape) # torch.Size([1, 1, 256, 256])3.5 参数规模怎么算出来的“手算”不只看形状还要能大致算出参数量。以down1为例它包含一个 DoubleConv(64, 128)也就是两个 3×3 卷积。第一个卷积权重形状是(128, 64, 3, 3)参数量128×64×9 73728偏置128第二个卷积权重形状是(128, 128, 3, 3)参数量128×128×9 147456偏置128。两层加起来约22 万再加上 BatchNorm 的gain和bias约2×128×2 512。其他层可以按同样公式推下去一个完整的经典 U-Net编码器从 64 通道开始总参数量在 3100 万左右。这个数字意味着什么如果输入是 512×512 的 RGB 图前向一次就能产生不小的显存占用训练时加上梯度和优化器状态一般建议显存至少 8GB。如果你用的显卡只有 4GB 甚至更少要么把输入降到 256×256要么把初始通道从 64 改成 32要么训练阶段开启混合精度这些都是立竿见影的手段。def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) print(f参数量: {count_parameters(model) / 1e6:.2f}M)4. 训练与验证让 U-Net 真正开始分割目标4.1 数据加载与预处理分割任务的数据和分类不太一样训练时需要同时读入原图和标注掩码而且掩码最常见的是单通道灰度图每个像素的数值是这个像素的类别编号。比如二分类分割任务中背景为 0前景为 1多分类任务则可能有 0、1、2、3 等。用 PyTorch 实现数据加载时掩码要转成torch.long而图像要转成torch.float并归一化到[0,1]。from torch.utils.data import Dataset class SegmentationDataset(Dataset): def __init__(self, image_paths, mask_paths, transformNone): 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 cv2.imread(self.image_paths[idx]) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) if self.transform is not None: augmented self.transform(imageimage, maskmask) image augmented[image] mask augmented[mask] # 处理成 tensor image torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 mask torch.from_numpy(mask).long() return image, mask这里有个非常容易犯错的点很多人的 mask 在可视化时看着是黑白图数值却只有 0 和 255而不是 0 和 1。如果直接拿 0、255 去算交叉熵损失模型永远学不对。所以在做数据预处理时必须先确认 mask 像素值范围必要时把 255 统一转成 1。数据增强是分割任务提点最有效的手段之一。翻转、旋转、随机裁剪、缩放都适合操作关键点是几何变换必须同时作用于 image 和 mask保证两者空间对齐。albumentations 库把这种同步处理封装得很干净import albumentations as A train_transform A.Compose([ A.RandomResizedCrop(height256, width256, scale(0.8, 1.0), p0.5), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.2), A.RandomRotate90(p0.5), ])需要留一个心眼随机裁剪成固定尺寸后一定要把所有图像统一成模型能处理的大小不同尺寸的 batch 是完全不能堆叠的。最稳妥的做法是训练前先确认所有图像尺寸一致或者在__getitem__里统一 resize 到固定尺寸。4.2 损失函数与评估指标语义分割里二分类和多分类的损失函数完全不同。二分类场景最常用的就是BCEWithLogitsLoss它内部已经包含了 Sigmoid所以模型最后的输出层不需要再手动加 Sigmoid多分类场景对应CrossEntropyLoss它在通道维度上自动做 Softmax。如果地平线比较极端比如前景只占整张图的 5%直接套 BCE 会被背景主导这种情况建议用 Dice 损失或者 BCE Dice 的组合损失。def dice_loss(pred, target): smooth 1e-6 pred torch.sigmoid(pred) pred pred.view(pred.size(0), -1) target target.view(target.size(0), -1).float() intersection (pred * target).sum(dim1) dice (2.0 * intersection smooth) / (pred.sum(dim1) target.sum(dim1) smooth) return 1.0 - dice.mean()Dice 损失的本质是直接优化预测与真值之间的重叠比例对类别不均衡非常友好但弊端是梯度在小目标上可能不够平滑所以实践中经常把它和交叉熵损失组合使用让交叉熵负责稳定收敛Dice 负责提升分割质量。验证阶段除了看损失值一定要额外看 Dice 系数或 mIoU。损失值下降不代表分割效果好比如极端不均衡的背景下模型全都预测背景BCE 损失可能也不高但分割任务完全失败。写一个简单的 Dice 评估函数每个 epoch 结束后在验证集上算一遍比一味盯着 loss 曲线靠谱。4.3 训练循环的关键设置训练循环的模板大多数项目大同小异但有几个点值得单独强调。第一优化器用 Adam 或 AdamW 都能有不错的表现初始学习率设 1e-3 到 3e-4 之间其中 AdamW 通常对权重衰减的处理更规范分割任务里推荐优先选择。第二训练时开启pin_memoryTrue和适当的num_workers可以明显缓解数据读取瓶颈尤其图像尺寸大、读取频繁的时候。第三混合精度训练不仅省显存还能加速PyTorch 自带torch.cuda.amp用起来不复杂。model UNet(in_channels3, num_classes1).cuda() optimizer torch.optim.AdamW(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemax, patience5) criterion torch.nn.BCEWithLogitsLoss() for epoch in range(100): model.train() epoch_loss 0.0 for images, masks in train_loader: images images.cuda() masks masks.cuda().float().unsqueeze(1) optimizer.zero_grad() outputs model(images) loss criterion(outputs, masks) loss.backward() optimizer.step() epoch_loss loss.item() # 每轮结束做验证 model.eval() with torch.no_grad(): # 计算验证集 Dice 或 mIoU pass这个循环里我额外做了两件事把 masks 用unsqueeze(1)加了一个通道维因为 BCEWithLogitsLoss 要求预测结果和目标的形状完全一致训练过程里保持model.train()验证过程里切换model.eval()否则 BatchNorm 的统计行为会导致验证结果异常抖动。epoch 结束后用ReduceLROnPlateau根据验证指标自动降低学习率能省去很多手动调参的烦恼。5. 踩坑实录U-Net 训练中的高频问题和解法5.1 张量形状不匹配的报错怎么定位RuntimeError: Sizes of tensors must match except in dimension 1是手写 U-Net 最常见的问题几乎每个从零实现的人都会遇到。报错位置通常在torch.cat([x2, x1], dim1)这一行原因就是上采样之后的分辨率和跳跃连接的特征图分辨率不一致。定位思路很简单在报错前把两个张量的 shape 都打印出来看到底差了多少个像素。解决手段有两个第一个是调整输入尺寸让它能被 2 整除足够多次。经典 U-Net 一般做 4 次下采样所以输入的高和宽最好是 16 的倍数而且为了避免中间出现奇数建议直接用 256、512 这样的尺寸第二是像我Up模块里写的那样对尺寸偏小的一边做对称 padding这样即使输入不是 2 的整数次幂也能跑通。实际项目里首选第一种因为固定输入尺寸不仅省心还能防止 BatchNorm 在不同尺寸下统计不一致的问题。5.2 损失值一直不下降损失从训练最开始就不下降甚至持续升高这个问题需要分几个方向排查。先看学习率Adam 的初始学习率 1e-3 是常见起步值如果数据比较简单、网络比较浅这个值可能偏高导致损失震荡不收敛降到 1e-4 再试一轮。再看数据预处理mask 如果是 0 和 255 而不是 0 和 1或者图像没有归一化到合理范围都会让模型无法有效收敛。还可以看标签和预测的类别数是否对得上多分类数据里类别数为 N模型输出通道必须是 N不能手动在通道维度上多写一个背景类。最容易被忽略的是类别不均衡。二分类分割里前景面积只有 5% 甚至更少模型很快学会全预测背景因为损失也低。这时候单单调学习率没用要从损失函数层面解决比如换成 Dice Loss或者给前景类加权重。criterion torch.nn.BCEWithLogitsLoss(pos_weighttorch.tensor([10.0]).cuda())这种pos_weight的用法在 PyTorch 里是专门为二分类正负样本不均衡准备的值可以理解为“过采样正样本的倍数”具体设多少需要通过验证集表现试出来一般从 2 到 10 之间开始调。5.3 显存不足与训练速度过慢U-Net 的显存占用大头集中在编码器的浅层因为这一层特征图分辨率是 256×256即使通道数只有 64也占了可观空间解码器同理。显存不足OOM时最有效的招数是减小输入分辨率256×256 比 512×512 的显存开销少了四倍效果往往只在精度上损失一点点。其次是减小 batch size但 batch size 太小会导致 BatchNorm 不稳定最好不低于 2。再就是混合精度开启后能在几乎不损失精度的情况下把显存占用砍掉将近一半。scaler torch.cuda.amp.GradScaler() for images, masks in train_loader: images images.cuda() masks masks.cuda().float().unsqueeze(1) optimizer.zero_grad() with torch.autocast(device_typecuda, dtypetorch.float16): outputs model(images) loss criterion(outputs, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()训练速度过慢的话除了开混合精度还可以把num_workers调到 4 或 8确保 GPU 在等数据时能少闲置。另一个容易被忽略的点是验证集评估循环里千万别梯度回传记得用torch.no_grad()包裹否则显存会被验证阶段的中间激活值直接占满。如果发现验证集指标在某个 epoch 后开始下降或停滞检查一下是否过拟合数据增强强度可以再加或者引入早停机制看验证集指标连续多个 epoch 不提升就保存当前模型并退出。5.4 二分类输出和损失怎么配合才正确很多新手在输出层加 Sigmoid又在损失里用BCELoss这样也能跑但数值上不够稳。更推荐的做法是输出层保持裸的nn.Conv2d配合BCEWithLogitsLoss因为后者内部做了数值稳定的 Sigmoid 计算。如果要多类别分割输出层同样不带 Softmax直接用CrossEntropyLoss处理它在内部同时完成了 Softmax 和交叉熵计算。只要记住一条规则损失函数要的是原始 logits 还是概率取决于你选哪个 loss别在两个地方同时做激活。做完这些检查U-Net 的训练稳定性和结果一般都能得到明显改善。我在实际项目中踩得最重的坑就是 mask 数值范围和类别权重这两类问题在可视化分割结果时往往看不出来只有直接看损失或者像素值分布才能抓到。最后再分享一个实际很有用的小习惯模型写好之后先用一张合成的小图做一次过拟合测试。比如拿一张 128×128 的纯色背景、中间画一个白色圆形的图让网络去拟合这一张图如果训练十几个 batch 后损失能降到很低说明网络结构本身没问题数据加载和损失也通了再换真实数据训练这样可以快速把代码层面的 bug 和模型层面的错误分开。从零实现模型最忌讳一上来就在几百张图上跑几个小时的训练中途才发现某个模块写错了浪费的时间完全可以用这个小实验避免。