前言
本文提出了用于低分辨率图像分割的MaskAttn - UNet框架,并将其核心的掩码注意力机制集成到YOLOv11中。传统U - Net类模型难以捕捉全局关联,Transformer类模型计算量大,而掩码注意力机制通过可学习的掩码,让模型选择性关注重要区域,融合了卷积的局部效率和注意力的全局视野。其工作流程包括特征适配、掩码生成、定向注意力计算和特征融合。我们将掩码注意力机制代码集成到YOLOv11中。
文章目录: YOLOv11改进大全:卷积层、轻量化、注意力机制、损失函数、Backbone、SPPF、Neck、检测头全方位优化汇总
专栏链接: YOLOv11改进专栏
介绍

摘要
低分辨率图像分割在机器人技术、增强现实和大规模场景理解等实际应用中至关重要。在这些场景中,由于计算资源限制,高分辨率数据往往难以获取。为解决这一挑战,我们提出了一种新颖的分割框架MaskAttn-UNet,它通过掩码注意力机制对传统UNet架构进行了优化。该模型能够选择性地突出重要区域,同时抑制无关背景,从而在杂乱复杂场景中提升分割精度。与传统UNet变体不同,MaskAttn-UNet有效平衡了局部特征提取与全局上下文感知能力,使其特别适用于低分辨率输入场景。我们在三个基准数据集上对该方法进行了评估,所有输入图像均调整为128×128分辨率,结果表明其在语义分割、实例分割和全景分割任务中均展现出具有竞争力的性能。实验结果显示,MaskAttn-UNet的精度可与当前最先进方法媲美,且计算成本远低于基于Transformer的模型,为资源受限场景下的低分辨率图像分割提供了高效且可扩展的解决方案。
文章链接
论文地址:论文地址
代码地址:代码地址
基本原理
掩码注意力模块是MaskAttn-UNet模型的核心创新组件,核心目标是在低分辨率图像分割场景中,高效平衡“局部细节捕捉”与“全局关联建模”,同时避免传统注意力机制的算力浪费,其原理可从核心设计逻辑、工作流程、关键特性三方面展开:
一、核心设计逻辑
该模块的核心思路是“选择性关注”——不像传统自注意力机制那样对图像中所有像素进行无差别全局计算,也不像纯卷积那样局限于局部区域,而是通过一个“可学习的掩码”(类似智能筛选器),让模型自动聚焦于对分割任务有用的区域(如物体轮廓、关键结构、前景目标),同时抑制无意义的背景噪音或冗余信息。
其设计初衷是解决两大痛点:
- 传统U-Net类模型:依赖卷积的局部性,难以捕捉图像中远距离物体的关联(如重叠物体、分散目标的整体特征),导致复杂场景分割模糊;
- Transformer类模型:全局自注意力计算量大(像素间两两匹配),内存和算力消耗极高,不适合低分辨率、资源受限的实际场景。
因此,掩码注意力模块本质是“卷积的局部效率”与“注意力的全局视野”的融合——用掩码筛选关键区域,只在有用区域内进行注意力计算,实现“精准且高效”的特征提取。
二、完整工作流程
模块的工作过程可拆解为4个关键步骤,全程围绕“筛选-计算-融合-优化”展开:
- 特征格式适配:先接收来自编码器或解码器的特征图(包含图像的局部细节和初步语义信息),并调整其格式,使其适配后续注意力计算的需求;
- 掩码生成与筛选:自动学习一个二进制掩码(可理解为一张“关注地图”),地图上的“高亮区域”对应图像中需要重点关注的部分(如物体边缘、前景目标),“暗区”对应无关背景。这个掩码是动态学习的,会根据不同图像、不同场景自适应调整,而非固定规则;
- 定向注意力计算:采用多头注意力机制(共4个注意力头,相当于从多个角度捕捉特征),但仅在掩码筛选后的“高亮区域”内计算像素间的关联。比如,对于低分辨率图像中的小物体,掩码会聚焦于该物体的像素范围,让这些像素间相互传递信息,从而强化物体的整体特征,同时忽略背景像素的无效关联;
特征融合与优化:将注意力计算后的特征,与原始输入的特征通过“残差连接”融合(保留初始的局部细节),再经过两层前馈网络进一步优化特征质量,最终输出“既包含局部精准细节,又融入全局关键关联”的增强特征。
鲁棒性强:掩码能有效抑制背景噪音,在复杂场景(如 clutter 杂乱环境、重叠物体、光线变化)中,仍能精准区分前景目标与背景,提升分割的稳定性。
核心代码
class Mask2FormerAttention(nn.Module):
def __init__(self, channels, size):
super(Mask2FormerAttention, self).__init__()
self.channels = channels
self.size = size
self.query = nn.Linear(channels, channels)
self.key = nn.Linear(channels, channels)
self.value = nn.Linear(channels, channels)
self.mask = None
self.norm = nn.LayerNorm([channels])
def forward(self, x):
batch_size, channels, height, width = x.size()
if channels != self.channels:
raise ValueError("Input channel size does not match initialized channel size.")
x = x.view(batch_size, channels, height * width).permute(0, 2, 1)
Q = self.query(x)
K = self.key(x)
V = self.value(x)
scores = torch.matmul(Q, K.transpose(-2, -1))
scores = scores / (self.channels ** 0.5)
if self.mask is None or self.mask.size(-1) != height * width:
binary_mask = torch.randint(0, 2, (batch_size, height, width), device=x.device)
binary_mask = binary_mask.view(batch_size, -1)
processed_mask = torch.where(binary_mask > 0.5, torch.tensor(0.0, device=x.device), torch.tensor(-float('inf'), device=x.device))
self.mask = processed_mask.unsqueeze(1).expand(-1, height * width, -1)
scores = scores + self.mask
attention_weights = F.softmax(scores, dim=-1)
attention_output = torch.matmul(attention_weights, V)
attention_output = attention_output + x
attention_output = self.norm(attention_output)
return attention_output.view(batch_size, channels, height, width)
结果
