OverLoCK项目中NATTEN运行时维度不匹配问题的分析与解决-CSDN博客

本文链接：https://blog.csdn.net/gitblog_07595/article/details/148647554

OverLoCK项目中NATTEN运行时维度不匹配问题的分析与解决

OverLoCK [CVPR 2025] OverLoCK: An Overview-first-Look-Closely-next ConvNet with Context-Mixing Dynamic Kernels 项目地址: https://gitcode.com/gh_mirrors/ove/OverLoCK

问题背景

在深度学习领域，注意力机制已成为现代神经网络架构中的重要组成部分。OverLoCK项目作为一种创新的神经网络架构，采用了NATTEN（Neural Attention）库来实现高效的注意力计算。然而，在实际应用过程中，开发者可能会遇到一个典型的运行时错误："Tensor shape mismatch at dimension"，特别是在将OverLoCK架构从编码器适配到解码器时。

问题现象

当尝试在解码器中使用OverLoCK架构进行上采样操作时，系统会抛出维度不匹配的运行时错误。错误信息显示在特定维度上出现了看似相同的数值比较（如"64 != 64"或"4 !=4"），这实际上反映了底层张量形状检查机制检测到的维度不一致问题。

技术分析

错误根源

经过深入分析，问题主要出现在ContMixDynamicConv模块的forward方法中。具体表现为：

注意力头数不一致：在将输入张量x拆分为value时，拆分后的头数与注意力权重张量的头数不匹配
张量重塑问题：在rearrange操作中，对value张量的重塑方式与后续NATTEN操作期望的输入形状不一致

关键代码段分析

在ContMixDynamicConv模块中，value张量通过以下方式生成：

value = rearrange(x, 'b (m g c) h w -> m b g h w c', m=2, g=self.num_heads//2)

而对应的注意力权重张量attn1和attn2的形状为：

attn1.shape: [8, 4, 16, 16, 25]  # 4个头
value[0].shape: [8, 2, 16, 16, 16]  # 只有2个头

这种头数的不匹配导致了NATTEN库在na2d_av函数中执行形状检查时抛出异常。

解决方案

修正方案一：统一注意力头数

最直接的解决方案是确保value张量的头数与注意力权重的头数一致。可以将value的重塑操作修改为：

value = rearrange(x, 'b (m g c) h w -> m b g h w c', m=2, g=self.num_heads)

这样value[0]和value[1]都会有4个头，与attn1和attn2的头数匹配。

修正方案二：调整注意力权重生成

另一种方法是在生成注意力权重时，确保其头数与value张量匹配：

# 修改attn1和attn2的生成逻辑
attn1 = weight[:, :self.num_heads//2, :, :, :effective_smk_size**2]
attn2 = weight[:, self.num_heads//2:, :, :, :effective_kernel_size**2]