PyTorch实现CIFAR10单头注意力机制报forward参数错误

我在CIFAR10数据集上实现注意力机制，目标是搭建仅含单个注意力头的注意力层，参考开源Transformer PyTorch实现中的多头注意力代码编写逻辑，运行时触发如下报错：

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1108         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110             return forward_call(*input, **kwargs)
   1111         # Do not call functions when jit is used
   1112         full_backward_hooks, non_full_backward_hooks = [], []

TypeError: forward() takes 1 positional argument but 2 were given

我清楚该报错属于基础编程范畴问题，但同时怀疑自己编写的单头注意力层逻辑本身存在错误，需要对应的问题排查提示，参考实现代码如下：

class ScaledDotProductAttention(nn.Module):
    def __init__(self, input_dim, output_dim,  attn_dropout=0.1):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        
        
        self.q = nn.Linear(input_dim, output_dim, bias=False)
        self.k = nn.Linear(input_dim, output_dim, bias=False)
        self.v = nn.Linear(input_dim, output_dim, bias=False)
        # print(self.q, self.k, self.v)
        # self.dropout = nn.Dropout(attn_dropout)
        # print(self.dropout)
        self.layer_norm = nn.LayerNorm(input_dim, eps=1e-6)
        # print(self.layer_norm)


    def forward(self, q, k, v, mask=None):

        batch = q.shape[0]
        #print(batch)

        dim_k, dim_v = self.k, self.v
        len_q, len_k, len_v  = q.size(1), k.size(1), v.size(1)

        
        q_s = self.q(q).view(batch, dim_k)
        k_s = self.k(k).view(batch, dim_k)
        v_s = self.v(v).view(batch, dim_v)
        print(q_s)
        #q, k, v = q.transpose(1,2), k.transpose(1,2), v.transpose(1,2)       
        attn = torch.matmul(q_s/dim_k , k_s.transpose(-1, -2))/np.sqrt(self.d_k)
        

        if mask is not None:
           attn = attn.masked_fill(mask == 0, -1e9)

        attn = self.dropout(F.softmax(attn, dim=-1))
        # print(attn)
        output = torch.matmul(attn, v_s)   
        # print(output)
        return output, attn     

  



class VGG(nn.Module):
  def __init__(self, num_classes=10, attention=False):
    super().__init__()

    self.num_classes = num_classes
    self.attention = attention

    vgg16 = models.vgg16(pretrained=True)
    self.feature_extractor = vgg16.features
    self.avg_pool = vgg16.avgpool
    self.clf = vgg16.classifier
    self.clf[6] = nn.Linear(in_features=4096, out_features=self.num_classes)
    
    if self.attention:
      self.attn_layer = ScaledDotProductAttention(512, 64)


  def forward(self, x):
    x = self.feature_extractor(x)

    if self.attention:
      x = self.attn_layer(x)

    x = self.avg_pool(x)
    x = torch.flatten(x, 1)
    x = self.clf(x)

    return x

1. 当前报错的直接触发原因 #

抛出的TypeError: forward() takes 1 positional argument but 2 were given来自注意力层调用逻辑与forward方法定义不匹配：

• 定义的ScaledDotProductAttention.forward除self参数外，要求传入q、k、v三个必填位置参数
• 但VGG的forward流程中调用注意力层时，仅传入了单个输入x：x = self.attn_layer(x)

PyTorch调用nn.Module实例时会自动将实例自身作为self参数传入forward，此时传入1个x加上self共2个参数，远少于forward要求的参数总数，直接触发参数数量不匹配报错。
如果是实现自注意力逻辑（q/k/v均来自输入特征本身），调用时需要改为x, attn_weight = self.attn_layer(x, x, x)，注意该forward返回output、attn两个值，不能仅用单个变量接收返回结果。

2. 注意力层内部的其他逻辑错误 #

修复调用问题后，代码中还有多处会触发运行错误的逻辑问题，按影响优先级排序：

• dropout层定义缺失：__init__方法中self.dropout的定义被注释，但forward流程中仍在调用self.dropout()，运行时会抛出属性不存在错误，要么取消dropout层的注释，要么删除forward中对应的dropout调用。
• 维度赋值逻辑错误：dim_k, dim_v = self.k, self.v这行代码是将nn.Linear层对象直接赋值给维度变量，并非读取维度数值，后续view、矩阵乘法操作都会触发类型错误。单头注意力场景下直接使用初始化时定义的self.output_dim作为q/k/v的投影维度即可，即dim_k = dim_v = self.output_dim。
• 张量形状转换错误：VGG特征提取层输出的是卷积特征图，形状为(batch_size, 通道数, 特征高, 特征宽)，直接做线性投影后用.view(batch, dim_k)会强制把所有特征压成二维，完全丢失空间结构，注意力计算的形状完全不符合要求。对视觉特征做自注意力时，需要先把特征图从(B, C, H, W)调整为(B, H*W, C)的序列格式，再做q/k/v投影；注意力计算完成后还要把序列形状还原回(B, C, H, W)，才能送入后续的平均池化层。
• 未定义属性引用：forward中用到了self.d_k，但__init__中没有定义该属性，会触发属性不存在错误，缩放点积的分母直接取np.sqrt(self.output_dim)即可。
• 维度不匹配：当前注意力层输出特征维度为64，和后续平均池化、分类头要求的512通道输入不匹配，要么将注意力层的output_dim设置为和input_dim一致的512，要么在注意力层后增加投影层将维度映射回512，否则后续全连接层计算会抛出形状不匹配错误。

• 修正注意力层调用方式，自注意力场景下传入三份相同输入作为q/k/v，同时用两个变量接收返回值
• 补全__init__中被注释的dropout层定义，修正维度变量的赋值逻辑，不要把层对象当成维度值使用
• 增加卷积特征图和序列格式的互相转换逻辑，适配视觉特征的空间结构
• 对齐注意力层输出维度和后续网络层的输入维度要求，避免形状报错

内容的提问来源于stack exchange，提问作者Dew