基于TensorLy的ResNet塔克分解维度不匹配RuntimeError求助

问题根源 #

你使用的ResNet仓库中，Lambda层负责shortcut分支的维度对齐（下采样+通道补零），当Partial Tucker分解后的主分支卷积层输出的特征图空间维度与原层不一致时，就会和shortcut分支的输出维度冲突，触发维度不匹配的RuntimeError。VBMF自动秩估计可能会调整空间维度的秩，导致分解后的层输出尺寸偏离原层。

具体解决步骤 #

1. 分解时强制对齐原层的输出维度 #

调用TensorLy的partial_tucker函数时，不要完全依赖VBMF的自动秩估计，手动指定空间维度的秩，确保分解后的卷积层输出的特征图尺寸与原卷积层完全一致：

• 先获取原卷积层的输出尺寸：用dummy输入过原层，得到输出的(H_out, W_out)。
• 分解时将塔克分解的空间模式秩设置为与H_out、W_out匹配的数值，保证输出尺寸不变。
• 示例代码片段：

  import tensorly as tl
  from tensorly.decomposition import partial_tucker
  import torch
  
  # 原卷积层
  original_conv = model.layer1[0].conv1
  # 获取原层输出尺寸（CIFAR10输入为32x32）
  dummy_input = torch.randn(1, original_conv.in_channels, 32, 32)
  original_output = original_conv(dummy_input)
  target_H, target_W = original_output.shape[2], original_output.shape[3]
  
  # 手动指定秩，确保空间维度匹配原层输出
  weights = original_conv.weight.data
  rank = [original_conv.in_channels, target_H, target_W, original_conv.out_channels]
  core, factors = partial_tucker(weights, modes=[0, 1, 2, 3], rank=rank, init='svd')
  

2. 重构分解层时严格继承原层的卷积参数 #

重构分解后的子层时，完全复用原卷积层的stride、padding、dilation、groups参数，避免因这些参数不一致导致输出尺寸变化：

• 比如原卷积层是nn.Conv2d(in_c, out_c, kernel_size=3, stride=2, padding=1)，分解后的第一个卷积层要设置相同的stride和padding，后续层保持stride=1、padding=0，同时通过调整核尺寸保证整体输出和原层一致。

3. 修正VBMF秩估计的空间维度约束 #

如果不想完全手动指定秩，可以修改VBMF的秩估计逻辑，限制空间维度的秩不能改变原层的输出尺寸：

• 先用VBMF估计通道模式的秩，再将空间模式的秩替换为原层输出的空间维度值，再进行分解。
• 示例代码片段：

  from tensorly.decomposition import partial_tucker
  from tensorly.contrib.sparse.decomposition import vbmf
  
  weights = original_conv.weight.data
  # 仅估计通道模式的秩，空间模式保留原尺寸
  estimated_rank = vbmf(weights, modes=[0, 3])
  # 修正空间模式的秩
  full_rank = [estimated_rank[0], target_H, target_W, estimated_rank[1]]
  core, factors = partial_tucker(weights, modes=[0, 1, 2, 3], rank=full_rank, init='svd')
  

4. 临时适配FLOPS计算的维度（不修改原模型架构） #

如果上述方法仍有问题，可以在计算FLOPS和参数时，给分解后的主分支添加一个临时的维度调整层，仅在计算时生效，不影响模型的实际训练和推理：

• 示例代码片段：

  from thop import profile
  import torch.nn as nn
  
  def wrap_model(model, target_size):
      class WrappedModel(nn.Module):
          def __init__(self, model):
              super().__init__()
              self.model = model
              self.adapt = nn.AdaptiveAvgPool2d(target_size)
  
          def forward(self, x):
              out = self.model(x)
              if out.shape[2:] != target_size:
                  out = self.adapt(out)
              return out
      return WrappedModel(model)
  
  # 计算FLOPS时使用包装后的模型
  flops, params = profile(wrap_model(decomposed_model, (target_H, target_W)), inputs=(dummy_input,))
  

注意事项 #

• 分解后的模型必须保证每一层的输出形状与原模型完全一致，这样shortcut的Lambda层才能正常工作，无需修改原架构。
• 优先通过调整塔克分解的秩和重构层的参数来对齐维度，临时适配层仅作为兜底方案，避免影响模型性能对比的准确性。

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