how can "Estimater.predict" do a real-time prediction ?

i want to use bert for sentiment classification mission, i fine-tuned bert on a dataset and get an available model, and then, i found it is very slow to predict one sample, someone said the reason is that Estimater.predict will reload graph on each call, which almost spent five seconds!!! god!!

some people say using tf.data.Dataset.from_generator() is a resolution, but i still cant succeed after trying many times

here is part of my code：

    def input_gen(self):
        while True:
            text = "这是一个测试"  # input()
            examples = [InputExample(guid=uuid.uuid4(), text_a=text, text_b=None, label="0")]
            features = BertSentiment.examples_to_features(examples, self.label_list, self.max_seq_length, self.tokenizer)

            all_input_ids = []
            all_input_mask = []
            all_segment_ids = []
            all_label_ids = []

            for feature in features:
                all_input_ids.append(feature.input_ids)
                all_input_mask.append(feature.input_mask)
                all_segment_ids.append(feature.segment_ids)
                all_label_ids.append(feature.label_id)

            num_examples = len(features)

            input_ids = tf.constant(all_input_ids, shape=[num_examples, self.max_seq_length], dtype=tf.int32)
            input_mask = tf.constant(all_input_mask, shape=[num_examples, self.max_seq_length], dtype=tf.int32)
            segment_ids = tf.constant(all_segment_ids, shape=[num_examples, self.max_seq_length], dtype=tf.int32)
            label_ids = tf.constant(all_label_ids, shape=[num_examples], dtype=tf.int32)
            print("here1:{}".format(input_ids.shape))

            yield input_ids, input_mask, segment_ids, label_ids

    def input_fn(self, params):
        """The actual input function."""
        d = tf.data.Dataset.from_generator(self.input_gen, output_types=(tf.int32, tf.int32, tf.int32, tf.int32))
        d = d.batch(batch_size=params["batch_size"], drop_remainder=False)
        iterator = d.make_one_shot_iterator()
        input_ids, input_mask, segment_ids, label_ids = iterator.get_next()
        return {'input_ids': input_ids, 'input_mask': input_mask, 'segment_ids': segment_ids, 'label_ids': label_ids}

    def predict(self):
        result = self.estimator.predict(input_fn=self.input_fn)
        result = list(result)
        print(result)

[sarnikowski]

Here is a simple alternative: Instead of using tf.estimator.predict, you can export your model and use predictor from tensorflow.contrib. This way you avoid loading the graph everytime. It roughly works like this:

First export your model after training:

estimator.export_saved_model(model_dir, serving_input_receiver_fn)

Then when you want to use it for inference reload your model using predictor:

from tensorflow.contrib import predictor
predict_fn = predictor.from_saved_model(model_dir)
result = predict_fn(...)

This should speed your inference up significantly.

[ApexPredator1]

Well~，After a day’s hard work, I found the answer： https://hanxiao.github.io/2019/01/02/Serving-Google-BERT-in-Production-using-Tensorflow-and-ZeroMQ/

Thank you very much Dr. Han Xiao