awasthiabhijeet / PIE

Fast + Non-Autoregressive Grammatical Error Correction using BERT. Code and Pre-trained models for paper "Parallel Iterative Edit Models for Local Sequence Transduction": www.aclweb.org/anthology/D19-1435.pdf (EMNLP-IJCNLP 2019)
MIT License
227 stars 40 forks source link

Try to export estimator for more efficient deployment #12

Closed raff7 closed 4 years ago

raff7 commented 4 years ago

Hi, I noticed that using estimator.evaluate is a very inefficient way to use this model on new datapoints. The best method i know of is to export the model with estimator.export_savedmodel

i tried by adding a few flags and this lines of code:

if FLAGS.do_export: estimator._export_to_tpu = False estimator.export_savedmodel(FLAGS.export_dir, serving_input_fn)

where serving_input_fn is

def serving_input_fn(): edit_sequence = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='edit_sequence') input_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_ids') input_mask = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_mask') segment_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='segment_ids') input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({ 'edit_sequence': edit_sequence, 'input_ids': input_ids, 'input_mask': input_mask, 'segment_ids': segment_ids, })() return input_fn

However i was unable to do so. I think it would be a very good addition to your code. but i get an error:ValueError: Couldn't find trained model at PIE_ckpt.

Im very confused by this error.

damosuzuki commented 4 years ago

Hi @raff7, Did you figure out how to export the pre-trained model?

FWIW, I was able to get something like this to work:

        def serving_input_receiver_fn():
            """Serving input_fn that builds features from placeholders

            Returns
            -------
            tf.estimator.export.ServingInputReceiver
            """
            placeholders = {
                "input_sequence": tf.placeholder(dtype=tf.int64, shape=[None, 128], name='input_sequence'),
                "input_mask": tf.placeholder(dtype=tf.int64, shape=[None, 128], name='input_mask'),
                "segment_ids": tf.placeholder(dtype=tf.int64, shape=[None, 128], name='segment_ids'),
                "edit_sequence": tf.placeholder(dtype=tf.int64, shape=[None, 128], name='edit_sequence'),
            }
            return tf.estimator.export.ServingInputReceiver(placeholders, placeholders)
        tf.contrib.tpu.export_estimator_savedmodel(estimator,
            export_path, serving_input_receiver_fn,
            checkpoint_path=predict_ckpt_dir)

This was derived from https://guillaumegenthial.github.io/serving-tensorflow-estimator.html#reload-and-predict-the-good-way

raff7 commented 4 years ago

@damosuzuki Ye at the end i was able without changing my code at all, but doing 1 step of training before exporting the model... I'm not sure why it worked