summaryrefslogtreecommitdiff
path: root/tools/tflkit/freeze_graph.sh
blob: c491ba4d2c8a161966bd11f325e2d2aea8f8cc5c (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
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"
}