使用freeze_graph将SavedModel生成的模型合并成一个pb文件

技巧:使用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 
    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 
    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

嗯...具体这个文件能不能直接用还有待研究

你可能感兴趣的:(使用freeze_graph将SavedModel生成的模型合并成一个pb文件)