Add freeze_graph script into tflkit (#3556)

freeze_graph tool will freeze a standalone graphdef with the stored variables.

Signed-off-by: Jiyoung Yun <jy910.yun@samsung.com>

diff --git a/tools/tflkit/README.md b/tools/tflkit/README.md
index 1136a22..805a6e3 100644 (file)
--- a/tools/tflkit/README.md
+++ b/tools/tflkit/README.md
@@ -238,3 +238,129 @@ $ ./transform_graph.sh --info=./transform.info
 $ ls *.pb
 inception_v3.pb  inception_v3.transform.pb
 ```
+
+## Freeze a TensorFlow model
+
+### TensorFlow
+
+TensorFlow provides methods to save and restore models. Each method stores related files in different ways. Here are two common ways to save the freeze stored models.
+  1. Use [tf.train.Saver](https://www.tensorflow.org/guide/saved_model#save_and_restore_variables)
+    In this way, it creates a `MetaGraphDef` file and checkpoint files that contain the saved variables. Saving this way will result in the following files in the exported directory:
+
+      ```
+      $ ls /tmp/saver/
+      checkpoint  model.ckpt.data-00000-of-00001  model.ckpt.index  model.ckpt.meta
+      ```
+
+  2. Use [SavedModel](https://www.tensorflow.org/guide/saved_model#build_and_load_a_savedmodel)
+    It is the easiest way to create a saved model. Saving this way will result in the following files in the exported directory:
+
+      ```
+      $ ls /tmp/saved_model/
+      saved_model.pb  variables
+      $ tree /tmp/saved_model/
+      /tmp/saved_model/
+      ├── saved_model.pb
+      └── variables
+          ├── variables.data-00000-of-00001
+          └── variables.index
+      ```
+
+The [freeze_graph](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/tools/freeze_graph.py) tool receives these files as input parameters and combines the stored variables and the standalone graphdef to generate a frozen graphdef file.
+
+### with tflkit
+
+The tflkit provides the simple way to create a frozen graph using [freeze_graph](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/tools/freeze_graph.py) tool. This tool requires an information file as a parameter. There is an [example file](info/freeze.info) for a freeze tool. Either `SAVED_MODEL` or `META_GRAPH` must be declared. And `META_GRAPH` is always used with `CKPT_PATH`. The `--tensorflow_path` can change the TensorFlow location. By default, it uses `externals/tensorflow` directory.
+
+Freeze information:
+  * SAVED_MODEL : Full directory path with TensorFlow `SavedModel` file and variables.
+  * META_GRAPH : Full filepath of file containing TensorFlow `MetaGraphDef`.
+  * CKPT_PATH : Full filepath of file containing TensorFlow variables. (ex. [...]/*.ckpt)
+  * FROZEN_PATH : Full filepath to save the output frozen graph.
+  * OUTPUT : Names of the output arrays, comma-separated.
+
+Usage (for example, `tf.train.Saver`):
+```
+$ cat > sample_saver.py << END
+import tensorflow as tf
+
+# Create some variables.
+v1 = tf.get_variable("v1", shape=[3], initializer = tf.zeros_initializer)
+v2 = tf.get_variable("v2", shape=[5], initializer = tf.zeros_initializer)
+
+inc_v1 = v1.assign(v1+1)
+dec_v2 = v2.assign(v2-1)
+
+# Add an op to initialize the variables.
+init_op = tf.global_variables_initializer()
+
+# Add ops to save and restore all the variables.
+saver = tf.train.Saver()
+
+# Later, launch the model, initialize the variables, do some work, and save the
+# variables to disk.
+with tf.Session() as sess:
+  sess.run(init_op)
+  # Do some work with the model.
+  inc_v1.op.run()
+  dec_v2.op.run()
+  # Save the variables to disk.
+  save_path = saver.save(sess, "/tmp/saver/model.ckpt")
+  print("Model saved in path: %s" % save_path)
+END
+$ python sample_saver.py
+$ ls /tmp/saver/
+checkpoint  model.ckpt.data-00000-of-00001  model.ckpt.index  model.ckpt.meta
+$ cat > freeze_saver.info << END
+#SAVED_MODEL=""                                  
+META_GRAPH="/tmp/saver/model.ckpt.meta"
+CKPT_PATH="/tmp/saver/model.ckpt"
+FROZEN_PATH="/tmp/saver/model.frozen.pb"
+OUTPUT="v2"
+END
+$ ./freeze_graph.sh --info=./freeze_saver.info
+$ ls /tmp/saver/*.pb
+/tmp/saver/model.frozen.pb
+```
+
+Usage (for example, `SavedModel`):
+```
+$ cat > sample_saved_model.py << END
+import tensorflow as tf
+
+# Create some variables.
+v1 = tf.get_variable("v1", shape=[3], initializer = tf.zeros_initializer)
+v2 = tf.get_variable("v2", shape=[5], initializer = tf.zeros_initializer)
+
+inc_v1 = v1.assign(v1+1)
+dec_v2 = v2.assign(v2-1)
+
+# Add an op to initialize the variables.
+init_op = tf.global_variables_initializer()
+
+# Later, launch the model, initialize the variables, do some work, and save the
+# variables to disk.
+with tf.Session() as sess:
+  sess.run(init_op)
+  # Do some work with the model.
+  inc_v1.op.run()
+  dec_v2.op.run()
+  # Save the variables to disk.
+  tf.saved_model.simple_save(sess, "/tmp/saved_model", inputs={'v1':v1}, outputs={'v2':v2})
+END
+$ python sample_saved_model.py
+$ ls /tmp/saved_model/
+saved_model.pb  variables
+$ cat > freeze_saved_model.info << END
+SAVED_MODEL="/tmp/saved_model"
+#META_GRAPH=
+#CKPT_PATH=
+FROZEN_PATH="/tmp/saved_model/model.frozen.pb"
+OUTPUT="v2"
+END
+$ ./freeze_graph.sh --info=./info/freeze_saved_model.info
+$ ls /tmp/saved_model/
+model.frozen.pb  saved_model.pb  variables
+$ ls /tmp/saved_model/*.frozen.pb
+/tmp/saved_model/model.frozen.pb
+```
\ No newline at end of file

diff --git a/tools/tflkit/freeze_graph.sh b/tools/tflkit/freeze_graph.sh
new file mode 100755 (executable)
index 0000000..c491ba4
--- /dev/null
+++ b/tools/tflkit/freeze_graph.sh
@@ -0,0 +1,98 @@
+#!/bin/bash
+
+usage()
+{
+  echo "usage : $0"
+  echo "       --info=Information file"
+  echo "       --tensorflow_path=TensorFlow path (Use externals/tensorflow by default)"
+}
+
+SCRIPT_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+TF_DIR="${SCRIPT_PATH}/../../externals/tensorflow"
+
+for i in "$@"
+do
+  case $i in
+    --info=*)
+      INFO=${i#*=}
+      ;;
+    --tensorflow_path=*)
+      TF_DIR=${i#*=}
+      ;;
+    -h|--help)
+      usage
+      exit 0
+      ;;
+    *)
+      usage
+      exit 1
+      ;;
+  esac
+  shift
+done
+
+if [ -z "$INFO" ]; then
+  echo "INFO is unset or set to the empty string"
+  usage
+  exit 1
+fi
+if [ -z "$TF_DIR" ]; then
+  echo "tensorflow_path is unset or set to the empty string"
+  usage
+  exit 1
+fi
+
+if [ ! -x "$(command -v bazel)" ]; then
+  echo "Cannot find bazel. Please install bazel."
+  exit 1
+fi
+
+source $INFO
+
+if [ -z "$SAVED_MODEL" ] && [ -z "$META_GRAPH" ]; then
+  echo "SAVED_MODEL or META_GRAPH + CKPT_PATH is unset or set to the empty string"
+  echo "Update the $INFO file"
+  exit 1
+fi
+if [ ! -z "$META_GRAPH" ] && [ -z "$CKPT_PATH" ]; then
+  echo "META_GRAPH is always used with CKPT_PATH"
+  echo "CKPT_PATH is unset or set to the empty string"
+  echo "Update the $INFO file"
+  exit 1
+fi
+if [ -z "$FROZEN_PATH" ]; then
+  echo "FROZEN_PATH is unset or set to the empty string"
+  echo "Update the $INFO file"
+  exit 1
+fi
+if [ -z "$OUTPUT" ]; then
+  echo "OUTPUT is unset or set to the empty string"
+  echo "Update the $INFO file"
+  exit 1
+fi
+
+CUR_DIR=$(pwd)
+{
+  echo "Enter $TF_DIR"
+  pushd $TF_DIR > /dev/null
+
+  if [ ! -z $SAVED_MODEL ]; then
+    bazel run tensorflow/python/tools:freeze_graph -- \
+    --input_saved_model_dir="$SAVED_MODEL" \
+    --input_binary=True \
+    --output_node_names="$OUTPUT" \
+    --output_graph="$FROZEN_PATH"
+  else
+    bazel run tensorflow/python/tools:freeze_graph -- \
+    --input_meta_graph="$META_GRAPH" \
+    --input_checkpoint="$CKPT_PATH" \
+    --input_binary=True \
+    --output_node_names="$OUTPUT" \
+    --output_graph="$FROZEN_PATH"
+  fi
+
+  popd
+
+  echo "OUTPUT FILE : $FROZEN_PATH"
+}

diff --git a/tools/tflkit/info/freeze.template b/tools/tflkit/info/freeze.template
new file mode 100644 (file)
index 0000000..076819c
--- /dev/null
+++ b/tools/tflkit/info/freeze.template
@@ -0,0 +1,5 @@
+SAVED_MODEL=
+META_GRAPH=
+CKPT_PATH=
+FROZEN_PATH=
+OUTPUT=