python - 估计器预测无限循环

Question

我不明白如何使用 TensorFlow Estimator API 进行单一预测 - 我的代码导致无限循环，不断预测相同的输入。

根据documentation ，预测应该在 input_fn 引发 StopIteration 异常时停止:

input_fn: Input function returning features which is a dictionary of string feature name to Tensor or SparseTensor. If it returns a tuple, first item is extracted as features. Prediction continues until input_fn raises an end-of-input exception (OutOfRangeError or StopIteration).

这是我代码中的相关部分:

classifier = tf.estimator.Estimator(model_fn=image_classifier, model_dir=output_dir,
                                    config=training_config, params=hparams)

def make_predict_input_fn(filename):
    queue = [ filename ]
    def _input_fn():
        if len(queue) == 0:
            raise StopIteration
        image = model.read_and_preprocess(queue.pop())
        return {'image': image}
    return _input_fn

predictions = classifier.predict(make_predict_input_fn('garden-rose-red-pink-56866.jpeg'))
for i, p in enumerate(predictions):
    print("Prediction %s: %s" % (i + 1, p["class"]))

我错过了什么？

Answer 1

那是因为 input_fn() 需要是一个生成器。将您的函数更改为(yield 而不是 return):

def make_predict_input_fn(filename):
    queue = [ filename ]
    def _input_fn():
        if len(queue) == 0:
            raise StopIteration
        image = model.read_and_preprocess(queue.pop())
        yield {'image': image}
    return _input_fn