Vision Transformer替换SelfAttention为ImageLinearAttention报错如何解决

你用到的ImageLinearAttention原本是为4维图像张量（格式为[batch, channel, height, width]）设计的，内部实现用了nn.Conv2d层，对应报错里提到的4维权重参数。而ViT中输入注意力模块的是3维序列张量，你输入的单批次无batch维度数据为[1983, 512]，加上batch维度、拼接class token后就变成了报错提到的[1, 1984, 512]格式，和卷积层要求的4维输入不匹配，即使删除了通道数相关逻辑也没有解决底层算子的输入维度要求问题。

报错信息：RuntimeError: Expected 4-dimensional input for 4-dimensional weight [384, 512, 1, 1], but got 3-dimensional input of size [1, 1984, 512] instead

• 方案1：适配3维序列输入，修改ImageLinearAttention内部算子
  将代码中所有的nn.Conv2d(in_channels, out_channels, kernel_size=1, ...)替换为适配3维输入的算子：
  • 可选替换为nn.Conv1d(in_channels, out_channels, kernel_size=1, ...)，需要注意调整前向传播时的维度顺序为[batch, embed_dim, seq_len]再输入卷积层，运算完成后再调换回原来的序列维度顺序
  • 也可以直接替换为nn.Linear(in_channels, out_channels, ...)，不需要调整维度顺序，直接适配[batch, seq_len, embed_dim]格式的输入
• 方案2：还原4维输入格式，适配原有ImageLinearAttention逻辑
  先将ViT的3维序列张量转换为4维特征图再输入注意力模块：
  1. 拆分输入序列的class token和patch token：cls_token, patch_tokens = x[:, :1], x[:, 1:]
  2. 根据ViT切分patch时的图像高宽，将patch token reshape为4维格式：patch_tokens = patch_tokens.reshape(batch_size, H, W, embed_dim).permute(0, 3, 1, 2)，其中H * W等于patch token的数量1983
  3. 将4维张量输入ImageLinearAttention得到输出，再将输出flatten为序列格式，和class token拼接回原有3维序列格式即可

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