我有一个模型需要自定义推理，因此我修改了 tf.keras.Model 类的 predict_step 方法。我希望推理能够根据某些参数进行修改，有没有简单的方法让 predict 方法接收参数并传递给 predict_step 函数？
类似这样：
class SimpleModel(tf.keras.Model): def __init__(self): super().__init__() self.threshold = None def call(self, inputs, training=None, mask=None): return inputs def predict(self, x, threshold=0.5, *args, **kwargs): self.threshold = threshold return super().predict(x, *args, **kwargs) def predict_step(self, data): return tf.greater(self(data, training=False), self.threshold)if __name__ == "__main__": x = tf.convert_to_tensor([0.0, 0.55, 0.85, 0.9]) model = SimpleModel() model.predict(x, threshold=0.5) model.predict(x, threshold=0.75)这种方法的问题在于，由于 predict_step 已经被创建，阈值不会改变。
更新 1：
这似乎有效，但不确定是否是最佳方法：
class SimpleModel(tf.keras.Model): def __init__(self): super().__init__() self.threshold = None def call(self, inputs, training=None, mask=None): return inputs def predict(self, x, threshold=0.5, *args, **kwargs): self.threshold = threshold self.predict_function = None return super().predict(x, *args, **kwargs) def predict_step(self, data): return tf.greater(self(data, training=False), self.threshold)if __name__ == "__main__": x = tf.convert_to_tensor([0.0, 0.55, 0.85, 0.9]) model = SimpleModel() pred = model(x) pred_1 = model.predict(x, threshold=0.5) pred_2 = model.predict(x, threshold=0.75) print(pred, pred_1, pred_2, sep="\n")更新 2：在这里我发布的问题中提到 predict_step 函数在图模式下运行，似乎解决这个问题的方法是设置模型的 self.run_eagerly = True。
class SimpleModel(tf.keras.Model): def __init__(self): super().__init__() self.run_eagerly = True self.threshold = None def call(self, inputs, training=None, mask=None): return inputs def predict(self, x, threshold=0.5, *args, **kwargs): self.threshold = threshold return super().predict(x, *args, **kwargs) def predict_step(self, data): return tf.greater(self(data, training=False), self.threshold)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")现在它在不使用 tf.Variable 的情况下也能工作（可能会因为启用了即时执行模式而运行得更慢）。
回答：
我对你想要的有了更好的理解。看看这个玩具示例，看看是否符合你的需求。
class SimpleModel(tf.keras.Model): def __init__(self): super().__init__() def call(self, inputs, training=None, mask=None): return inputs def custom_predict(func): def threshold_handler(self, x, threshold=None, *args, **kwargs): if threshold is None: return func(self, x, *args, **kwargs) else: vals = func(self, x, *args, **kwargs) return list(filter(lambda x: x > threshold, vals)) return threshold_handler # 这是一种花哨的说法，意思是 predict = custom_predict(predict) # 实际上，它是在伪装成 predict 的 custom_predict @custom_predict def predict(self, x, *args, **kwargs): return super().predict(x, *args, **kwargs)x = tf.convert_to_tensor([0.0, 0.55, 0.85, 0.9])model = SimpleModel()pred = model(x)pred_0 = model.predict(x, steps=1)pred_1 = model.predict(x, threshold=0.5, steps=1)pred_2 = model.predict(x, threshold=0.75, steps=1)print(pred, pred_0, pred_1, pred_2, sep="\n")当然，使用装饰器在你可以在自己的 predict 函数中处理逻辑时是完全多余的，但也许更高层次的想法会激发你自己如何处理的灵感。另一个自定义的选项是使用回调（例如，参考 fastai 或 Pytorch Lightning）。
