在 tf.keras.Model 中将参数传递给 model.predict

Question

我有一个模型需要自定义推理，因此我修改了 tf.keras.Model class 的 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")

更新二： 在我发布关于 predict_step 函数运行ning 在图形模式下的问题之后，似乎另一种解决问题的方法是设置模型的 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 的情况下工作（可能运行由于急切模式而变慢）。

Answer 1

由于不允许我发表评论，所以我只能直接提供一个可能非常离谱的答案。

我对你的问题的理解与其说是如何多次调用 predict_step()，不如说是如何使阈值可变。我的建议是使 self.threshold 成为 [1 x 1] 不可训练的变量。

我正在考虑添加类似

的内容

threshold = tf.Variable(.65,trainable=False, dtype='float32')

我认为你可以如何实现它如下，

class SimpleModel(tf.keras.Model):
    def __init__(self,threshold=.5):
        super().__init__()
        self.threshold = tf.Variable(threshold,trainable=False, dtype='float32')

    def call(self, inputs, training=None, mask=None):
        return inputs

    def predict(self, x, threshold=0.5, *args, **kwargs):
        self.threshold.assign(tf.convert_to_tensor(threshold,dtype='float32'))
        self.predict_function = None
        return super().predict(x, *args, **kwargs)

    def predict_step(self, data):
        out = tf.greater(self(data, training=False), self.threshold)
        # self.threshold.assign( <Calculate new threshold here as a float32 tensor>)
        return out


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")

事情发生了变化：

threshold __init__() 中的参数（可能没有必要）
- 如果不希望出现上述情况，则必须将变量的实例化从 __init__() 移动到 predict()
在 __init__()

tf.variable

self.threshold

不再立即返回 predict_step()
在 predict_step

此代码已编译，但因为我是在解决我自己对您的查询的解释，所以我可能与您正在寻找的内容相去甚远。

Answer 2

我对你想要的东西有了更好的了解。查看这个玩具示例，看看它是否是您想要的。

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
    
    # fancy way of saying predict = custom_predict(predict)
    # really, it's running custom_predict masquerading as 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")

当然，当您本可以在自己的预测函数中处理逻辑时，装饰器就太过分了，但也许更高层次的想法会让您自己的想法流向您想要处理的方式。可定制性的另一种选择是使用回调（例如，参见 fastai 或 Pytorch Lightning）。

在 tf.keras.Model 中将参数传递给 model.predict

Passing parameters to model.predict in tf.keras.Model

tensorflow

machine-learning

keras

tensorflow2.0

tf.keras