当前位置：首页 > news >正文

052、HAT 模型详解：混合注意力 Transformer 在超分中的创新与代码实现

news 2026/7/5 5:54:55

052、HAT 模型详解：混合注意力 Transformer 在超分中的创新与代码实现

从一次让人抓狂的调试说起

去年秋天，我在一个4倍超分项目上被卡了整整两周。当时用的是SwinIR，效果已经不错了，但老板非要再提0.2dB PSNR。我试了各种trick——加深网络、加通道注意力、换损失函数，结果要么过拟合要么训练崩了。直到某天深夜，我盯着TensorBoard上那条死活上不去的曲线，突然意识到一个问题：SwinIR的窗口注意力虽然高效，但它在局部窗口内做自注意力，天然丢失了跨窗口的长程依赖。而RCAN那种通道注意力虽然能全局建模，但空间细节又不够精细。

这不就是典型的“既要又要”吗？HAT（Hybrid Attention Transformer）就是来解决这个矛盾的。它把通道注意力和空间注意力揉在一起，用了一种很巧妙的方式——不是简单拼接，而是让它们互相补充。今天这篇笔记，我就把HAT的完整实现和踩过的坑都摊开来讲。

HAT的核心思想：别让注意力打架

先看HAT的整体结构。它延续了SwinIR的U型架构，但每个Transformer Block里塞了两个注意力模块：一个通道注意力（Channel Attention），一个空间注意力（Spatial Attention）。这两个模块是串行连接的，但内部设计有讲究。

通道注意力用的是SE-like的结构，但加了一个小trick——它把输入特征先做全局平均池化，然后经过两个全连接层，最后用sigmoid激活得到通道权重。这里有个细节：第一个全连接层做降维（减少参数量），第二个恢复维度。降维比例我一般设4或8，太小了通道间交互不够，太大了参数量爆炸。

空间注意力部分，HAT没有用常见的卷积加sigmoid那种简单方案，而是用了自注意力机制。具体来说，它把特征图分成若干窗口，在每个窗口内做自注意力。但这里有个关键区别：窗口大小和SwinIR的窗口大小可以不一样。我试过把空间注意力的窗口设成8x8，而SwinIR的窗口是7x7，这样能捕捉不同尺度的空间关系。

# 这里踩过坑：通道注意力和空间注意力的顺序不能乱classHybridAttention(nn.Module):def__init__(self,dim,num_heads,window_size):super().__init__()self.channel_attn=ChannelAttention(dim)# 先做通道self.spatial_attn=SpatialAttention(dim,num_heads,window_size)# 再做空间defforward(self,x):# 别这样写：先空间后通道，效果会差0.1-0.2dBx=self.channel_attn(x)x=self.spatial_attn(x)returnx

为什么通道注意力要放在前面？我的理解是：通道注意力先做全局重标定，相当于给每个通道打上重要性标签，这样空间注意力在后续处理时就能更聚焦于重要通道的细节。如果反过来，空间注意力先做，它可能会被噪声通道干扰，导致注意力图不干净。

代码实现中的三个关键细节

1. 通道注意力的降维比例

通道注意力的核心代码很简单，但降维比例的选择有讲究。我见过有人直接用dim//16，结果小模型效果还行，大模型直接崩了。经验值是：当dim小于256时，比例用4；dim在256-512之间用8；dim大于512用16。

classChannelAttention(nn.Module):def__init__(self,dim,reduction=8):super().__init__()# 这里踩过坑：reduction不能太大，否则信息丢失严重self.fc=nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(dim,dim//reduction,1,bias=False),nn.ReLU(inplace=True),nn.Conv2d(dim//reduction,dim,1,bias=False),nn.Sigmoid())defforward(self,x):# 别这样写：直接用nn.Linear代替Conv2d，会丢失空间结构信息b,c,h,w=x.shape y=self.fc(x)returnx*y

注意这里我用的是Conv2d而不是Linear，因为Conv2d能保持特征图的形状，避免reshape操作带来的额外开销。而且Conv2d的1x1卷积本质上就是全连接，但更高效。

2. 空间注意力的窗口划分

空间注意力部分，我直接复用了SwinIR的窗口划分逻辑，但窗口大小单独设置。这里有个容易忽略的点：窗口大小必须能整除特征图尺寸，否则需要做padding。我一般设成8或16，这样大多数特征图都能整除。

classSpatialAttention(nn.Module):def__init__(self,dim,num_heads,window_size):super().__init__()self.window_size=window_size self.num_heads=num_heads# 这里踩过坑：qkv的投影维度必须能被num_heads整除self.qkv=nn.Linear(dim,dim*3,bias=False)self.proj=nn.Linear(dim,dim)defforward(self,x):b,c,h,w=x.shape# 别这样写：直接对整个特征图做自注意力，显存会爆炸# 正确的做法是划分窗口x=window_partition(x,self.window_size)# 窗口内的自注意力计算x=self.window_attention(x)x=window_reverse(x,self.window_size,h,w)returnx

窗口划分的代码我直接抄的SwinIR，但加了一个小优化：如果特征图尺寸小于窗口大小，就退化为全局自注意力。这个情况在浅层特征中很少出现，但深层特征（比如下采样后）可能会遇到。

3. 混合注意力的残差连接

HAT的每个Block都有两个残差连接：一个在通道注意力之后，一个在空间注意力之后。但这两个残差连接的缩放系数不同。通道注意力的残差系数是0.1，空间注意力的是0.2。这个系数是我调参调出来的，太小了梯度传不过去，太大了训练不稳定。

classHATBlock(nn.Module):def__init__(self,dim,num_heads,window_size):super().__init__()self.norm1=nn.LayerNorm(dim)self.attn=HybridAttention(dim,num_heads,window_size)self.norm2=nn.LayerNorm(dim)self.ffn=FeedForward(dim)# 这里踩过坑：残差系数不能一样，否则通道注意力的效果会被淹没self.ca_scale=0.1self.sa_scale=0.2defforward(self,x):shortcut=x x=self.norm1(x)x=self.attn(x)# 别这样写：直接x = x + shortcut，梯度会爆炸x=shortcut+self.ca_scale*x# 通道注意力残差shortcut=x x=self.norm2(x)x=self.ffn(x)x=shortcut+self.sa_scale*x# 空间注意力残差returnx