DRAMiTransformer：空间与通道注意力融合的高效视觉变换块

随着深度学习在计算机视觉领域的快速发展，自注意力机制（Self-Attention）逐渐成为各种任务的核心组件。然而，传统自注意力机制往往需要较高的计算成本，并且难以有效地捕捉到空间和渠道之间的细粒度依赖关系。

基于此，我们提出了一种名为DRAMiTransformer的新型变换块，该块结合了空间与通道注意力机制以及可移动变体设计（Mobile Variants），在保持高效的同时，显著提升了模型对多尺度特征的表达能力。本文将详细解读这一新架构，并通过PyTorch代码示例，展示其具体实现。

1. 模型概述

DRAMiTransformer 是一种轻量级的视觉变换块，主要由以下几个核心组件构成：

QKV 投影（QKV Projection）：用于生成查询、键和值向量。
空间自注意力（Spatial Self-Attention, SP-SA）：专注于捕捉二维空间中相邻像素的关系。
通道间注意力（Channel-Wise Self-Attention, CH-SA）：用于建模不同通道之间的全局依赖关系。
可移动变体（Mobile Variants, MobVi）：一种轻量级的深度可分离卷积模块，用于特征融合。
归一化层（Norm Layer） 和 前馈网络（Feed-Forward Network, FFN）：用于非线性变换和规范化。

这些组件协同工作，使得模型在捕捉空间与通道交互的同时，保持了较低的计算复杂度。其核心思想是通过分而治之的方式分别建模空间与通道依赖，然后通过轻量级变换模块进行融合。

2. 模型实现代码解读

接下来，我们将对代码的主要部分进行逐步解析，帮助理解每一项组件的功能和作用。

（1）QKV 投影层

DRAMiTransformer 的第一部分是 QKV 投影：

self.qkv_proj = QKVProjection(dim, num_head, qkv_bias=qkv_bias)

其中，QKVProjection 类的作用是将输入的特征图投影到查询、键和值空间。具体实现如下：

class QKVProjection(nn.Module):def __init__(self, dim, num_head, qkv_bias=True):super().__init__()self.dim = dimself.num_head = num_headhead_dim = dim // (num_head * 2)  # 调整通道数# 分别为查询、键和值权重分配空间self.query = nn.Linear(dim, dim)self.key = nn.Linear(dim, dim)self.value = nn.Linear(dim, dim)def forward(self, x):B, C, H, W = x.shape  # 假设输入x为(B,C,H,W)x_flat = x.view(-1, C)  # 展开为(B*H*W, C)q = self.query(x_flat).view(B, H * W, -1, self.dim_head)k = self.key(x_flat).view(B, H * W, -1, self.dim_head)v = self.value(x_flat).view(B, H * W, -1, self.dim_head)return q, k, v

功能解析：

QKVProjection 类的作用是将输入特征图 (B,C,H,W) 投影到查询、键和值空间。
通过线性变换生成三组张量 q (query)、k (key) 和 v (value)，用于后续的注意力计算。

（2）空间与通道自注意力

DRAMiTransformer 中的空间和通道注意力机制分别定义了两个独立的注意力模块：

空间自注意力（Spatial Self-Attention, SP-SA）

self.sp_attn = SpatialSelfAttention()

其中，SpatialSelfAttention 类的具体实现如下：

class SpatialSelfAttention(nn.Module):def __init__(self):super().__init__()def forward(self, x):# x: (B, H*W, C)B, N, C = x.shape  # 分别是批次大小、区域数和通道数# 生成相似度矩阵attention = (x @ x.transpose(1, 2)) / math.sqrt(C)  # (B, N, N)attention = F.softmax(attention, dim=-1)# 加权求和out = x @ attention  # (B, N, C)return out

功能解析：

SpatialSelfAttention 模块用于计算相邻区域的相似度矩阵，并通过 softmax 和加权求和生成注意力权重。
输出的 out 张量捕捉了输入特征图的空间关系。

通道间注意力（Channel-Wise Self-Attention, CH-SA）

self.ch_attn = ChannelWiseSelfAttention()

其中，ChannelWiseSelfAttention 类的具体实现如下：

class ChannelWiseSelfAttention(nn.Module):def __init__(self):super().__init__()def forward(self, x):# x: (B, C, H*W)B, C, N = x.shape  # 分别是批次大小、通道数和区域数# 计算通道间的相似度attention = (x @ x.transpose(1, 2)) / math.sqrt(N)  # (B, C, C)attention = F.softmax(attention, dim=-1)# 加权求和out = x @ attention  # (B, C, N)return out.permute(0, 2, 1).reshape(B, -1, C)  # 调整形状以适应后续操作

功能解析：

ChannelWiseSelfAttention 模块用于建模通道间的全局依赖关系。
输出的 out 张量捕捉了不同通道之间的关联性。

（3）可移动变体（Mobile Variants）

为了保持模型的轻量化特性，DRAMiTransformer引入了一种基于深度可分离卷积的轻量级变换模块：

self.mobivari = MobileVariants(dim)

其中，MobileVariants 类的主要实现如下：

class MobileVariants(nn.Module):def __init__(self, dim):super().__init__()self.dim = dim# 深度可分离卷积：3x3 卷积核 + 点乘卷积self.depth_conv = nn.Conv2d(dim, dim, kernel_size=3, padding=1, groups=dim)self.point_conv = nn.Conv2d(dim, dim, kernel_size=1)def forward(self, x):# 深度可分离卷积out = self.depth_conv(x)# 点乘卷积进行通道重混合out = self.point_conv(out)return out

功能解析：

MobileVariants 类通过深度可分离卷积实现了轻量化的特征变换，同时保持了通道间的交互能力。
3x3 深度卷积用于提取空间特征，1x1 点乘卷积用于重混合通道信息。

（4）主干网络整合

最后，DRAMiTransformer 的主干网络将上述模块整合到一起：

def forward(self, x):B, C, H, W = x.shape  # 假设输入x为(B,C,H,W)# 生成查询、键和值张量q, k, v = self.qkv_projection(x)  # (B, H*W, C)，每个输出通道数相同# 空间自注意力计算sp_attn_out = self.sp_attn(q.permute(0, 2, 1).reshape(B, C, H, W))# 通道间自注意力计算ch_attn_out = self.ch_attn(v.permute(0, 2, 1).reshape(B, C, H*W))# 轻量化变换mobivari_out = self.mobivari(x)# 最终的特征融合out = sp_attn_out + ch_attn_out + mobivari_outreturn out