技巧:使用freeze_graph将SavedModel生成的模型合并成一个pb文件
在训练mnist模型的时候,增加一个参数,可以在训练完毕后导出SavedModel
python mnist.py --export_dir /tmp/mnist_saved_model
训练完成后,会以以下文件结构在/tmp/mnist_saved_model生成模型文件
── 1530159051
├── saved_model.pb
└── variables
├── variables.data-00000-of-00001
└── variables.index
saved_model.pb
保存的是计算图信息
variables
保存的是计算中的变量信息
我的目是想要在Android上运行测试这个模型,但是从官方给的Demo中可以看到,所加载的模型都是单个文件,暂时也没看到加载SavedModel的接口
为了方便Android平台使用,用freeze_graph
脚本将SavedModel
合并
查看这个脚本的源码,在最开头可以看到
Converts checkpoint variables into Const ops in a standalone GraphDef file.
This script is designed to take a GraphDef proto, a SaverDef proto, and a set of
variable values stored in a checkpoint file, and output a GraphDef with all of
the variable ops converted into const ops containing the values of the
variables.
It's useful to do this when we need to load a single file in C++, especially in
environments like mobile or embedded where we may not have access to the
RestoreTensor ops and file loading calls that they rely on.
这个脚本就是用来将变量转换为常量合并到到计算图中的
具体怎么使用呢?
脚本中给出的例子是这样的:
bazel build tensorflow/python/tools:freeze_graph && \
bazel-bin/tensorflow/python/tools/freeze_graph \
--input_graph=some_graph_def.pb \
--input_checkpoint=model.ckpt-8361242 \
--output_graph=/tmp/frozen_graph.pb --output_node_names=softmax
因为我电脑是完整安装了tensorflow
的,所以可以直接使用freeze_graph
而不用通过bazel
去创建freeze_graph
freeze_graph --input_graph=saved_model.pb --input_checkpoint=variables/variables --output_graph=merge_graph.pb --output_node_names=softmax
执行,然后挂了
from ._conv import register_converters as _register_converters
Traceback (most recent call last):
File "/home/xxx/anaconda2/bin/freeze_graph", line 11, in <module>
sys.exit(run_main())
File "/home/xxx/anaconda2/lib/python2.7/site-packages/tensorflow/python/tools/freeze_graph.py", line 379, in run_main
app.run(main=my_main, argv=[sys.argv[0]] + unparsed)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/tensorflow/python/platform/app.py", line 125, in run
_sys.exit(main(argv))
File "/home/xxx/anaconda2/lib/python2.7/site-packages/tensorflow/python/tools/freeze_graph.py", line 378, in <lambda>
my_main = lambda unused_args: main(unused_args, flags)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/tensorflow/python/tools/freeze_graph.py", line 272, in main
flags.saved_model_tags, checkpoint_version)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/tensorflow/python/tools/freeze_graph.py", line 231, in freeze_graph
input_graph_def = _parse_input_graph_proto(input_graph, input_binary)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/tensorflow/python/tools/freeze_graph.py", line 174, in _parse_input_graph_proto
text_format.Merge(f.read(), input_graph_def)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/google/protobuf/text_format.py", line 536, in Merge
descriptor_pool=descriptor_pool)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/google/protobuf/text_format.py", line 590, in MergeLines
return parser.MergeLines(lines, message)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/google/protobuf/text_format.py", line 623, in MergeLines
self._ParseOrMerge(lines, message)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/google/protobuf/text_format.py", line 638, in _ParseOrMerge
self._MergeField(tokenizer, message)
File "/home/xxx/anaconda2/lib/python2.7/site-packages/google/protobuf/text_format.py", line 706, in _MergeField
name = tokenizer.ConsumeIdentifierOrNumber()
File "/home/xxx/anaconda2/lib/python2.7/site-packages/google/protobuf/text_format.py", line 1166, in ConsumeIdentifierOrNumber
raise self.ParseError('Expected identifier or number, got %s.' % result)
google.protobuf.text_format.ParseError: 1:1 : Expected identifier or number, got.
什么鬼错,Expected identifier or number
字面意思应该是传入的文件少了点什么东西
再看看freeze_graph
脚本的源码
在run_main()
函数里看到了一个参数
parser.add_argument(
"--input_saved_model_dir",
type=str,
default="",
help="Path to the dir with TensorFlow \'SavedModel\' file and variables.")
原来是用错参数了,噗...
修改一下命令
freeze_graph \
--input_saved_model_dir=1530159051 \
--output_node_names=Softmax,ArgMax \
--output_graph=merge1_graph.pb
再执行,屁都没冒一个,说明执行成功了,没报错,哈哈
看下指定的输出文件有没有生成
ww3:~/source/TFmodels_out/mnist_saved_model$ ls -l
总用量 12804
drwxr-xr-x 3 ww3 ww3 4096 6月 28 12:10 1530159051
-rw-rw-r-- 1 ww3 ww3 13105521 6月 28 17:58 merge1_graph.pb
嗯...具体这个文件能不能直接用还有待研究