在TensorFlow Keras自定义层的call方法中获取批量大小

核心解决方案 #

使用tf.shape(inputs)获取运行时的动态张量形状，通过索引提取batch_size，避免直接解构赋值符号张量：

class TestClass(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(TestClass, self).__init__(**kwargs)

    def get_config(self):
        config = super(TestClass, self).get_config()
        return config

    def call(self, inputs: tf.Tensor):
        if inputs.dtype.base_dtype != self._compute_dtype_object.base_dtype:
            inputs = tf.cast(inputs, dtype=self._compute_dtype_object)
        # 获取动态形状，避免解构赋值符号张量
        shape = tf.shape(inputs)
        record_count = shape[0]  # 此处为实际运行时的batch_size
        n = shape[1]
        tf.print("Dynamic batch size", record_count)
        return inputs

关键原理说明 #

1. 静态形状 vs 动态形状

   • inputs.shape返回静态形状：模型构建阶段输入层的batch维度默认设为None（支持动态batch），因此静态形状无法反映运行时的真实batch大小。
   • tf.shape(inputs)返回动态形状张量：会在模型运行（如predict、fit）时根据实际输入数据计算真实维度，适合获取动态变化的batch_size。

2. 避免符号张量解构赋值
   TensorFlow图执行模式下，tf.shape(inputs)返回的是符号张量，不支持像普通Python元组那样直接解构赋值（record_count, n = tf.shape(inputs)），必须通过索引（shape[0]）逐个提取维度值，否则会触发OperatorNotAllowedInGraphError。

错误原因解析 #

之前尝试直接解构tf.shape(inputs)时的错误，本质是图执行模式不允许迭代符号张量。添加@tf.function装饰器无法解决问题——因为@tf.function本身就是触发图执行的条件，问题根源在于赋值方式而非执行模式。

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