UNet改进（5）：线性注意力机制（Linear Attention）-原理详解与代码实现

拾忆-eleven2025-06-24 14:22

引言

在计算机视觉领域，UNet架构因其在图像分割任务中的卓越表现而广受欢迎。近年来，注意力机制的引入进一步提升了UNet的性能。本文将深入分析一个结合了线性注意力机制的UNet实现，探讨其设计原理、代码实现以及在医学图像分割等任务中的应用潜力。

UNet架构概述

UNet最初由Ronneberger等人提出，主要用于生物医学图像分割。其独特的U形结构由编码器（下采样路径）和解码器（上采样路径）组成，通过跳跃连接将低层特征与高层特征相结合，既保留了空间信息又利用了深层的语义信息。

传统的UNet结构简单有效，但随着研究的深入，人们发现引入注意力机制可以显著提升模型性能，特别是在处理复杂场景和微小结构时。

线性注意力机制

注意力机制的基本概念

注意力机制的核心思想是让模型能够"关注"输入数据中最相关的部分。在传统的自注意力机制中，计算复杂度通常是O(N²)，这对于高分辨率图像来说计算成本很高。

线性注意力实现

在我们的实现中，采用了线性注意力机制来降低计算复杂度。以下是关键的LinearAttention类实现：

python 复制代码 
class LinearAttention(nn.Module):
    def __init__(self, channels):
        super(LinearAttention, self).__init__()
        self.query = nn.Conv2d(channels, channels // 8, kernel_size=1)
        self.key = nn.Conv2d(channels, channels // 8, kernel_size=1)
        self.value = nn.Conv2d(channels, channels, kernel_size=1)
        self.gamma = nn.Parameter(torch.zeros(1))
        
    def forward(self, x):
        batch_size, C, height, width = x.size()
        
        # 计算query, key, value
        q = self.query(x).view(batch_size, -1, height * width).permute(0, 2, 1)  # (B, N, C')
        k = self.key(x).view(batch_size, -1, height * width)  # (B, C', N)
        v = self.value(x).view(batch_size, -1, height * width)  # (B, C, N)
        
        # 线性注意力计算
        kv = torch.bmm(k, v)  # (B, C', C)
        z = 1 / (torch.bmm(q, k.sum(dim=2, keepdim=True)) + 1e-6)  # (B, N, 1)
        attn = torch.bmm(q, kv)  # (B, N, C)
        out = attn * z  # (B, N, C)
        
        out = out.view(batch_size, C, height, width)
        return self.gamma * out + x

这个实现有几个关键特点：

  1. 通道缩减：通过将通道数减少到1/8来降低计算复杂度

  2. 线性复杂度：通过矩阵乘法的重新排列，将复杂度从O(N²)降低到O(N)

  3. 可学习的gamma参数：控制注意力特征与原始特征的混合比例

网络组件详解

双卷积块

双卷积块是UNet的基本构建模块，包含两个连续的3x3卷积层，每个卷积层后接批量归一化和ReLU激活函数。我们的实现增加了可选的注意力机制：

python 复制代码 
class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels, use_attention=False):
        super(DoubleConv, self).__init__()
        self.use_attention = use_attention
        
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        
        if use_attention:
            self.attention = LinearAttention(out_channels)
    
    def forward(self, x):
        x = self.double_conv(x)
        if self.use_attention:
            x = self.attention(x)
        return x

下采样模块

下采样模块由最大池化层和双卷积块组成：

python 复制代码 
class Down(nn.Module):
    def __init__(self, in_channels, out_channels, use_attention=False):
        super(Down, self).__init__()
        self.downsampling = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            DoubleConv(in_channels, out_channels, use_attention)
        )
    
    def forward(self, x):
        return self.downsampling(x)

上采样模块

上采样模块使用转置卷积进行上采样，然后与编码路径的特征图拼接，最后通过双卷积块：

python 复制代码 
class Up(nn.Module):
    def __init__(self, in_channels, out_channels, use_attention=False):
        super(Up, self).__init__()
        self.upsampling = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_channels, out_channels, use_attention)
    
    def forward(self, x1, x2):
        x1 = self.upsampling(x1)
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

完整的UNet架构

结合上述组件，我们构建了完整的UNet模型：

python 复制代码 
class UNet(nn.Module):
    def __init__(self, in_channels=1, num_classes=1):
        super(UNet, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        
        # 编码器部分
        self.in_conv = DoubleConv(in_channels, 64, use_attention=True)
        self.down1 = Down(64, 128, use_attention=True)
        self.down2 = Down(128, 256, use_attention=True)
        self.down3 = Down(256, 512, use_attention=True)
        self.down4 = Down(512, 1024)
        
        # 解码器部分
        self.up1 = Up(1024, 512, use_attention=True)
        self.up2 = Up(512, 256, use_attention=True)
        self.up3 = Up(256, 128, use_attention=True)
        self.up4 = Up(128, 64, use_attention=True)
        
        self.out_conv = OutConv(64, num_classes)
    
    def forward(self, x):
        # 编码路径
        x1 = self.in_conv(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.out_conv(x)

这个架构有几个值得注意的特点：

  1. 对称结构：编码器和解码器基本对称，但最深层的下采样块没有使用注意力机制

  2. 渐进式通道变化：通道数从64开始，每次下采样翻倍，直到1024

  3. 广泛的注意力应用：除了最深层的下采样，其他所有层都应用了注意力机制

注意力机制的应用策略

在我们的实现中，注意力机制的应用策略值得关注：

  1. 编码路径：前四个下采样块中，前三个使用了注意力机制

  2. 解码路径：所有上采样块都使用了注意力机制

  3. 输入输出：输入卷积和最终输出卷积没有使用注意力机制

这种策略基于以下考虑：

  • 深层特征已经具有高度抽象性，可能不需要额外的注意力

  • 解码路径需要精确的定位，注意力机制尤为重要

  • 输入输出层结构简单，注意力机制的收益可能不明显

性能优化考虑

  1. 内存效率：线性注意力显著降低了内存消耗

  2. 计算效率：通过通道缩减和线性复杂度计算保持高效

  3. 数值稳定性：在注意力计算中添加了小常数(1e-6)防止除零错误

实际应用建议

  1. 医学图像分割：这种结构特别适合CT/MRI图像分割任务

  2. 参数调整：可以根据任务复杂度调整注意力层的位置和数量

  3. 输入通道：当前设置为1通道输入，适用于灰度医学图像

扩展可能性

  1. 多模态输入：修改输入通道数以适应RGB或多模态医学图像

  2. 深度监督：在解码路径中添加辅助输出

  3. 注意力变体：尝试其他类型的注意力机制如通道注意力

结论

本文详细分析了一个结合线性注意力机制的UNet实现。这种架构在保持UNet原有优势的同时，通过精心设计的注意力机制提升了模型对重要特征的关注能力。线性注意力的引入使得模型在高分辨率图像上也能高效运行，为医学图像分割等任务提供了有力的工具。

代码实现展示了如何将现代注意力机制与传统UNet架构有机结合，这种模式也可以应用于其他视觉任务的网络设计中。读者可以根据具体任务需求调整注意力层的位置和数量，找到最佳的性能平衡点。

随着注意力机制的不断发展，我们期待看到更多高效、精准的UNet变体出现，推动医学图像分析和其他视觉任务的进步。

完整代码

如下：

python 复制代码 
import torch.nn as nn
import torch
import math

class LinearAttention(nn.Module):
    def __init__(self, channels):
        super(LinearAttention, self).__init__()
        self.query = nn.Conv2d(channels, channels // 8, kernel_size=1)
        self.key = nn.Conv2d(channels, channels // 8, kernel_size=1)
        self.value = nn.Conv2d(channels, channels, kernel_size=1)
        self.gamma = nn.Parameter(torch.zeros(1))
        
    def forward(self, x):
        batch_size, C, height, width = x.size()
        
        # 计算query, key, value
        q = self.query(x).view(batch_size, -1, height * width).permute(0, 2, 1)  # (B, N, C')
        k = self.key(x).view(batch_size, -1, height * width)  # (B, C', N)
        v = self.value(x).view(batch_size, -1, height * width)  # (B, C, N)
        
        # 线性注意力计算
        kv = torch.bmm(k, v)  # (B, C', C)
        z = 1 / (torch.bmm(q, k.sum(dim=2, keepdim=True)) + 1e-6)  # (B, N, 1)
        attn = torch.bmm(q, kv)  # (B, N, C)
        out = attn * z  # (B, N, C)
        
        out = out.view(batch_size, C, height, width)
        return self.gamma * out + x

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels, use_attention=False):
        super(DoubleConv, self).__init__()
        self.use_attention = use_attention
        
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        
        if use_attention:
            self.attention = LinearAttention(out_channels)
    
    def forward(self, x):
        x = self.double_conv(x)
        if self.use_attention:
            x = self.attention(x)
        return x

class Down(nn.Module):
    def __init__(self, in_channels, out_channels, use_attention=False):
        super(Down, self).__init__()
        self.downsampling = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            DoubleConv(in_channels, out_channels, use_attention)
        )
    
    def forward(self, x):
        return self.downsampling(x)

class Up(nn.Module):
    def __init__(self, in_channels, out_channels, use_attention=False):
        super(Up, self).__init__()
        self.upsampling = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_channels, out_channels, use_attention)
    
    def forward(self, x1, x2):
        x1 = self.upsampling(x1)
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

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

class UNet(nn.Module):
    def __init__(self, in_channels=1, num_classes=1):
        super(UNet, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        
        # 编码器部分
        self.in_conv = DoubleConv(in_channels, 64, use_attention=True)
        self.down1 = Down(64, 128, use_attention=True)
        self.down2 = Down(128, 256, use_attention=True)
        self.down3 = Down(256, 512, use_attention=True)
        self.down4 = Down(512, 1024)
        
        # 解码器部分
        self.up1 = Up(1024, 512, use_attention=True)
        self.up2 = Up(512, 256, use_attention=True)
        self.up3 = Up(256, 128, use_attention=True)
        self.up4 = Up(128, 64, use_attention=True)
        
        self.out_conv = OutConv(64, num_classes)
    
    def forward(self, x):
        # 编码路径
        x1 = self.in_conv(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.out_conv(x)

model = UNet(in_channels=1, num_classes=1)
相关推荐
xingshanchang4 小时前
PyTorch 不支持旧GPU的异常状态与解决方案：CUDNN_STATUS_NOT_SUPPORTED_ARCH_MISMATCH
人工智能·pytorch·python
reddingtons5 小时前
Adobe Firefly AI驱动设计：实用技巧与创新思维路径
大数据·人工智能·adobe·illustrator·photoshop·premiere·indesign
CertiK5 小时前
IBW 2025: CertiK首席商务官出席，探讨AI与Web3融合带来的安全挑战
人工智能·安全·web3
Deepoch6 小时前
Deepoc 大模型在无人机行业应用效果的方法
人工智能·科技·ai·语言模型·无人机
Deepoch6 小时前
Deepoc 大模型：无人机行业的智能变革引擎
人工智能·科技·算法·ai·动态规划·无人机
kngines6 小时前
【字节跳动】数据挖掘面试题0003：有一个文件，每一行是一个数字，如何用 MapReduce 进行排序和求每个用户每个页面停留时间
人工智能·数据挖掘·mapreduce·面试题
Binary_ey6 小时前
AR衍射光波导设计遇瓶颈，OAS 光学软件来破局
人工智能·软件需求·光学软件·光波导
昵称是6硬币6 小时前
YOLOv11: AN OVERVIEW OF THE KEY ARCHITECTURAL ENHANCEMENTS目标检测论文精读（逐段解析）
图像处理·人工智能·深度学习·yolo·目标检测·计算机视觉
平和男人杨争争7 小时前
机器学习2——贝叶斯理论下
人工智能·机器学习
静心问道7 小时前
XLSR-Wav2Vec2：用于语音识别的无监督跨语言表示学习
人工智能·学习·语音识别
热门推荐
01Coze实战第13讲：飞书多维表格读取+豆包生图模型，轻松批量生成短剧封面02如何将markdown生成pdf03Coze扣子平台完整体验和实践（附国内和国际版对比）04华为昇腾 910B 部署 DeepSeek-R1 蒸馏系列模型详细指南05手机电脑之间快速传输图片视频文件，不压缩画质、不限制大小的方法！06扣子（coze）实战|我用扣子搭建了一个自动分析小红薯笔记内容的AI应用|详细步骤拆解07微软账户和本地账户有什么区别？如何切换登录账户？08鸽了六年的某大厂面试题：你会手写一个模板引擎吗？09DeepSeek各版本说明与优缺点分析10GPU 进阶笔记（二）：华为昇腾 910B GPU