PyTorch将输入图像尺寸改为64×64×3时出现张量形状不匹配报错的问题

原网络是完全针对32×32×3输入尺寸设计的，修改输入分辨率后，卷积、池化层输出的特征图尺寸发生变化，但没有同步修改全连接层的输入维度、拉平张量时的手动计算值，才连续触发维度不匹配的RuntimeError。

两个报错的对应问题：

• 第一次报错：手动写死拉平维度为10817，但实际特征张量的总元素数刚好是10816，多算了1个维度
• 第二次报错：把拉平维度改成了800，但fc1层的初始化输入维度还是原网络的400，线性层计算时两个矩阵形状无法相乘

按照现有网络的层设置，输入64×64×3图像时，逐步骤计算特征图尺寸：

• 第一层卷积Conv2d(3,6,5)（无padding、步长1）：输出尺寸 = 输入尺寸 - 卷积核尺寸 +1 = 64-5+1=60，对应特征图形状6×60×60
• 第一次最大池化MaxPool2d(2,2)：输出尺寸 = 特征图尺寸 //2 =60//2=30，对应特征图形状6×30×30
• 第二层卷积Conv2d(6,16,5)：输出尺寸 =30-5+1=26，对应特征图形状16×26×26
• 第二次最大池化：输出尺寸=26//2=13，对应最终特征图形状16×13×13

• 最终拉平后的特征维度为 16*13*13 =2704，这就是全连接层需要接收的输入值

1. 修正transform的顺序（推荐） #

原代码把ToTensor放在最前面会导致Resize/RandomCrop作用在张量上，兼容性不如先做图像变换再转张量：

transform = transforms.Compose(
    [transforms.Resize(64),
     transforms.RandomCrop(64),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

2. 修正网络结构 #

把原来写死的旧维度替换成计算得到的2704，同时注意当前任务是猫狗二分类，原来的fc3输出10类不符合任务需求，对应调整为2：

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 原输入维度16*5*5替换为计算得到的16*13*13
        self.fc1 = nn.Linear(16 * 13 * 13, 120)
        self.fc2 = nn.Linear(120, 84)
        # 原输出10类替换为2类（猫/狗）
        self.fc3 = nn.Linear(84, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        # 用自动拉平代替手动写死维度的x.view，避免后续调整参数时算错维度
        x = x.flatten(1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

• 不要无依据给拉平维度加固定数值（比如之前写的10816+1），所有维度必须和卷积输出的特征图大小严格对应
• 不要只修改x.view的参数、不同步修改全连接层的初始化输入维度，否则还是会触发矩阵相乘维度不匹配的报错

内容的提问来源于stack exchange，提问作者BIT-REAPER