blob: ae771cf80a9f1bdca84d402bf14cfbc7259d3f95 (
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] (If omitted, the module installed in system will be used by default.)"
}
SCRIPT_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
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 [ ! -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)
{
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
$FREEZE_GRAPH \
--input_saved_model_dir="$SAVED_MODEL" \
--input_binary=True \
--output_node_names="$OUTPUT" \
--output_graph="$FROZEN_PATH"
else
$FREEZE_GRAPH \
--input_meta_graph="$META_GRAPH" \
--input_checkpoint="$CKPT_PATH" \
--input_binary=True \
--output_node_names="$OUTPUT" \
--output_graph="$FROZEN_PATH"
fi
if [ -e "$TF_DIR" ]; then
popd
fi
echo "OUTPUT FILE : $FROZEN_PATH"
}
|