TensorFlow自定义ImagePreprocessingLayer输入形状报错求助

你犯了一个常见的维度索引错误：Keras模型的输入张量是4维（[batch_size, height, width, channels]），但你在自定义层里把tf.shape(inputs)[0]当成了图像高度——这其实是batch维度的大小（一次输入的样本数量），而非单张图像的高度。

当你单独用该层处理单张样本张量（3维：[height, width, channels]）时，tf.shape(inputs)[0]确实对应图像高度，所以代码能正常运行；但在模型中，输入是带batch的4维张量，逻辑就完全偏离了：

• 当batch_size < 100时，你创建的padding是3维张量（[100 - batch_size, 543, 3]），而输入是4维，直接在axis=0拼接就会出现维度不匹配的错误。
• 就算batch_size > 100，你对整个batch做resize，也会把batch维度改成100，这完全不是你想要的（你要的是把每张图像的高度统一为100）。

要正确处理带batch的4维输入，需要逐个处理batch里的每个样本，可以用tf.map_fn实现。以下是修正后的自定义层代码：

class ImagePreprocessingLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(ImagePreprocessingLayer, self).__init__()
        self.trainable = False
    
    def process_single_image(self, image):
        # 单张图像是3维：[height, width, channels]
        n_rows = tf.shape(image)[0]
        
        if n_rows > 100:
            image = tf.image.resize(image, size=(100, 543))
        elif n_rows < 100:
            padding = tf.zeros(shape=(100 - n_rows, 543, 3), dtype=image.dtype)
            image = tf.concat([image, padding], axis=0)
        
        # 填充NaN值
        image = tf.where(tf.math.is_nan(image), tf.zeros_like(image), image)
        return image
    
    def call(self, inputs):
        # 对batch中的每个样本单独处理，再合并回4维张量
        return tf.map_fn(self.process_single_image, inputs)

• 在tf.image.resize中可以指定method参数（比如tf.image.ResizeMethod.BILINEAR），让下采样的效果更可控。
• 既然输入图像宽度固定为543，你在Input层指定shape=(None, 543, 3)的做法是正确的，这能帮助Keras更好地跟踪张量形状。

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