PyTorch插值报错：输入输出空间维度不匹配问题求助

报错核心原因 #

F.interpolate要求输入张量必须是**(N, C, d1, d2,...dK)**格式（批量数、通道数、空间维度），但你在插值前将张量压缩成了[256, 1296]（仅批量+展平特征的2D张量），没有空间维度，导致与目标输出的维度要求不匹配。同时你指定的插值目标尺寸(x.shape[0], x.shape[1])也错误——插值不需要改变批量数，只需要恢复空间维度。

修改方案 #

保留张量的空间维度完成插值后，再展平成目标形状：

def forward(self, x):
    print(f'dim of x is {x.shape}')
    src_dims = (x.shape[0], 1, 128, 128)
    z = self.encoder(torch.reshape(x, src_dims))
    out = self.decoder(z[::-1][0], z[::-1][1:])
    out = self.head(out)
    
    # 先执行插值：保留4D格式(N,C,H,W)，目标空间尺寸为原输入的128x128
    if self.retain_dim:
        out = F.interpolate(out, size=(128, 128), mode='bilinear', align_corners=False)
    
    # 展平成与输入一致的2D形状[256, 16384]
    out = torch.flatten(out, start_dim=1)
    print(f'shape of OUT after interpolate is {out.shape}')
    
    z = z[0]
    z = self.head(z)
    z = torch.flatten(z, start_dim=1)
    return out, z

关键修改说明 #

1. 保留空间维度：移除了提前重塑为[256,1,1024]并torch.squeeze的步骤，让out保持[256,1,32,32]的4D格式，符合F.interpolate的输入要求。
2. 修正插值目标尺寸：指定目标空间尺寸为(128,128)，对应原输入的空间分辨率，而不是错误的批量数+特征数。
3. 简洁展平操作：用torch.flatten(start_dim=1)替代手动重塑+压缩，从通道维度后开始展平，直接得到[256,16384]的目标形状。

可选调整 #

• 插值模式mode可根据任务需求修改，比如'nearest'、'bicubic'等；
• align_corners参数在PyTorch 1.10+默认是False，旧版本需手动设置避免警告。

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