Unet模型Padding尺寸不匹配求助：RuntimeError张量尺寸不符

问题描述 #

构建Siamese UNet模型时，解码器第三层出现张量尺寸不匹配错误。已在UpBlock类中尝试用padding方法对齐尺寸，前两个张量尺寸正常（torch.Size([2, 512, 139, 139])、torch.Size([2, 256, 267, 267])），但后续触发报错：

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 519 but got size 518 for tensor number 1 in the list

相关代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=2),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=2),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, out):
        x = self.double_conv(out)
        return x

class DownBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DownBlock, self).__init__()
        self.double_conv = DoubleConv(in_channels, out_channels)
        self.down_sample = nn.MaxPool2d(2)

    def forward(self, x):
        skip_out = self.double_conv(x)
        down_out = self.down_sample(skip_out)
        return (down_out, skip_out)

def pad_tensor(source, target):
    """
    Pad source tensor to match target tensor size
    :param source: tensor that need to get padding
    :param target: tensor of the desired shape
    :return: source tensor with shape equal to target
    """
    diff_y = target.size()[2] - source.size()[2]
    diff_x = target.size()[3] - source.size()[3]
    source = F.pad(source, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2])
    return source

class UpBlock(nn.Module):
    """Upscaling then double conv"""
    def __init__(self, in_channels, out_channels, up_sample_mode):
        super(UpBlock, self).__init__()
        if up_sample_mode == 'conv_transpose':
            self.up_sample = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2,  padding=1, output_padding=1)        
        elif up_sample_mode == 'bilinear':
            self.up_sample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            raise ValueError("Unsupported `up_sample_mode` (can take one of `conv_transpose` or `bilinear`)")
        self.double_conv = DoubleConv(in_channels, out_channels)

    def forward(self, down_input, skip_input):
        down_input = self.up_sample(down_input)
        down_input = pad_tensor(down_input, skip_input)
        x = torch.cat([down_input, skip_input], dim=1)
        return self.double_conv(x)

class Siamese_UNet_Assement(nn.Module):
    def __init__(self, out_classes, up_sample_mode='conv_transpose'):
        super(Siamese_UNet_Assement, self).__init__()
        self.up_sample_mode = up_sample_mode
        # Downsampling Path
        self.down_conv1 = DownBlock(3, 64)
        self.down_conv2 = DownBlock(64, 128)
        self.down_conv3 = DownBlock(128, 256)
        self.down_conv4 = DownBlock(256, 512)
        # Bottleneck
        self.double_conv = DoubleConv(512, 1024)
        # Upsampling Path
        self.up_conv4 = UpBlock(1024, 512, self.up_sample_mode)
        self.up_conv3 = UpBlock(512, 256,self.up_sample_mode)
        self.up_conv2 = UpBlock(256,128,self.up_sample_mode)
        self.up_conv1 = UpBlock(128 , 64, self.up_sample_mode)
        # Final Convolution
        self.conv_last = nn.Conv2d(64, out_classes, kernel_size=1)

    def forward(self, input1,input2):
        # Unet1
        # Encoder
        enc1_1, skip1_out1 = self.down_conv1(input1)
        print(skip1_out1.size())
        enc1_2, skip2_out1 = self.down_conv2(enc1_1)
        print(skip2_out1.size())
        enc1_3, skip3_out1 = self.down_conv3(enc1_2)
        print(skip3_out1.size())
        enc1_4, skip4_out1 = self.down_conv4(enc1_3)
        print(skip4_out1.size())

        bottleneck_1  = self.double_conv(enc1_4)
        print(bottleneck_1.size())

        #Decoder
        dec4_1= self.up_conv4(bottleneck_1, skip4_out1)
        print(dec4_1.size())
        dec3_1= self.up_conv3(dec4_1, skip3_out1)
        print(dec3_1.size())
        dec2_1= self.up_conv2(dec3_1, skip2_out1)
        print(dec2_1.size())
        dec1_1 = self.up_conv1(dec2_1, skip1_out1)
        print(dec1_1.size())
        dec1_1 = self.conv_last(dec1_1)
        print(dec1_1.size())
        
        # Unet2
        # Encoder
        enc2_1, skip1_out2 = self.down_conv1(input2)
        enc2_2, skip2_out2 = self.down_conv2(enc2_1)
        enc2_3, skip3_out2 = self.down_conv3(enc2_2)
        enc2_4, skip4_out2 = self.down_conv4(enc2_3)
        bottleneck_2  = self.double_conv(enc2_4)
        
        #Decoder
        dec4_2= self.up_conv4(bottleneck_2 , skip4_out2)
        dec3_2_= self.up_conv3(dec4_2, skip3_out2)
        dec2_2= self.up_conv2(dec3_2_, skip2_out2)
        dec1_2 = self.up_conv1(dec2_2, skip1_out2)
        dec1_2 = self.conv_last(dec1_2)
        
        # Siamese
        output = torch.cat((dec1_1,dec1_2), 1)
        return output

问题原因 #

1. DoubleConv的padding设置错误：使用padding=2时，每次3x3卷积会让特征图尺寸增加2（公式：H_out = H_in - kernel_size + 2*padding + 1），叠加MaxPool2d的下采样后，编码器各层skip tensor的尺寸计算出现偏差，导致后续上采样后的张量与skip tensor的尺寸差为奇数。
2. ConvTranspose2d参数不合理：padding=1和output_padding=1的设置破坏了2倍上采样的尺寸逻辑，进一步放大了尺寸偏差。
3. pad_tensor的奇数差处理不足：当尺寸差为奇数时，diff_x // 2和diff_x - diff_x//2的总和虽然等于diff_x，但拼接前的DoubleConv会再次改变尺寸，最终引发不匹配。

解决方案 #

1. 修正DoubleConv的padding

将DoubleConv中的padding=2改为padding=1，确保卷积后特征图尺寸与输入一致（3x3卷积+padding=1时，H_out = H_in），从根源上保证编码器各层尺寸计算准确：

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, out):
        x = self.double_conv(out)
        return x

2. 修正ConvTranspose2d的参数

调整UpBlock中ConvTranspose2d的参数，确保上采样后尺寸为输入的2倍：

class UpBlock(nn.Module):
    """Upscaling then double conv"""
    def __init__(self, in_channels, out_channels, up_sample_mode):
        super(UpBlock, self).__init__()
        if up_sample_mode == 'conv_transpose':
            # 修正参数：padding=0，output_padding=0
            self.up_sample = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2, padding=0, output_padding=0)        
        elif up_sample_mode == 'bilinear':
            self.up_sample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            raise ValueError("Unsupported `up_sample_mode` (can take one of `conv_transpose` or `bilinear`)")
        self.double_conv = DoubleConv(in_channels, out_channels)

    def forward(self, down_input, skip_input):
        down_input = self.up_sample(down_input)
        down_input = pad_tensor(down_input, skip_input)
        x = torch.cat([down_input, skip_input], dim=1)
        return self.double_conv(x)

3. 优化pad_tensor处理奇数偏差（可选）

如果仍存在尺寸差为奇数的情况，修改pad_tensor确保padding后尺寸完全匹配：

def pad_tensor(source, target):
    """
    Pad source tensor to match target tensor size
    :param source: tensor that need to get padding
    :param target: tensor of the desired shape
    :return: source tensor with shape equal to target
    """
    diff_y = target.size()[2] - source.size()[2]
    diff_x = target.size()[3] - source.size()[3]
    # 处理奇数偏差，确保左右、上下的padding总和等于diff值
    pad_left = diff_x // 2
    pad_right = diff_x - pad_left
    pad_top = diff_y // 2
    pad_bottom = diff_y - pad_top
    source = F.pad(source, [pad_left, pad_right, pad_top, pad_bottom])
    return source

验证 #

修改后，编码器各层的skip tensor尺寸与解码器上采样后的尺寸会完全匹配，拼接操作不会再触发尺寸不匹配错误。

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