Export and Import Models with Tensorflow SavedModelBuilder: an LSTM Example

NOTE: (7/23/2018) I’m primarily a PyTorch dev and am new to TensorFlow, and this is my first attempt to get it working. I will update this post to reflect changes in my understanding of the framework.

While the official docs favor SavedModelBuilder over checkpoint tf.train.Saver and tf.saved_model.simple_save, detailed explanations and examples are sparse.

My project is to translate Russian into IPA (International Phonetic Alphabet) which is basically an RNN application using sequence-to-sequence (seq2seq) LSTM. I was able to cleanly export a model, import it, and use it in code.

1. Export

My first approach failed because it did not export any variable. I was using the official docs approach:

# Failing Code
# After loading checkpoint separately
import time
import os
export_dir = os.path.join('models', time.strftime("%Y%m%d-%H%M%S"))
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
with tf.Session(graph=tf.Graph()) as sess:
    builder.add_meta_graph_and_variables(sess,     [tf.saved_model.tag_constants.TRAINING.TRAINING],
                                         strip_default_attrs=True)
# Add a second MetaGraphDef for inference.
with tf.Session(graph=tf.Graph()) as sess:
    builder.add_meta_graph([tag_constants.SERVING],
    strip_default_attrs=True)
builder.save()

I suspected that the reason was that the session was not populated with any actual graph, so I loaded a checkpoint into sess:

# Minimal code to successfully export tf model
import time
import os
import tensorflow as tf
trained_checkpoint_prefix = 'checkpoints/dev'
export_dir = os.path.join('models', time.strftime("%Y%m%d-%H%M%S"))
loaded_graph = tf.Graph()
with tf.Session(graph=loaded_graph) as sess:
    # Restore from checkpoint
    loader = tf.train.import_meta_graph(trained_checkpoint_prefix + '.meta')
    loader.restore(sess, trained_checkpoint_prefix)
    
    # Export checkpoint to SavedModel
    builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
    builder.add_meta_graph_and_variables(sess,
                                         [tf.saved_model.tag_constants.TRAINING],
                                         strip_default_attrs=True)
builder.add_meta_graph([tf.saved_model.tag_constants.SERVING], strip_default_attrs=True)
    builder.save()

This appears to be working, as can be seen by loading and using the export model.

2. Load and Use Exported Model

After much tinkering, I found this saved model contains all named tensors for my task, as I was able to access `loaded_graph.get_tensor_by_name(‘input:0’)`, and other named tensors:

# load input and output vocabularies
import helper
_, (source_vocab_to_int, target_vocab_to_int), (source_int_to_vocab, target_int_to_vocab) = helper.load_preprocess()
load_path = helper.load_params()
# load and use exported model
import time
import os
import tensorflow as tf
# model_to_use = '20180723-125517' # 20180723-125517 is guaranteed to work
model_to_use = '20180723-133127'
export_dir = os.path.join('models', model_to_use)
batch_size = 32
text = \
"""
Прислушайся как дружественно струны
Вступают в строй и голос подают
Как будто мать отец и отрок юный
В счастливом единении поют
Нам говорит согласье струн в концерте
Что одинокий путь подобен смерти
"""
words = text.split()
loaded_graph = tf.Graph()
with tf.Session(graph=loaded_graph) as sess:
    tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.TRAINING], export_dir)
input_data = loaded_graph.get_tensor_by_name('input:0')
    logits = loaded_graph.get_tensor_by_name('predictions:0')
    target_sequence_length = loaded_graph.get_tensor_by_name('target_sequence_length:0')
    source_sequence_length = loaded_graph.get_tensor_by_name('source_sequence_length:0')
    keep_prob = loaded_graph.get_tensor_by_name('keep_prob:0')
    
    
    # Translate all words
    print('Russian, IPA')
    for word in words:
        prediction_x = word
        prediction_x = "\n".join([" ".join(line) for line in prediction_x.split("\n")])
translate_sentence = sentence_to_seq(prediction_x, source_vocab_to_int)
    
        translate_logits = sess.run(logits, {input_data: [translate_sentence]*batch_size,
                                             target_sequence_length: [len(translate_sentence)*2]*batch_size,
                                             source_sequence_length: [len(translate_sentence)]*batch_size,
                                             keep_prob: 1.0})[0]
x = "".join([source_int_to_vocab[i] for i in translate_sentence])
        y = "".join([target_int_to_vocab[i] for i in translate_logits])
        print(x + ", " + y)

The output is an expected

Russian, IPA
прислушайся, prʲɪslʊˈʂajsʲə<EOS>
как, kɐˈk<EOS>
дружественно, drʊˈʐɛstvʲɪn(ː)ə<EOS>
струны, ˈstrunɨ<EOS>
вступают, fstʊˈpajʊt<EOS>
в, f<EOS>
строй, stroj<EOS>
и, ɪˈ
голос, ɡɐˈlos<EOS>
подают, pɐˈdajʊt<EOS>
как, kɐˈk<EOS>
будто, bʊtˈda<EOS>
мать, matʲ<EOS>
отец, ɐˈtʲesʲ<EOS>
и, ɪˈ
отрок, ɐtrɐˈk<EOS>
юный, ˈjʉnɨj<EOS>
в, f<EOS>
счастливом, ˈɕːasʲtʲɪvəf<EOS>
единении, (j)ɪdʲɪˈnʲenʲɪje
поют, pɐˈjut<EOS>
нам, nɐm<EOS>
говорит, ɡəvɐˈrʲit<EOS>
согласье, sɐˈɡlasʲje<EOS>
струн, strʊnʲ<EOS>
в, f<EOS>
концерте, kɐnˈt͡sɛrtʲe<EOS>
что, t͡ɕˈta
одинокий, ɐdʲɪˈnokʲɪj<EOS>
путь, pʊtʲ<EOS>
подобен, pədɐˈbʲen<EOS>
смерти, smʲɪrˈtʲi<EOS>

Thank you for reading! I don’t fully understand how the model builder works, but I will continually update this post with more insights. Please leave a comment if you have suggestions or feedback!