Imported Upstream version 0.3upstream/0.3

1 files changed, 98 insertions, 0 deletions

diff --git a/tools/tflkit/freeze_graph.sh b/tools/tflkit/freeze_graph.sh
new file mode 100755
index 000000000..c491ba4d2
--- /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"
+}