自编基于层结构（Layer）的添加自注意力机制

自编基于层结构（Layer）的添加自注意力机制

直接开撕！传统神经网络层结构那套全连接+激活函数的组合拳早就看腻了，今天咱们整点刺激的——给网络层装个自注意力插件。这玩意儿能让网络自己决定哪些信息重要，比无脑全连接不知道高到哪里去了。

先看这个基础层结构怎么改：

class AttentionLayer(nn.Module): def __init__(self, dim, heads=4): super().__init__() self.heads = heads self.scale = dim ** -0.5 # 这个缩放因子千万别忘 self.to_qkv = nn.Linear(dim, dim*3, bias=False) # 输出前再加个全连接 self.proj = nn.Sequential( nn.Linear(dim, dim), nn.Dropout(0.1) )

注意看to_qkv这行，一石三鸟直接把输入转换成查询、键、值三个向量。这里有个骚操作——用单个线性层同时生成QKV，比分开写三个层省事儿多了，实测还能减少参数冲突。

核心计算部分才是重头戏：

def forward(self, x): b, n, _, h = *x.shape, self.heads # 生成QKV并拆分成多头 [重要！] qkv = self.to_qkv(x).chunk(3, dim=-1) q, k, v = map(lambda t: t.reshape(b, n, h, -1).transpose(1, 2), qkv) # 注意力能量计算（矩阵乘法搞起） dots = (q @ k.transpose(-2, -1)) * self.scale attn = dots.softmax(dim=-1) # 信息聚合与还原形状 out = (attn @ v).transpose(1, 2).reshape(b, n, -1) return self.proj(out)

这里有几个坑要注意：1) chunk拆解时维度要对齐；2) 多头reshape的顺序影响计算效率；3) 缩放因子不加模型直接爆炸。建议在调试时先print下各维度变化，别问我怎么知道的。

实际使用时可以像乐高积木一样插入网络：

class SuperNet(nn.Module): def __init__(self): super().__init__() self.layers = nn.Sequential( nn.Linear(256, 512), AttentionLayer(512), # 这里插入！ nn.ReLU(), nn.Linear(512, 10) )

注意输入维度要和注意力层的dim参数对齐。实测在NLP任务中，这种结构对长距离依赖捕捉效果拔群，比单纯堆LSTM省显存不说，在GPU上还能并行加速。

最后说个骚操作：把传统卷积和自注意力混搭使用，前几层用CNN抓局部特征，后面接注意力层搞全局关系。这种组合拳在图像分类任务中效果意外的好，不信你试试？代码改起来也简单，把上面的AttentionLayer直接插到卷积后面就完事。

遇到维度不匹配别慌，记住万能调试三步法：1) print各层输入输出形状；2) 检查矩阵乘法维度对齐；3) 梯度裁剪别超过1e3。自注意力虽好，可不要贪杯哦，head数太多小心显存爆炸！