Skip to content

Instantly share code, notes, and snippets.

@dcalsky
Forked from zhanwenchen/export_tf_model.py
Created October 14, 2019 23:51
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save dcalsky/43239a21b7da23fc988742edbfbec3cb to your computer and use it in GitHub Desktop.
Save dcalsky/43239a21b7da23fc988742edbfbec3cb to your computer and use it in GitHub Desktop.
Minimal code to load a trained TensorFlow model from a checkpoint and export it with SavedModelBuilder
import os
import tensorflow as tf
trained_checkpoint_prefix = 'checkpoints/dev'
export_dir = os.path.join('models', '0') # IMPORTANT: each model folder must be named '0', '1', ... Otherwise it will fail!
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, tf.saved_model.tag_constants.SERVING],
strip_default_attrs=True)
builder.save()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment