Estimator predict infinite loop

Question

I don't understand how to make a single prediction using TensorFlow Estimator API - my code results in an endless loop that keeps predicting for the same input.

According to the documentation, the prediction is supposed to stop when input_fn raises a StopIteration exception:

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

Here's the relevant part in my code:

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

What am I missing?

score 0 · Answer 1 · answered Nov 27 '18 at 00:54

That's because input_fn() needs to be a generator. Change your function to (yield instead of 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

score 0 · Answer 2 · answered Sep 19 '19 at 18:12

one solution is to use itertools.islice:

import itertools

predictions = itertools.islice(predictions, number_of_samples)

for i, p in enumerate(predictions):
    print("Prediction %s: %s" % (i + 1, p["class"]))

number_of_samples is an int which is the stop point of the iterator.

Estimator predict infinite loop

2 Answers2