TensorFlow中Unet奇数分辨率问题及可变输入方案咨询

问题根源 #

你的拼接错误是因为编码器池化操作和解码器上采样操作的尺寸不匹配：

• 当输入尺寸为奇数（比如81）时，默认padding='valid'的MaxPooling2D会向下取整计算输出尺寸（81//2=40）；
• 解码器用strides=(2,2)的Conv2DTranspose上采样时，40x40的特征图会被还原为80x80，和编码器中对应的81x81特征图（如drop4、conv3等）尺寸不一致，导致concatenate操作失败。

解决方案（无需调整输入尺寸） #

下面提供一种通用的修改方案，通过自动对齐特征图尺寸实现可变尺寸输入：

步骤1：定义尺寸对齐函数

创建一个Lambda层，自动裁剪或填充上采样后的特征图，使其与编码器对应的特征图尺寸完全匹配：

def align_shape(x):
    encoder_feat, up_feat = x
    h_enc, w_enc = encoder_feat.shape[1], encoder_feat.shape[2]
    h_up, w_up = up_feat.shape[1], up_feat.shape[2]
    
    # 处理高度差异
    h_diff = h_enc - h_up
    if h_diff > 0:
        up_feat = tf.keras.layers.Cropping2D(cropping=((0, h_diff), (0, 0)))(up_feat)
    elif h_diff < 0:
        up_feat = tf.keras.layers.ZeroPadding2D(padding=((0, -h_diff), (0, 0)))(up_feat)
    
    # 处理宽度差异
    w_diff = w_enc - w_up
    if w_diff > 0:
        up_feat = tf.keras.layers.Cropping2D(cropping=((0, 0), (0, w_diff)))(up_feat)
    elif w_diff < 0:
        up_feat = tf.keras.layers.ZeroPadding2D(padding=((0, 0), (0, -w_diff)))(up_feat)
    
    return up_feat

步骤2：修改解码器的拼接逻辑

在每个concatenate操作前，先通过上面的Lambda层对齐上采样特征图与编码器特征图的尺寸，以下是完整修改后的模型代码：

import tensorflow as tf
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Conv2DTranspose, concatenate, Dropout, Lambda, ZeroPadding2D, Cropping2D
from tensorflow.keras.models import Model

def align_shape(x):
    encoder_feat, up_feat = x
    h_enc, w_enc = encoder_feat.shape[1], encoder_feat.shape[2]
    h_up, w_up = up_feat.shape[1], up_feat.shape[2]
    
    h_diff = h_enc - h_up
    if h_diff > 0:
        up_feat = Cropping2D(cropping=((0, h_diff), (0, 0)))(up_feat)
    elif h_diff < 0:
        up_feat = ZeroPadding2D(padding=((0, -h_diff), (0, 0)))(up_feat)
    
    w_diff = w_enc - w_up
    if w_diff > 0:
        up_feat = Cropping2D(cropping=((0, 0), (0, w_diff)))(up_feat)
    elif w_diff < 0:
        up_feat = ZeroPadding2D(padding=((0, 0), (0, -w_diff)))(up_feat)
    
    return up_feat

def build_generator():
    input = Input((None,None,3))
    
    # Encoder path
    conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(input)
    conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1)
    conv2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2)
    conv3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool3)
    conv4 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
    drop4 = tf.keras.layers.Dropout(0.5)(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)

    conv5 = Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool4)
    conv5 = Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5)
    drop5 = tf.keras.layers.Dropout(0.5)(conv5)

    # Decoder path
    up6 = Conv2DTranspose(512, 2, strides=(2, 2), padding='same', kernel_initializer='he_normal')(drop5)
    up6_aligned = Lambda(align_shape)([drop4, up6])
    merge6 = concatenate([drop4, up6_aligned], axis=3)
    conv6 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge6)
    conv6 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)

    up7 = Conv2DTranspose(256, 2, strides=(2, 2), padding='same', kernel_initializer='he_normal')(conv6)
    up7_aligned = Lambda(align_shape)([conv3, up7])
    merge7 = concatenate([conv3, up7_aligned], axis=3)
    conv7 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge7)
    conv7 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7)

    up8 = Conv2DTranspose(128, 2, strides=(2, 2), padding='same', kernel_initializer='he_normal')(conv7)
    up8_aligned = Lambda(align_shape)([conv2, up8])
    merge8 = concatenate([conv2, up8_aligned], axis=3)
    conv8 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge8)
    conv8 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8)

    up9 = Conv2DTranspose(64, 2, strides=(2, 2), padding='same', kernel_initializer='he_normal')(conv8)
    up9_aligned = Lambda(align_shape)([conv1, up9])
    merge9 = concatenate([conv1, up9_aligned], axis=3)
    conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge9)
    conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)

    conv10 = Conv2D(3, 1, activation='sigmoid')(conv9)

    return Model(inputs=input,outputs=conv10,name='generator')

效果说明 #

修改后的模型会自动处理奇数/偶数尺寸的输入，通过裁剪或填充极小的边缘区域（最多1个像素）来保证特征图尺寸匹配，完全不需要调整原始输入图像的尺寸，同时保留Unet的跳跃连接结构。

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