PyTorch图像自编码器训练报错：张量维度不匹配求助

错误根源是编码器输出的特征图经解码器还原后，尺寸与原始输入(3,347,400)不匹配，导致MSE损失计算时张量维度冲突。以下是具体的尺寸推导和修正后的模型代码：

编码器尺寸推导（输入：(N,3,347,400)，N为batch大小） #

1. 第一层Conv2d：输出尺寸为(N,16,174,200)
2. 第二层Conv2d：输出尺寸为(N,32,87,100)
3. 第三层Conv2d：输出尺寸为(N,64,81,94)

解码器修正方案 #

解码器需严格对应编码器的下采样逻辑，调整转置卷积的步长、输出填充参数，确保最终输出与输入尺寸完全一致：

class Autoencoder(nn.Module):
    def __init__(self):
        super().__init__()        
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.ReLU(True),
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
            nn.ReLU(True),
            nn.Conv2d(32, 64, 7) 
        )
        
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 32, 7), 
            nn.ReLU(True),
            # 对应编码器第二层的stride=2，添加output_padding=1补全尺寸
            nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1), 
            nn.ReLU(True),
            # 对应编码器第一层的stride=2，output_padding=(0,1)分别匹配高度/宽度的尺寸差
            nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1, output_padding=(0,1)), 
            nn.Sigmoid()
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

修正后尺寸验证 #

解码器每一步输出：

1. 第一层转置卷积：(N,32,87,100)（匹配编码器第二层输出）
2. 第二层转置卷积：(N,16,174,200)（匹配编码器第一层输出）
3. 第三层转置卷积：(N,3,347,400)（与原始输入尺寸完全一致）

此时MSE损失可正常计算，不会再出现维度不匹配的错误。

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