Exporting and Importing a MetaGraph

A MetaGraph contains both a TensorFlow GraphDef as well as associated metadata necessary for running computation in a graph when crossing a process boundary. It can also be used for long term storage of graphs. The MetaGraph contains the information required to continue training, perform evaluation, or run inference on a previously trained graph.

The APIs for exporting and importing the complete model are in the tf.train.Saver class: export_meta_graph and import_meta_graph.

What's in a MetaGraph

The information contained in a MetaGraph is expressed as a MetaGraphDef protocol buffer. It contains the following fields:

  • MetaInfoDef for meta information, such as version and other user information.
  • GraphDef for describing the graph.
  • SaverDef for the saver.

  • CollectionDef map that further describes additional components of the model, such as Variables, QueueRunners, etc. In order for a Python object to be serialized to and from MetaGraphDef, the Python class must implement to_proto() and from_proto() methods, and register them with the system using register_proto_function.

    For example,

    def to_proto(self):
  """Converts a `Variable` to a `VariableDef` protocol buffer.

  Returns:
    A `VariableDef` protocol buffer.
  """
  var_def = variable_pb2.VariableDef()
  var_def.variable_name = self._variable.name
  var_def.initializer_name = self.initializer.name
  var_def.snapshot_name = self._snapshot.name
  if self._save_slice_info:
    var_def.save_slice_info_def.MergeFrom(self._save_slice_info.to_proto())
  return var_def

@staticmethod
def from_proto(variable_def):
  """Returns a `Variable` object created from `variable_def`."""
  return Variable(variable_def=variable_def)

ops.register_proto_function(ops.GraphKeys.VARIABLES,
                            proto_type=variable_pb2.VariableDef,
                            to_proto=Variable.to_proto,
                            from_proto=Variable.from_proto)

Exporting a Complete Model to MetaGraph

The API for exporting a running model as a MetaGraph is export_meta_graph().

  def export_meta_graph(filename=None, collection_list=None, as_text=False):
    """Writes `MetaGraphDef` to save_path/filename.

    Args:
      filename: Optional meta_graph filename including the path.
      collection_list: List of string keys to collect.
      as_text: If `True`, writes the meta_graph as an ASCII proto.

    Returns:
      A `MetaGraphDef` proto.
    """

A collection can contain any Python objects that users would like to be able to uniquely identify and easily retrieve. These objects can be special operations in the graph, such as train_op, or hyper parameters, such as "learning rate". Users can specify the list of collections they would like to export. If no collection_list is specified, all collections in the model will be exported.

The API returns a serialized protocol buffer. If filename is specified, the protocol buffer will also be written to a file.

Here are some of the typical usage models:

  • Export the default running graph:

    # Build the model
...
with tf.Session() as sess:
# Use the model
...
# Export the model to /tmp/my-model.meta.
meta_graph_def = tf.train.export_meta_graph(filename='/tmp/my-model.meta')

  • Export the default running graph and only a subset of the collections.

    meta_graph_def = tf.train.export_meta_graph(
  filename='/tmp/my-model.meta',
  collection_list=["input_tensor", "output_tensor"])

The MetaGraph is also automatically exported via the save() API in tf.train.Saver.

Import a MetaGraph

The API for importing a MetaGraph file into a graph is import_meta_graph().

Here are some of the typical usage models:

  • Import and continue training without building the model from scratch.

    ...
# Create a saver.
saver = tf.train.Saver(...variables...)
# Remember the training_op we want to run by adding it to a collection.
tf.add_to_collection('train_op', train_op)
sess = tf.Session()
for step in xrange(1000000):
    sess.run(train_op)
    if step % 1000 == 0:
        # Saves checkpoint, which by default also exports a meta_graph
        # named 'my-model-global_step.meta'.
        saver.save(sess, 'my-model', global_step=step)

    Later we can continue training from this saved meta_graph without building the model from scratch.

    with tf.Session() as sess:
  new_saver = tf.train.import_meta_graph('my-save-dir/my-model-10000.meta')
  new_saver.restore(sess, 'my-save-dir/my-model-10000')
  # tf.get_collection() returns a list. In this example we only want the
  # first one.
  train_op = tf.get_collection('train_op')[0]
  for step in xrange(1000000):
    sess.run(train_op)

  • Import and extend the graph.

    For example, we can first build an inference graph, export it as a meta graph:

    # Creates an inference graph.
# Hidden 1
images = tf.constant(1.2, tf.float32, shape=[100, 28])
with tf.name_scope("hidden1"):
  weights = tf.Variable(
      tf.truncated_normal([28, 128],
                          stddev=1.0 / math.sqrt(float(28))),
      name="weights")
  biases = tf.Variable(tf.zeros([128]),
                       name="biases")
  hidden1 = tf.nn.relu(tf.matmul(images, weights) + biases)
# Hidden 2
with tf.name_scope("hidden2"):
  weights = tf.Variable(
      tf.truncated_normal([128, 32],
                          stddev=1.0 / math.sqrt(float(128))),
      name="weights")
  biases = tf.Variable(tf.zeros([32]),
                       name="biases")
  hidden2 = tf.nn.relu(tf.matmul(hidden1, weights) + biases)
# Linear
with tf.name_scope("softmax_linear"):
  weights = tf.Variable(
      tf.truncated_normal([32, 10],
                          stddev=1.0 / math.sqrt(float(32))),
      name="weights")
  biases = tf.Variable(tf.zeros([10]),
                       name="biases")
  logits = tf.matmul(hidden2, weights) + biases
  tf.add_to_collection("logits", logits)

init_all_op = tf.initialize_all_variables()

with tf.Session() as sess:
  # Initializes all the variables.
  sess.run(init_all_op)
  # Runs to logit.
  sess.run(logits)
  # Creates a saver.
  saver0 = tf.train.Saver()
  saver0.save(sess, saver0_ckpt)
  # Generates MetaGraphDef.
  saver0.export_meta_graph('my-save-dir/my-model-10000.meta')

    Then later import it and extend it to a training graph.

    with tf.Session() as sess:
  new_saver = tf.train.import_meta_graph('my-save-dir/my-model-10000.meta')
  new_saver.restore(sess, 'my-save-dir/my-model-10000')
  # Addes loss and train.
  labels = tf.constant(0, tf.int32, shape=[100], name="labels")
  batch_size = tf.size(labels)
  labels = tf.expand_dims(labels, 1)
  indices = tf.expand_dims(tf.range(0, batch_size), 1)
  concated = tf.concat(1, [indices, labels])
  onehot_labels = tf.sparse_to_dense(
      concated, tf.pack([batch_size, 10]), 1.0, 0.0)
  logits = tf.get_collection("logits")[0]
  cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits,
                                                          onehot_labels,
                                                          name="xentropy")
  loss = tf.reduce_mean(cross_entropy, name="xentropy_mean")

  tf.scalar_summary(loss.op.name, loss)
  # Creates the gradient descent optimizer with the given learning rate.
  optimizer = tf.train.GradientDescentOptimizer(0.01)

  # Runs train_op.
  train_op = optimizer.minimize(loss)
  sess.run(train_op)

  • Retrieve Hyper Parameters

    filename = ".".join([tf.latest_checkpoint(train_dir), "meta"])
tf.train.import_meta_graph(filename)
hparams = tf.get_collection("hparams")

results matching ""

    No results matching ""