如何在Keras自定义层中处理输入形状含None时的输出尺寸计算？

核心逻辑 #

在Keras编写自定义层时，不管输入的Batch、Height、Width是不是None，只要对确定的维度做整除2操作，不确定的维度直接保留None即可。最规范的方式是重写compute_output_shape方法，也可以在call方法里结合静态/动态形状处理。

具体实现 #

方式1：重写compute_output_shape（推荐） #

这个方法专门用来定义输入到输出的形状映射，逻辑清晰直观：

from tensorflow.keras.layers import Layer

class CustomDownsampleLayer(Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    def call(self, inputs):
        # 这里写你的下采样逻辑，比如步长为2的卷积、平均池化，或者直接切片
        return inputs[:, ::2, ::2, :]
    
    def compute_output_shape(self, input_shape):
        batch, h, w, c = input_shape
        # 对非None的维度做整除2，None维度直接保留
        out_h = h // 2 if h is not None else None
        out_w = w // 2 if w is not None else None
        return (batch, out_h, out_w, c)

方式2：在call里动态处理（TF2.x适用） #

TF2.x中可以直接利用张量的形状属性动态计算，还能给输出设置静态形状提示，帮助Keras做形状推断：

import tensorflow as tf
from tensorflow.keras.layers import Layer

class CustomDownsampleLayer(Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    def call(self, inputs):
        # 获取运行时的动态形状
        out_h = tf.shape(inputs)[1] // 2
        out_w = tf.shape(inputs)[2] // 2
        
        # 执行下采样操作
        output = inputs[:, ::2, ::2, :]
        
        # 设置静态形状提示，方便模型构建阶段的形状推断
        output.set_shape((
            inputs.shape[0],
            inputs.shape[1] // 2 if inputs.shape[1] is not None else None,
            inputs.shape[2] // 2 if inputs.shape[2] is not None else None,
            inputs.shape[3]
        ))
        return output

注意事项 #

• 如果输入的高/宽是奇数，//2会向下取整（比如5→2），要确保层的实际操作和这个形状计算逻辑匹配，比如步长为2的卷积对奇数尺寸的处理逻辑也是一致的。
• 动态维度（None）直接保留即可，Keras会自动处理动态形状的传递，完全不影响模型适配不同尺寸的输入。

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