Update `freeze_graph.sh` script (#6345)
author이성재/On-Device Lab(SR)/Principal Engineer/삼성전자 <sj925.lee@samsung.com>
Fri, 9 Aug 2019 10:23:06 +0000 (19:23 +0900)
committerGitHub Enterprise <noreply-CODE@samsung.com>
Fri, 9 Aug 2019 10:23:06 +0000 (19:23 +0900)
* Change to use the 'tensorflow' module installed on the system by default.

Signed-off-by: Sung-Jae Lee <sj925.lee@samsung.com>
tools/tflkit/freeze_graph.sh

index c491ba4..ae771cf 100755 (executable)
@@ -4,13 +4,11 @@ usage()
 {
   echo "usage : $0"
   echo "       --info=Information file"
-  echo "       --tensorflow_path=TensorFlow path (Use externals/tensorflow by default)"
+  echo "       [--tensorflow_path=TensorFlow path] (If omitted, the module installed in system will be used by default.)"
 }
 
 SCRIPT_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
-TF_DIR="${SCRIPT_PATH}/../../externals/tensorflow"
-
 for i in "$@"
 do
   case $i in
@@ -37,11 +35,6 @@ if [ -z "$INFO" ]; then
   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."
@@ -74,17 +67,22 @@ fi
 
 CUR_DIR=$(pwd)
 {
-  echo "Enter $TF_DIR"
-  pushd $TF_DIR > /dev/null
+  if [ -e "$TF_DIR" ]; then
+    echo "Enter $TF_DIR"
+    pushd $TF_DIR > /dev/null
+    FREEZE_GRAPH="bazel run tensorflow/python/tools:freeze_graph -- "
+  else
+    FREEZE_GRAPH="python -m tensorflow.python.tools.freeze_graph "
+  fi
 
   if [ ! -z $SAVED_MODEL ]; then
-    bazel run tensorflow/python/tools:freeze_graph -- \
+    $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 -- \
+    $FREEZE_GRAPH \
     --input_meta_graph="$META_GRAPH" \
     --input_checkpoint="$CKPT_PATH" \
     --input_binary=True \
@@ -92,7 +90,9 @@ CUR_DIR=$(pwd)
     --output_graph="$FROZEN_PATH"
   fi
 
-  popd
+  if [ -e "$TF_DIR" ]; then
+    popd
+  fi
 
   echo "OUTPUT FILE : $FROZEN_PATH"
 }