diff options
Diffstat (limited to 'tools/tflkit/freeze_graph.sh')
| -rwxr-xr-x | tools/tflkit/freeze_graph.sh | 26 |
1 files changed, 13 insertions, 13 deletions
diff --git a/tools/tflkit/freeze_graph.sh b/tools/tflkit/freeze_graph.sh index c491ba4d2..ae771cf80 100755 --- a/tools/tflkit/freeze_graph.sh +++ b/tools/tflkit/freeze_graph.sh @@ -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" } |