U-Net 模型 Python 代码实现

正义的彬彬侠

936人浏览 · 2025-10-21 00:33:03

正义的彬彬侠 · 2025-10-21 00:33:03 发布

import torch
from torch import nn
import torch.nn.functional as F

# 一次卷积操作，包括 卷积 + BN（可选） + ReLU
class ConvBNReLU(nn.Module):
    """
    Conv + BN[optional] + ReLU
    """
    def __init__(self, in_ch, out_ch, isBN=True):
        super(ConvBNReLU, self).__init__()
        self.isBN = isBN
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1)
        self.relu = nn.ReLU(inplace=True)
        if isBN:
            self.bn = nn.BatchNorm2d(out_ch)        

    def forward(self, x):
        x = self.conv(x)
        if self.isBN:
            x = self.bn(x)
        x = self.relu(x)
        return x

# 两次卷积，可以选择是否预先进行MaxPool
class DoubleConv(nn.Module):
    """
    MaxPool[optional] + ConvBNReLU + ConvBNReLU
    """
    def __init__(self, in_ch, out_ch, isBN=True, is_pool=False):
        super(DoubleConv, self).__init__()
        self.is_pool = is_pool
        if is_pool:
            self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv1 = ConvBNReLU(in_ch, out_ch, isBN)
        self.conv2 = ConvBNReLU(out_ch, out_ch, isBN)

    def forward(self, x):
        if self.is_pool:
            x = self.maxpool(x)
        x = self.conv1(x)
        x = self.conv2(x)
        return x

# 上采样，包括跳线连接encoder部分的feature map
class UpsampleConv(nn.Module):
    """
    skip feature ---------------|
    [x2 Upsample | Deconv] -- concat -- DoubleConv
    """
    def __init__(self, in_ch, skip_ch, out_ch, isDeconv=True, isBN=True):
        super(UpsampleConv, self).__init__()
        if isDeconv:
            self.up = nn.ConvTranspose2d(in_ch, in_ch, kernel_size=2, stride=2)
        else:
            self.up = nn.UpsamplingBilinear2d(scale_factor=2)
        self.double_conv = DoubleConv(in_ch + skip_ch, out_ch, isBN, is_pool=False)

    def forward(self, x, x_skip):
        x = self.up(x)
        # 对于奇数尺寸的下采样，需要进行padding才能与跳连的特征图同尺寸
        # 如 5x5 -> 2x2，上采样后 4x4，diffX = diffY = 1, pad: (0, 1, 0, 1)
        diffX = x_skip.size()[3] - x.size()[3]
        diffY = x_skip.size()[2] - x.size()[2]
        x = F.pad(x, (diffX // 2, diffX - diffX//2,  diffY // 2, diffY - diffY//2))
        x = torch.cat([x, x_skip], dim=1)
        x = self.double_conv(x)
        return x

# 激活函数选择，支持sigmoid、softmax和无激活函数
def select_activation(name=None):
    if not name or name.lower() == 'identity':
        return nn.Identity()
    elif name.lower() == 'sigmoid':
        return nn.Sigmoid()
    elif name.lower() == 'softmax':
        return nn.Softmax()
    else:
        raise ValueError(f"current unsupport activation name {name}")


# 最后一层，用于输出大小和通道数符合任务需求的特征图
class SegmentHead(nn.Module):
    """
    Conv + Upsample[optional] + Activation[optional]
    """
    def __init__(self, in_ch, out_ch, act=None, upsample=1):
        super(SegmentHead, self).__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=1)
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=upsample)
        self.act = select_activation(act)

    def forward(self, x):
        x = self.conv(x)
        x = self.upsample(x)
        x = self.act(x)
        return x


# 利用上面定义的模块，搭建UNet结构
class UNet(nn.Module):

    def __init__(self, img_ch=3, base_ch=16, num_class=2, isBN=True, isDeconv=False):
        super(UNet, self).__init__()
        # Encoder 部分
        self.conv_in = DoubleConv(img_ch, base_ch, isBN, is_pool=False)
        self.down1 = DoubleConv(base_ch, base_ch * 2, isBN, is_pool=True)
        self.down2 = DoubleConv(base_ch * 2, base_ch * 4, isBN, is_pool=True)
        self.down3 = DoubleConv(base_ch * 4, base_ch * 8, isBN, is_pool=True)
        self.down4 = DoubleConv(base_ch * 8, base_ch * 16, isBN, is_pool=True)
        # Decoder 部分
        self.up1 = UpsampleConv(base_ch * 16, base_ch * 8, base_ch * 8, isDeconv, isBN)
        self.up2 = UpsampleConv(base_ch * 8, base_ch * 4, base_ch * 4, isDeconv, isBN)
        self.up3 = UpsampleConv(base_ch * 4, base_ch * 2, base_ch * 2, isDeconv, isBN)
        self.up4 = UpsampleConv(base_ch * 2, base_ch, base_ch, isDeconv, isBN)
        self.seghead = SegmentHead(base_ch, num_class, act=None, upsample=1)

    def forward(self, x_in):
        feat = self.conv_in(x_in)
        print("feat size : ", feat.size())
        d1 = self.down1(feat)
        print("d1 size : ", d1.size())
        d2 = self.down2(d1)
        print("d2 size : ", d2.size())
        d3 = self.down3(d2)
        print("d3 size : ", d3.size())
        d4 = self.down4(d3)
        print("d4 size : ", d4.size())
        u1 = self.up1(d4, d3)
        print("u1 size : ", u1.size())
        u2 = self.up2(u1, d2)
        print("u2 size : ", u2.size())
        u3 = self.up3(u2, d1)
        print("u3 size : ", u3.size())
        u4 = self.up4(u3, feat)
        print("u4 size : ", u4.size())
        out = self.seghead(u4)
        return out


x_in = torch.randn((1, 3, 224, 224))
unet = UNet(img_ch=3, base_ch=64, num_class=3, isBN=True, isDeconv=False)
x_out = unet(x_in)
print("unet input size : ", x_in.size())
print("unet output size : ", x_out.size())

代码实现了U-Net模型，这是一种经典的语义分割网络架构，广泛用于医疗图像分割等领域。U-Net的核心特点是U形结构，包括编码器（Encoder，下采样提取特征）、解码器（Decoder，上采样恢复分辨率）和跳跃连接（Skip Connections，融合多尺度特征）。该实现支持可选的Batch Normalization（BN）和上采样方式（转置卷积或双线性插值），并通过模块化设计（如ConvBNReLU、DoubleConv等）提高了代码的可读性和灵活性。

下面将按照代码的逻辑结构，从导入库开始，一步一步逐段分析。每段解释包括代码的目的、关键概念、数学公式（如果适用）、实现细节，以及为什么这样设计。代码整体分为辅助模块定义（ConvBNReLU、DoubleConv、UpsampleConv、select_activation、SegmentHead）和主模型类（UNet），以及测试部分。假设输入图像为[1, 3, 224, 224]（batch=1，RGB通道=3，高度/宽度=224），base_ch=64，num_class=3，isBN=True，isDeconv=False。

基于计算，特征图尺寸变化如下（所有尺寸均为[batch, channels, height, width]）：

feat: [1, 64, 224, 224]
d1: [1, 128, 112, 112]
d2: [1, 256, 56, 56]
d3: [1, 512, 28, 28]
d4: [1, 1024, 14, 14]
u1: [1, 512, 28, 28]
u2: [1, 256, 56, 56]
u3: [1, 128, 112, 112]
u4: [1, 64, 224, 224]
输入: [1, 3, 224, 224]
输出: [1, 3, 224, 224]

这些尺寸通过卷积公式计算得出：卷积输出尺寸 $Hout=Hin+2×padding−kernel_sizestride+1H_{out} = \frac{H_{in} + 2 \times padding - kernel\_size}{stride} + 1$ （stride默认为1）；池化/上采样为2倍变化。实际运行时，代码中的print语句会输出这些尺寸。

1. 导入库

import torch
from torch import nn
import torch.nn.functional as F

目的：导入PyTorch核心库，用于定义神经网络和张量操作。
详细解释：
- import torch：PyTorch主库，提供张量（如torch.randn）和自动微分。
- from torch import nn：神经网络模块，提供Conv2d、ReLU、BatchNorm2d、ConvTranspose2d、UpsamplingBilinear2d等层。
- import torch.nn.functional as F：函数式模块，提供pad（填充）和cat（拼接）等操作，无需实例化。
为什么需要这些库：U-Net是一个深度学习模型，需要PyTorch的模块化构建。F用于forward中的pad操作，避免边界尺寸不匹配。

2. 单次卷积模块：ConvBNReLU

# 一次卷积操作，包括 卷积 + BN（可选） + ReLU
class ConvBNReLU(nn.Module):
    """
    Conv + BN[optional] + ReLU
    """
    def __init__(self, in_ch, out_ch, isBN=True):
        super(ConvBNReLU, self).__init__()
        self.isBN = isBN
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1)
        self.relu = nn.ReLU(inplace=True)
        if isBN:
            self.bn = nn.BatchNorm2d(out_ch)        

    def forward(self, x):
        x = self.conv(x)
        if self.isBN:
            x = self.bn(x)
        x = self.relu(x)
        return x

目的：定义基本卷积单元，包括卷积、可选BN和ReLU激活，用于构建更复杂的块。
详细解释：
- __init__：初始化层。nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)：3x3卷积，padding=1保持输入输出尺寸相同。isBN控制是否添加nn.BatchNorm2d（批标准化，加速训练、稳定梯度）。nn.ReLU(inplace=True)：就地激活，节省内存。
- forward：前向传播。顺序：Conv → [BN] → ReLU。
数学概念：卷积： $yi,j,c=∑k=0Cin−1∑m,nxi+m,j+n,k⋅wm,n,k,c+bcy_{i,j,c} = \sum_{k=0}^{C_{in}-1} \sum_{m,n} x_{i+m, j+n, k} \cdot w_{m,n,k,c} + b_c$ 。BN： $x^=x−μσ2+ϵ⋅γ+β\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta$ 。
为什么这样设计：模块化，便于复用。isBN可选，支持实验（如无BN的轻量模型）。U-Net原版无BN，但添加BN可提升性能。

3. 双卷积模块：DoubleConv

# 两次卷积，可以选择是否预先进行MaxPool
class DoubleConv(nn.Module):
    """
    MaxPool[optional] + ConvBNReLU + ConvBNReLU
    """
    def __init__(self, in_ch, out_ch, isBN=True, is_pool=False):
        super(DoubleConv, self).__init__()
        self.is_pool = is_pool
        if is_pool:
            self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv1 = ConvBNReLU(in_ch, out_ch, isBN)
        self.conv2 = ConvBNReLU(out_ch, out_ch, isBN)

    def forward(self, x):
        if self.is_pool:
            x = self.maxpool(x)
        x = self.conv1(x)
        x = self.conv2(x)
        return x

目的：定义U-Net的基本构建块，双Conv单元，可选预池化，用于编码器下采样。
详细解释：
- __init__：is_pool控制是否添加nn.MaxPool2d（2x2最大池化，stride=2，下采样到1/2）。然后两个ConvBNReLU：第一个改变通道（in_ch → out_ch），第二个保持（out_ch → out_ch）。
- forward：顺序：[MaxPool] → Conv1 → Conv2。
尺寸变化：如果is_pool=True，尺寸减半；Conv保持尺寸（padding=1）。
为什么：U-Net每个编码层使用双Conv增强非线性表达。is_pool用于下采样路径，便于编码器构建。

4. 上采样模块：UpsampleConv

# 上采样，包括跳线连接encoder部分的feature map
class UpsampleConv(nn.Module):
    """
    skip feature ---------------|
    [x2 Upsample | Deconv] -- concat -- DoubleConv
    """
    def __init__(self, in_ch, skip_ch, out_ch, isDeconv=True, isBN=True):
        super(UpsampleConv, self).__init__()
        if isDeconv:
            self.up = nn.ConvTranspose2d(in_ch, in_ch, kernel_size=2, stride=2)
        else:
            self.up = nn.UpsamplingBilinear2d(scale_factor=2)
        self.double_conv = DoubleConv(in_ch + skip_ch, out_ch, isBN, is_pool=False)

    def forward(self, x, x_skip):
        x = self.up(x)
        # 对于奇数尺寸的下采样，需要进行padding才能与跳连的特征图同尺寸
        # 如 5x5 -> 2x2，上采样后 4x4，diffX = diffY = 1, pad: (0, 1, 0, 1)
        diffX = x_skip.size()[3] - x.size()[3]
        diffY = x_skip.size()[2] - x.size()[2]
        x = F.pad(x, (diffX // 2, diffX - diffX//2,  diffY // 2, diffY - diffY//2))
        x = torch.cat([x, x_skip], dim=1)
        x = self.double_conv(x)
        return x

目的：定义解码器上采样单元，包括上采样、跳跃连接拼接和双Conv精炼。
详细解释：
- __init__：isDeconv选择上采样方式：True为nn.ConvTranspose2d（转置卷积，可学习）；False为nn.UpsamplingBilinear2d（双线性插值，固定）。然后DoubleConv处理拼接后特征（通道in_ch + skip_ch → out_ch）。
- forward：上采样（2倍）→ pad填充（处理奇数尺寸不匹配，如224偶数无pad）→ torch.cat拼接（dim=1通道维）→ DoubleConv。
数学概念：转置卷积输出： $Hout=(Hin−1)×2+2H_{out} = (H_{in} - 1) \times 2 + 2$ （k=2, s=2, p=0）。pad：对称填充确保尺寸匹配。
为什么：跳跃连接融合浅层细节（x_skip）和深层语义（x），提升边界精度。pad处理池化导致的尺寸问题（U-Net常见）。

5. 激活函数选择：select_activation

# 激活函数选择，支持sigmoid、softmax和无激活函数
def select_activation(name=None):
    if not name or name.lower() == 'identity':
        return nn.Identity()
    elif name.lower() == 'sigmoid':
        return nn.Sigmoid()
    elif name.lower() == 'softmax':
        return nn.Softmax()
    else:
        raise ValueError(f"current unsupport activation name {name}")

目的：提供灵活的输出激活函数选择，支持二元（sigmoid）、多类（softmax）或无激活（identity）。
详细解释：根据name返回对应模块。nn.Identity()：恒等映射，无变化。nn.Sigmoid()： $σ(x)=11+e−x\sigma(x) = \frac{1}{1 + e^{-x}}$ 。nn.Softmax()： $Softmax(xi)=exi∑exj\text{Softmax}(x_i) = \frac{e^{x_i}}{\sum e^{x_j}}$ （默认dim=1，但代码未指定，实际调用需注意）。
为什么：语义分割任务多样（如二元用sigmoid，多类用softmax）。默认identity，便于在损失函数中处理（如CrossEntropyLoss内置softmax）。

6. 输出头：SegmentHead

# 最后一层，用于输出大小和通道数符合任务需求的特征图
class SegmentHead(nn.Module):
    """
    Conv + Upsample[optional] + Activation[optional]
    """
    def __init__(self, in_ch, out_ch, act=None, upsample=1):
        super(SegmentHead, self).__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=1)
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=upsample)
        self.act = select_activation(act)

    def forward(self, x):
        x = self.conv(x)
        x = self.upsample(x)
        x = self.act(x)
        return x

目的：最终层，将解码器输出转换为目标通道和尺寸的分割图。
详细解释：
- __init__：1x1卷积（in_ch → out_ch，降维到num_class）。nn.UpsamplingBilinear2d(scale_factor=upsample)：可选上采样（默认1，无变化）。act从select_activation选择。
- forward：Conv → Upsample → Act。
为什么：适应不同任务（如多类分割）。upsample支持额外放大（虽默认1，但可扩展）。

7. 主模型：UNet

# 利用上面定义的模块，搭建UNet结构
class UNet(nn.Module):

    def __init__(self, img_ch=3, base_ch=16, num_class=2, isBN=True, isDeconv=False):
        super(UNet, self).__init__()
        # Encoder 部分
        self.conv_in = DoubleConv(img_ch, base_ch, isBN, is_pool=False)
        self.down1 = DoubleConv(base_ch, base_ch * 2, isBN, is_pool=True)
        self.down2 = DoubleConv(base_ch * 2, base_ch * 4, isBN, is_pool=True)
        self.down3 = DoubleConv(base_ch * 4, base_ch * 8, isBN, is_pool=True)
        self.down4 = DoubleConv(base_ch * 8, base_ch * 16, isBN, is_pool=True)
        # Decoder 部分
        self.up1 = UpsampleConv(base_ch * 16, base_ch * 8, base_ch * 8, isDeconv, isBN)
        self.up2 = UpsampleConv(base_ch * 8, base_ch * 4, base_ch * 4, isDeconv, isBN)
        self.up3 = UpsampleConv(base_ch * 4, base_ch * 2, base_ch * 2, isDeconv, isBN)
        self.up4 = UpsampleConv(base_ch * 2, base_ch, base_ch, isDeconv, isBN)
        self.seghead = SegmentHead(base_ch, num_class, act=None, upsample=1)

    def forward(self, x_in):
        feat = self.conv_in(x_in)
        print("feat size : ", feat.size())
        d1 = self.down1(feat)
        print("d1 size : ", d1.size())
        d2 = self.down2(d1)
        print("d2 size : ", d2.size())
        d3 = self.down3(d2)
        print("d3 size : ", d3.size())
        d4 = self.down4(d3)
        print("d4 size : ", d4.size())
        u1 = self.up1(d4, d3)
        print("u1 size : ", u1.size())
        u2 = self.up2(u1, d2)
        print("u2 size : ", u2.size())
        u3 = self.up3(u2, d1)
        print("u3 size : ", u3.size())
        u4 = self.up4(u3, feat)
        print("u4 size : ", u4.size())
        out = self.seghead(u4)
        return out

目的：组装U-Net模型，包括编码器（5层下采样）、解码器（4层上采样）和输出头。
详细解释：
- __init__：参数包括输入通道img_ch、基础通道base_ch（通道倍增，如64→128→…→1024）、类别数num_class、isBN、isDeconv。编码器：conv_in（初始双Conv，无池化）；down1-4（带池化，下采样）。解码器：up1-4（上采样+跳连）。seghead输出。
- forward：编码：x_in → feat（初始） → d1-d4（下采样）。解码：u1（d4+d3） → u2（u1+d2） → … → u4（u3+feat）。out：seghead处理。print调试尺寸。
为什么：标准U-Net结构，base_ch控制模型大小（小值轻量）。跳连从浅到深融合。print便于验证尺寸匹配。

8. 测试部分

x_in = torch.randn((1, 3, 224, 224))
unet = UNet(img_ch=3, base_ch=64, num_class=3, isBN=True, isDeconv=False)
x_out = unet(x_in)
print("unet input size : ", x_in.size())
print("unet output size : ", x_out.size())

目的：创建随机输入，实例化模型，运行前向传播，并打印输入/输出尺寸。
详细解释：torch.randn生成模拟输入。unet实例化（base_ch=64，num_class=3，双线性上采样）。x_out = unet(x_in)调用forward。print验证输出恢复原尺寸。
为什么：调试模型，确保端到端工作。输出通道=num_class（3），尺寸与输入相同。

总体总结

代码流程：辅助模块 → UNet组装 → 测试验证。
适用场景：语义分割，如医疗图像（肿瘤分割）。灵活参数支持自定义（BN提升稳定性，双线性上采样减少伪影）。
潜在改进：添加Dropout防过拟合；用ResNet替换编码器；训练时用DiceLoss或CrossEntropyLoss。
局限：对奇数尺寸需pad；无act的输出适合logits输入损失函数。

AI Agent技术社区

Agent 垂直技术社区，欢迎活跃、内容共建。

更多推荐

ego-lite：这浏览器让AI Agent和你能同时用，再也不抢标签页了

上周末逛GitHub Trending，前十里面有五个跟AI Agent有关，其中有个项目我点进去之后愣是看了半小时——egolite。它不是那种"又一个浏览器自动化框架"的套娃项目。它解决了一个特别具体但特别烦人的问题。用Claude Code或者Codex做过浏览器自动化的人应该都懂那种痛。你让Agent帮你填个表单，它就在你浏览器里一顿乱跳，标签页飞来飞去，你啥也干不了。然后你想，那我给

AI Agent技术社区

AI Agnet

AI Agent。Agent 和普通聊天机器人有什么区别？Workflow 和 Agentic Workflow 是什么？Tool、Function Calling、Plugin、Skill 又有什么不同？RAG 算不算记忆？MCP 和 A2A 分别是干什么的？一个 Agent 不够，为什么还要搞多智能体？这些名词单独看都不难，难的是它们经常混在一起出现。所以这篇文章不讲太复杂的理论，而是尽量用通

AI Agent技术社区

reverse-skill：用 AI Agent 路由器统一调度安全研究工具链的工程实践

不是一个安全工具，而是一个。它试图解决的核心问题是：当 Claude Code、Cursor 这类代码 AI 客户端面对一个安全分析任务时，它并不知道该用jadx还是apktool，该走 APK 逆向路径还是 ELF 二进制路径。这个项目通过预定义的路由规则、标准化的工作流 Playbook 和自动化的工具索引，让 AI Agent 不再靠猜，而是按结构化流程执行。这处于这一演进阶段的典型产物——