tf.keras.Model自定义predict_step时如何实现predict参数透传

核心思路是不要将动态推理参数存储为模型实例属性——这类属性会被tf.function编译静态图时捕获为闭包常量，后续修改不会触发图重新编译，也就不会生效。
Keras原生支持将自定义推理参数通过输入结构透传的机制，不需要重写内部缓存逻辑、也不需要开启eager模式。

实现原理 #

Keras的Model.predict()方法会将你传入的输入结构（元组、列表、字典）经过数据适配、拆batch后，原封不动传入每一步的predict_step方法。
我们可以把动态参数和原始输入打包为统一结构传给父类predict，在predict_step中拆包获取参数。此时参数是静态图的正式输入张量，不是外部捕获的闭包变量，每次传入不同值时图可以直接处理，不会出现缓存失效的问题。

代码实现 #

class SimpleModel(tf.keras.Model):
    def call(self, inputs, training=None, mask=None):
        return inputs

    def predict_step(self, data):
        # 按传入predict的结构拆包，拿到原始输入和自定义参数
        x, threshold = data
        return tf.greater(self(x, training=False), threshold)

    def predict(self, x, threshold=0.5, *args, **kwargs):
        # 将参数统一转成张量，避免数据适配层推断问题
        threshold = tf.convert_to_tensor(threshold, dtype=tf.float32)
        # 把自定义参数和输入打包为元组传给父类predict，不修改任何模型实例属性
        return super().predict((x, threshold), *args, **kwargs)


if __name__ == "__main__":
    x = tf.convert_to_tensor([0.0, 0.55, 0.85, 0.9])
    model = SimpleModel()
    pred_1 = model.predict(x, threshold=0.5)
    pred_2 = model.predict(x, threshold=0.75)
    print(pred_1, pred_2, sep="\n")

方案优势 #

• 全程保持静态图推理性能：首次调用后自动缓存计算图，后续不同参数的推理不需要重新编译图，没有临时方案一中反复清空predict_function带来的编译开销
• 无性能损失：不需要开启run_eagerly=True，保留TensorFlow图模式带来的算子融合、内存优化等收益
• 无状态安全问题：自定义参数不再存储在模型实例属性上，多线程/多并发场景下调用不会出现属性被覆盖的竞态问题
• 适配全场景输入：无论输入是numpy数组、Tensor，还是自定义tf.data.Dataset，只要对应调整Dataset的输出结构为(x, threshold)，即可无缝兼容

注意事项 #

• 如果你的模型是多输入结构，只需要按照原有结构把自定义参数附加即可，无需修改原有输入的适配逻辑。例如多输入模型的predict可以写为：

  def predict(self, x, threshold=0.5, *args, **kwargs):
      # x为(input1, input2, input3)的多输入结构
      threshold = tf.convert_to_tensor(threshold, dtype=tf.float32)
      return super().predict((x, threshold), *args, **kwargs)
  

• 如果需要为每个样本设置独立的动态阈值，只需要让threshold和输入x保持相同的第一维长度，Keras会自动按batch拆分传入predict_step。

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