PyTorch将卷积层输出传入FCC层的维度匹配及可复现性问题

维度不匹配报错修复 #

你遇到的mat1 dim 1 must match mat2 dim 0报错核心由两个问题导致：

• 替换AdaptiveAvgPool2d为自定义Identity层后，输出的是维度为[10, 256, 9, 9]的4维张量，全连接层需要最后一维的维度和in_features参数匹配，你没有对卷积输出做展平操作，最后一维仅为9，和你修改后的20736完全不匹配
• 需同步确保全连接层第一层的in_features参数确实修改为20736

对应修复代码示例：

# 替换原有自适应池化层
model.avgpool = Identity()
# 重构分类头，新增展平层，调整全连接层输入维度
model.classifier = Sequential(
    nn.Flatten(start_dim=1), # 将4维张量展平为[batch_size, 256*9*9]的2维张量
    nn.Dropout(p=0.5, inplace=False),
    nn.Linear(in_features=20736, out_features=4096, bias=True)
    # 其余分类层保持原有配置即可
)

确定性可复现需求的替代方案 #

如果需要保留池化操作提升模型泛化性，可以用支持torch.use_deterministic_algorithms(True)配置的普通AvgPool2d替换AdaptiveAvgPool2d，不需要直接移除池化层：
你当前最后一层池化输出为[10, 256, 9, 9]，如果要和原有自适应池化逻辑对齐得到[6,6]的输出，可按如下配置设置普通平均池化参数：

model.avgpool = nn.AvgPool2d(kernel_size=4, stride=1)
# 输出维度为[10, 256, 6, 6]，和原有AdaptiveAvgPool2d输出完全一致

此时全连接层可以保留原有in_features=9216的配置，不需要修改。

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