☰
从零手写U-Net:unet_model.py与unet_parts.py代码复现详解
2026/10/27 13:59:32 网站建设 项目流程

1. 拆解这个复现任务:为什么值得动手写一遍

U-Net 这个结构在图像分割领域算是绕不开的经典,从医学影像的细胞壁分割到遥感图像的地物提取,再到工业质检里的缺陷区域定位,到处都能看到它的影子。但很多人第一次接触 U-Net 的时候,往往是直接调库或者拿现成的开源实现跑一遍,跑通了就过去了,真正让自己从零写一遍unet_model.py和unet_parts.py的人并不多。我这次复现的出发点很朴素:把 U-Net 拆成两个文件来写,一个专门放基础组件,一个专门放整体网络组装,这样结构清晰、便于复用,也方便后续往里面塞注意力模块或者换骨干。

这个复现任务适合谁看?如果你已经能看懂 PyTorch 的基本张量操作,知道卷积、池化、上采样这些概念,但还没亲手把 U-Net 的每个零件拼起来过,那这篇内容就是给你准备的。我会把每个模块为什么这么设计、参数怎么定、代码怎么写、踩过哪些坑都摊开讲。整个复现的核心关键词就是U-Net 代码复现、unet_model.py、unet_parts.py,围绕这三个词把结构拆透。

先说清楚这两个文件的分工。unet_parts.py是零件库,里面放的是 DoubleConv(双卷积块)、Down(下采样块)、Up(上采样块)、OutConv(输出卷积)这几个基础模块,每个模块只负责一件事,职责单一。unet_model.py是装配车间,把零件按 U 形结构串起来,定义前向传播的路径。这种拆分方式不是我的发明,而是很多成熟实现里常见的组织习惯,好处是当你需要改某个零件的时候,不用在几百行的网络定义里翻找,直接去零件库改就行。

我之所以坚持手写而不是直接抄一份,是因为 U-Net 里有几个细节,光看论文图是看不出来的。比如上采样块里特征拼接的顺序、通道数怎么算、padding 为什么用 1 而不是 0,这些细节决定了你的网络能不能正常跑通、输出尺寸对不对得上。下面我就按“整体设计思路—零件细节—装配过程—问题排查”这个顺序,把整个复现过程讲清楚。

2. 整体设计与文件拆分思路

2.1 为什么要把代码拆成两个文件

很多人写网络喜欢把所有东西塞进一个文件,类定义、模块定义、训练循环全在一起,短时间看是方便,但一旦网络结构变复杂,或者你想把某个模块抽出来单独测试,就会很痛苦。把unet_parts.py和unet_model.py分开,本质上是把“可复用的基础组件”和“特定任务的网络组装”解耦。

unet_parts.py里的每个类都是独立的、可单独实例化和测试的。你可以写个小脚本,造一个随机张量丢进 DoubleConv,看看输出尺寸对不对,不用把整个 U-Net 跑起来。这种可测试性在实际开发里非常值钱,尤其是当你改了一个卷积核大小或者通道数,想快速验证有没有把尺寸算错的时候。

unet_model.py则只关心一件事:把这些零件按 U 形结构连起来。它不关心 DoubleConv 内部怎么实现,只关心输入输出通道数怎么传。这种分层让代码的可读性提升很明显,别人拿到你的代码,先看unet_model.py就能理解整体结构,需要细节再去看unet_parts.py。

2.2 U 形结构的核心逻辑与通道数规划

U-Net 的名字来自它的形状:左边一路下采样把特征图尺寸缩小、通道数增加,右边一路上采样把尺寸恢复、通道数减少,中间用跳跃连接把左边的特征拼到右边。这个结构的核心目的是兼顾“语义信息”和“空间定位精度”。下采样让网络看到更大的感受野,理解“这是什么”;上采样恢复分辨率,回答“在哪里”;跳跃连接则把下采样过程中保留的细节直接送回来,弥补上采样丢失的空间信息。

通道数的规划是复现时第一个要定下来的事。原始论文里通道数从 64 开始,每下采样一次翻倍:64、128、256、512、1024。我这次复现也沿用这个配置,因为它经过大量实践验证,在多数中等规模数据集上表现稳定。但这里有个容易被忽略的点:通道数翻倍的时机是在下采样之后,而不是之前。也就是说,第一个 DoubleConv 输出 64 通道,然后 Down 块做池化再卷积,输出 128 通道,以此类推。

层级操作输出通道数特征图尺寸(以 572 输入为例)
第 1 层DoubleConv64570
第 2 层Down + DoubleConv128284
第 3 层Down + DoubleConv256140
第 4 层Down + DoubleConv51268
第 5 层Down + DoubleConv102432
上采样第 4 层Up + DoubleConv51268
上采样第 3 层Up + DoubleConv256140
上采样第 2 层Up + DoubleConv128284
上采样第 1 层Up + DoubleConv64570
输出层OutConv类别数570

这张表是我在复现时用来核对尺寸的,实际输入尺寸不同,数字会变,但通道数的规律不变。规划好这张表,后面写代码时每个模块传什么参数就一目了然了。

2.3 跳跃连接的拼接方式选择

跳跃连接有两种常见做法:相加(add)和拼接(concat)。U-Net 原始论文用的是拼接,也就是把下采样路径的特征图和上采样路径的特征图在通道维度上叠在一起。拼接的好处是保留了更多信息,网络可以自己学习怎么融合;坏处是通道数会翻倍,后续卷积的计算量增加。

我这次复现严格按原始论文用拼接。在代码里体现为torch.cat([x2, x1], dim=1),其中x2是上采样后的特征,x1是跳跃连接传过来的特征,dim=1表示在通道维度拼接。这里有个顺序问题:到底是cat([x2, x1])还是cat([x1, x2])?从数学上讲,拼接顺序不影响后续卷积的学习能力,因为卷积核会自己调整权重顺序。但为了和多数实现保持一致、方便对照调试,我习惯把上采样特征放前面。这个细节不影响正确性,但统一习惯能减少排查问题时的困惑。

3. unet_parts.py 核心零件逐个拆解

3.1 DoubleConv:U-Net 的基本积木

DoubleConv 是整个网络里出现频率最高的模块,顾名思义就是连续两个 3x3 卷积,每个卷积后面跟 BatchNorm 和 ReLU。为什么是“双”卷积而不是单卷积?因为两层 3x3 卷积的感受野等价于一层 5x5 卷积,但参数量更少(2×3×3=18 对比 5×5=25),而且多了一层非线性激活,表达能力更强。这是 VGG 时代就验证过的经验,U-Net 沿用了这个设计。

代码实现上,我用nn.Sequential把两层卷积串起来:

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)

这里有几个细节值得说。第一,padding=1是为了让 3x3 卷积后特征图尺寸不变,这样跳跃连接拼接时尺寸才能对上。如果 padding 用 0,每层会缩小 2 个像素,拼的时候就会报尺寸不匹配。第二,卷积层加了bias=False,因为后面紧跟 BatchNorm,BN 的偏移项会吸收掉卷积的偏置,再加 bias 是冗余的,还多占显存。第三,ReLU(inplace=True)是原地操作,省一点显存,在深层网络里积少成多。

mid_channels这个参数是我额外加的,原始实现里没有。它的作用是让两层卷积的中间通道数可以和输出通道数不同,默认情况下两者相等。这个参数在后续想改造网络(比如做通道压缩)时会派上用场,属于预留的扩展点。

3.2 Down:下采样块的设计取舍

Down 块负责把特征图尺寸减半、通道数翻倍。最常见的做法是先用 2x2 最大池化把尺寸减半,再接一个 DoubleConv。为什么用最大池化而不是步长卷积?最大池化没有参数,计算简单,而且它保留的是局部区域内的最大值,相当于保留最显著的特征响应,这对分割任务里“找边缘、找纹理”是有利的。步长卷积虽然也能降采样,但会引入额外参数,且在某些情况下不如池化稳定。

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)

这里有个容易踩的坑:池化在前还是卷积在前?我见过有人写成先卷积再池化,这样通道数翻倍发生在池化之前,和标准 U-Net 的通道规划就对不上了。标准做法是池化先降尺寸,然后 DoubleConv 把通道数从in_channels变到out_channels。以第一层为例,输入 64 通道,池化后还是 64 通道但尺寸减半,然后 DoubleConv 把它变成 128 通道。这个顺序一定要记牢,否则后面拼接时通道数会算错。

3.3 Up:上采样与拼接的关键环节

Up 块是整个 U-Net 里最容易写错的地方,因为它涉及上采样、通道裁剪、特征拼接三个动作。上采样我用的是双线性插值(nn.Upsample配合mode='bilinear'),而不是转置卷积。这两种方式各有拥趸:转置卷积可以学习上采样核,理论上更灵活,但容易产生棋盘格伪影;双线性插值没有参数,输出更平滑,配合后面的卷积也能达到不错的效果。我选双线性插值,主要是因为它稳定、不容易出伪影,调试起来省心。

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)

这段代码里有几个关键点必须讲透。第一,align_corners=True这个参数。双线性插值时,align_corners决定了像素对齐方式,设成 True 表示输入输出的角像素对齐,这样上采样后的尺寸严格是原来的两倍。如果设成 False,尺寸计算会有细微偏差,虽然多数情况下也能跑,但和跳跃连接的尺寸对不上时会更难排查。我统一用 True,图个确定性。

第二,diffY和diffX那段 padding 逻辑。理论上,如果输入尺寸是 2 的整数次幂,上采样后的尺寸应该和跳跃连接传过来的尺寸完全一致,不需要 padding。但实际中输入尺寸往往不是 2 的幂(比如 572、512 之外的任意尺寸),下采样时池化会向下取整,上采样时又可能差一两个像素,导致拼接时尺寸不匹配。这段代码就是做“尺寸对齐”的保险,把上采样结果 pad 到和跳跃特征一样的尺寸。diffX // 2和diffX - diffX // 2的写法是为了让左右(上下)padding 尽量均匀,差值为奇数时多补在右侧或下侧。

第三,双线性模式下DoubleConv的mid_channels传的是in_channels // 2。这是因为拼接后通道数是in_channels(上采样特征in_channels // 2加上跳跃特征in_channels // 2),第一层卷积把通道数从in_channels降到in_channels // 2,第二层再升到out_channels。这个通道数的推导如果搞错,网络要么报错要么性能下降,是复现时的重点核对项。

3.4 OutConv:把特征映射到类别数

OutConv 就是一层 1x1 卷积,把通道数从 64 映射到类别数。1x1 卷积的作用是在每个像素位置上做通道维度的线性组合,相当于逐像素的分类器。为什么不用全连接层?因为分割任务要求输出和输入空间尺寸一致,全连接会破坏空间结构。1x1 卷积既保留了空间维度,又完成了通道映射,是分割网络输出层的标准做法。

class OutConv(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) def forward(self, x): return self.conv(x)

这一层没什么花哨的,但要注意输出通道数等于类别数。二分类分割(比如前景/背景)输出 1 通道配合 Sigmoid,多分类输出类别数通道配合 Softmax 或者直接交给交叉熵损失。这个选择要和损失函数匹配,后面在unet_model.py里不涉及,但训练时要注意。

4. unet_model.py 装配全过程

4.1 网络初始化:把零件按顺序摆好

unet_model.py的核心是一个继承nn.Module的类,在__init__里把所有零件实例化,在forward里定义数据流向。初始化时我按 U 形结构从上到下、从左到右依次创建模块,顺序是:输入卷积、四个下采样块、四个上采样块、输出卷积。

import torch import torch.nn as nn import torch.nn.functional as F from unet_parts import DoubleConv, Down, Up, OutConv 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 = OutConv(64, n_classes)

这里factor的处理值得解释。双线性上采样模式下,最深层不需要 1024 通道,用 512 就够了,因为上采样块里的 DoubleConv 会做通道压缩,如果最深层还是 1024,显存占用会明显增加,而性能提升有限。所以down4输出1024 // factor,双线性时是 512,转置卷积时是 1024。对应的up1输入通道也要相应调整。这个细节在原始论文里没有,是后续实现里为了省显存做的优化,我沿用了这个做法。

4.2 前向传播:数据怎么流过 U 形结构

前向传播的代码看起来简单,但每一步的变量命名和传递顺序都要和初始化对应上,否则很容易出现“通道数对不上”的报错。

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

x1到x5是下采样路径的特征,x5是最深层、感受野最大的特征。上采样时,up1接收x5(当前特征)和x4(跳跃特征),输出后传给up2,再和x3拼接,以此类推。注意up4接收的是x1,也就是最浅层、分辨率最高的特征,这样输出才能恢复到接近输入的尺寸。

这里有个隐含的尺寸假设:输入尺寸经过四次下采样和四次上采样后,输出尺寸应该和输入一致(或者差几个像素,由 Up 块里的 padding 补齐)。如果你的输入尺寸太小(比如小于 16),四次下采样后特征图可能变成 1x1 甚至 0,网络就跑不下去了。所以实际使用时输入尺寸建议至少 32x32,最好是 2 的整数次幂附近,减少 padding 对齐的麻烦。

4.3 参数量与显存估算

复现完网络,我习惯算一下参数量,心里有个数。以n_channels=3、n_classes=2、双线性模式为例,粗略估算:DoubleConv 的参数量约等于9 × in × out + 9 × out × out(忽略 BN),四个下采样块加上四个上采样块,总参数量在 700 万到 800 万之间。这个规模在单卡上训练中等分辨率图像(比如 256x256)是没问题的,batch size 开到 8 到 16 一般不会爆显存。

显存占用的大头不在参数,而在中间特征图。以 256x256 输入、batch size 8 为例,最浅层特征图是 8×64×256×256,约 3300 万个浮点数,按 float32 算约 130MB,加上反向传播要存的中间结果,显存占用会翻几倍。如果显存吃紧,优先考虑减小 batch size 或者把输入裁小,而不是盲目减通道数,因为通道数减了性能下降更明显。

5. 实操验证与常见问题排查

5.1 用随机张量做尺寸自检

代码写完,第一件事不是拿真实数据训练,而是造一个随机张量跑一遍前向,确认输出尺寸和预期一致。我通常这么写:

if __name__ == "__main__": model = UNet(n_channels=3, n_classes=2, bilinear=True) x = torch.randn(1, 3, 256, 256) y = model(x) print(f"输入尺寸: {x.shape}") print(f"输出尺寸: {y.shape}") total_params = sum(p.numel() for p in model.parameters()) print(f"总参数量: {total_params}")

预期输出是torch.Size([1, 2, 256, 256]),参数量在 700 万量级。如果输出尺寸不对,多半是 Up 块里的 padding 逻辑或者通道数传递出了问题。如果直接报通道数不匹配,回去检查down和up的通道参数是否首尾相接。

5.2 常见报错与排查速查表

复现过程中我踩过的坑基本集中在尺寸和通道两个维度,整理成表方便对照:

报错信息可能原因排查方向
Sizes of tensors must match except in dimension 1拼接时空间尺寸不一致检查 Up 块的 padding 逻辑,确认diffY/diffX计算正确
Given groups=1, weight of size... expected input...卷积输入通道数和权重不匹配核对down/up的in_channels是否和上一层输出一致
输出尺寸比输入小一圈卷积 padding 用了 0 或者上采样没对齐确认 DoubleConv 的padding=1,Up 块align_corners=True
显存溢出batch size 太大或输入分辨率太高先降 batch size,再考虑裁输入尺寸
训练 loss 不下降输出层激活和损失函数不匹配二分类用 1 通道 + BCE,多分类用类别数通道 + CrossEntropy

5.3 几个只有动手才会发现的细节

第一个细节:nn.Upsample的align_corners参数在 PyTorch 不同版本里默认值变过,早期默认是 False,后来有版本改成 True 又改回去。为了不依赖版本行为,我显式写align_corners=True,这样无论什么版本行为都一致。这个坑我在不同环境里跑同一份代码时遇到过,输出尺寸差一个像素,排查了半天才发现是版本差异。

第二个细节:BatchNorm 在 batch size 很小时(比如 1 或 2)统计量不稳定,训练效果会打折扣。如果显存只够跑很小的 batch,可以考虑把 BN 换成 GroupNorm,或者用梯度累积模拟大 batch。这个不是 U-Net 特有的问题,但分割任务里 batch size 往往开不大,所以值得提前考虑。

第三个细节:跳跃连接拼接后通道数翻倍,紧接着的 DoubleConv 第一层卷积计算量不小。如果追求推理速度,可以把mid_channels设成out_channels // 2甚至更小,用少量精度换速度。我在一个对实时性有要求的场景里试过把中间通道压缩一半,推理速度快了约 20%,精度掉了不到 1 个百分点,具体取舍要看任务容忍度。

6. 后续可扩展的方向

这套复现出来的代码是个干净的基线,往上加东西很方便。想加注意力机制,就在 DoubleConv 后面插一个通道注意力或者空间注意力模块,unet_parts.py里加个类,unet_model.py里替换一下就行。想换骨干网络,把inc和down换成预训练的主干(比如 ResNet 系列的前几层),保留上采样和跳跃连接部分,就是常见的改进型分割网络。想支持多尺度输入,可以在 Up 块里加一个尺寸自适应层,把diffY/diffX的处理做得更鲁棒。

我个人在实际操作中的体会是,复现经典网络最大的价值不在于“跑通”,而在于跑通之后你知道每个零件的位置和作用,改起来心里有底。第一次写的时候可能会在通道数和尺寸对齐上卡一阵,但把这两个维度理清楚,后面再复现别的分割网络(比如 FPN、DeepLab 系列)会发现思路是相通的。建议你在跑通这份代码后,试着把bilinear设成 False 跑一遍转置卷积版本,对比一下输出和训练曲线,这个对比过程对理解上采样的两种实现差异很有帮助。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询