PyTorch实现UNet时torch.cat报张量尺寸不匹配RuntimeError如何解决

报错核心原因 #

报错出在上采样特征与下采样对应特征做通道维拼接时，二者的高、宽维度不匹配（报错显示维度值为15、10）。根源是上采样模块逻辑错误、上下采样尺寸计算不对齐、数据预处理维度处理偏差三个问题叠加。

具体问题与修复步骤 #

1. 上采样模块逻辑颠倒（核心问题） #

当前expand_block的结构和调用逻辑完全不符合UNet常规设计：

• 常规UNet上采样路径逻辑为：先对低分辨率特征做转置卷积上采样→和同分辨率下采样特征拼接→再做两次3×3卷积提取特征
• 你当前的expand_block是先做两次3×3卷积、最后才接转置卷积上采样，且forward调用时在拼接前就传入expand_block，还多传了冗余参数1

修复方法：

• 重写expand_block结构，把转置卷积放在模块最前面，后续卷积层适配拼接后的通道数：

def expand_block(self, in_channels, out_channels, kernel_size, padding):
    expand = nn.Module()
    # 转置卷积上采样层：通道数从in_channels压缩为out_channels
    expand.trans_conv = nn.ConvTranspose2d(
        in_channels, out_channels, 
        kernel_size=3, stride=2, padding=1, output_padding=1
    )
    # 拼接后通道数翻倍，因此两次卷积的输入通道设为out_channels*2
    expand.conv_layers = nn.Sequential(
        nn.Conv2d(out_channels*2, out_channels, kernel_size=kernel_size, stride=1, padding=padding),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    )
    # 实现上采样→拼接→卷积的完整流程
    def forward(x, skip_feat):
        up_feat = expand.trans_conv(x)
        # 强制对齐尺寸，避免微小误差导致拼接失败
        up_feat = torch.nn.functional.interpolate(
            up_feat, size=skip_feat.shape[2:], mode='bilinear', align_corners=True
        )
        cat_feat = torch.cat([up_feat, skip_feat], dim=1)
        return expand.conv_layers(cat_feat)
    expand.forward = forward
    return expand

• 修改forward的上采样调用逻辑，去掉冗余传参：

def forward(self, x):
    # 下采样路径保持原有逻辑
    conv1 = self.downconv1(x)
    conv2 = self.downconv2(conv1)
    conv3 = self.downconv3(conv2)
    conv4 = self.downconv4(conv3)
    conv5 = self.downconv5(conv4)

    # 上采样路径按新的expand逻辑调用
    upconv1 = self.upconv1(conv5, conv4)
    upconv2 = self.upconv2(upconv1, conv3)
    upconv3 = self.upconv3(upconv2, conv2)
    upconv4 = self.upconv4(upconv3, conv1)
    upconv5 = self.upconv5(upconv4)
    self.out = upconv5

2. 尺寸不对齐补全处理 #

即便参数正确，奇数尺寸的特征图在池化、转置卷积时也可能出现1-2像素的尺寸误差，上述重写的expand逻辑里已经通过interpolate强制对齐了上采样和下采样的特征图高宽，从根源避免尺寸不匹配的报错，不需要手动计算每层的尺寸公式。

3. 数据预处理维度修正 #

当前预处理和训练循环的维度处理存在错误，导致模型输入形状不符合要求：

• 灰度图经numpy转张量后是(H, W)的2维形状，缺少通道维度
• dataloader输出的inputs已经自带batch维度（形状为(batch_size, H, W)），你在训练循环里加inputs[None,...]会多套一层冗余维度，把通道维挤成了batch维

修复方法：

• 预处理代码里给张量加通道维度，替换废弃的Scale接口：

def getImageAndTransform(self,item):
    data = Image.open(self.datalist[item])
    label = Image.open(self.labellist[item])
    p = torchvision.transforms.Compose([torchvision.transforms.Resize((256,256))])
    data = p(data)
    label = p(label)
    # 加通道维度，变成(1, H, W)
    data = torch.from_numpy(np.array(data)).unsqueeze(0)
    label = torch.from_numpy(np.array(label)).unsqueeze(0)
    return data, label

• 训练循环里去掉冗余的None维度，直接输入模型：

outputs = model(inputs.float())

调试验证方法 #

修复后先构造假输入跑通前向传播，不用加载全量数据：

# 构造batch=1、单通道、256x256的假输入，形状符合模型要求
x = torch.randn(1, 1, 256, 256)
# 打印前向传播过程中每一层的输出形状，定位残留的尺寸问题
y = model(x)
print(y.shape)

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