PyTorch-Image-Models中MultiQueryAttention2d模块的Upsample问题解析

在分析PyTorch-Image-Models项目中的MobileNet v4实现时，我们发现其MultiQueryAttention2d模块存在一个值得注意的实现细节问题，特别是在处理query_strides大于1的情况时。

问题本质

MultiQueryAttention2d模块中的上采样操作使用了不正确的参数传递方式。原始代码将query_strides直接作为size参数传递给nn.Upsample，而实际上应该使用scale_factor参数。这个差异会导致上采样行为与预期不符。

技术细节

在PyTorch中，nn.Upsample有两个关键参数：

• size：指定输出的确切尺寸
• scale_factor：指定相对于输入尺寸的缩放比例

在注意力机制的上下文中，我们通常希望按比例放大特征图，因此scale_factor才是正确的选择。使用size参数会导致输出尺寸被固定为query_strides值，而不是按比例放大。

影响范围

这个问题主要影响以下场景：

1. 当query_strides参数大于1时
2. 在需要按比例放大特征图的注意力计算中
3. 在构建自定义模型时使用这个模块的stride功能

值得注意的是，在MobileNet v4的默认配置中，这个问题不会显现，因为当前实现只使用了kv_stride而没有使用query_strides功能。

解决方案

正确的实现应该将代码修改为使用scale_factor参数：

nn.Upsample(scale_factor=self.query_strides, mode='bilinear', align_corners=False)

相关修复

在修复过程中还发现了一个配套问题：平均池化操作添加了额外的padding，导致尺寸不匹配。这些问题共同影响了模块在stride模式下的正确行为。

对开发者的建议

1. 在使用自定义stride参数时，务必验证特征图的尺寸变化
2. 对于注意力机制中的上采样操作，明确区分size和scale_factor的使用场景
3. 在修改类似核心模块时，建议构建测试用例验证各种stride组合下的行为

这个问题提醒我们，在实现复杂的注意力机制时，尺寸变换相关的操作需要特别小心，确保各阶段的特征图尺寸符合预期。