Exporting an Inference Graph - 1.2 English

First, you need to write the codes of building a TensorFlow graph for training and evaluation with each part written in a separate script. If you have trained a baseline model before, you should already have the training codes, then you only need to prepare the codes for evaluation. The evaluation script must contain a function named model_fn that creates all needed nodes from input to output. The function should return a dictionary that maps output nodes’ names to their operations or a tf.estimator.Estimator. For example, if your network is an image classifier, then the returned dictionary will usually include ops to calculate top-1 and top-5 accuracy like this:

def model_fn():
  # graph definition codes here
  # ……
return {
      'top-1': slim.metrics.streaming_accuracy(predictions, labels),
      'top-5': slim.metrics.streaming_recall_at_k(logits, org_labels, 5)
  }

Or, if you use TensorFlow Estimator API to train and evaluate your network, your model_fn must return an instance of tf.estimator. At the same time, you also need to provide a function called eval_input_fn, which the Estimator uses to get the data used in the evaluation.

def cnn_model_fn(features, labels, mode):
  # codes for building graph here
…
eval_metric_ops = {
      "accuracy": tf.metrics.accuracy(
          labels=labels, predictions=predictions["classes"])}
  return tf.estimator.EstimatorSpec(
      mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

def model_fn():
  return tf.estimator.Estimator(
      model_fn=cnn_model_fn, model_dir="./models/train/")

mnist = tf.contrib.learn.datasets.load_dataset("mnist")
train_data = mnist.train.images # Returns np.array
train_labels = np.asarray(mnist.train.labels, dtype=np.int32)
eval_data = mnist.test.images # Returns np.array
eval_labels = np.asarray(mnist.test.labels, dtype=np.int32)

def eval_input_fn():
  return tf.estimator.inputs.numpy_input_fn(
      x={"x": eval_data},
      y=eval_labels,
      num_epochs=1,
      shuffle=False)

The evaluation codes are used to export an inference GraphDef file and evaluate network’s performance during pruning.

Use the following code to export a GraphDef proto file:

import tensorflow as tf
from google.protobuf import text_format
from tensorflow.python.platform import gfile

with tf.Graph().as_default() as graph:
# your graph definition here
# ……
    graph_def = graph.as_graph_def()
    with gfile.GFile(‘inference_graph.pbtxt’, 'w') as f:
      f.write(text_format.MessageToString(graph_def))