Slim生成的ckpt文件转PB文件

Posted 2022-02-09 17岁boy想当攻城狮

前言

搭建Slim框架方法请参考这篇文章：搭建迁移训练Slim框架环境_17岁boy的博客-CSDN博客

生成PB文件

转换的方式也很简单，Slim自带了一个graph.py文件，文件存在于models/research/slim/export_inference_graph.py，这个文件可以生成对应包含神经网络算子布局的PB文件，但是没有包含神经元权重这些信息，仅仅包含对应神经网络的网格结构布局（layout），说白了就是生成一个包含神经网络的空壳结构文件，只是没有实际的训练出来的权重数据。

输入如下命令可以开始转换：

命令的参数这里就不介绍了，参数作用可以参考“前言”里的链接

python3.7 models/research/slim/export_inference_graph.py         \\
--alsologtostderr                                                \\
--dataset_name=flowers                                           \\
--model_name=mobilenet_v1_025                                    \\
--image_size=96                                                  \\
--use_grayscale=True                                             \\
--output_file=./flowers_graph.pb

执行完成之后当前目录下会生成一个flowers_graph.pb的文件，这个文件里包含了神经网络的网格结构信息，但是缺少权重信息，所以我们还需要将刚刚生成的ckpt文件里的权重信息提取出来放入到pb文件里去。

这里我们需要Tensor Flow的模块，需要下载Tensor Flow

git clone https://github.com/tensorflow/tensorflow

转换文件在：tensorflow/tensorflow/python/tools/freeze_graph.py，这个文件专门用于将ckpt神经网络里的权重信息copy到pb文件中

输入如下命令开始copy，这一步也叫冻结权重

python3.7 tensorflow/tensorflow/python/tools/freeze_graph.py     \\
--input_graph=./flowers_graph.pb                                 \\
--input_checkpoint=./train/model.ckpt-0                          \\
--input_binar=true                                               \\
--output_graph=flowers_frozen.pb                                 \\
--output_node_names=MobilenetV1/Predictions/Reshape_1

参数介绍：

参数	作用
--input_graph	graph pb文件路径，就是刚刚生成的pb文件
--input_checkpoint	要导出的ckpt路径，这里文件名给到ckpt-0就可以了后面不用给，因为ckpt文件是多个文件存储不同的网格数据，所以给前半部分，后半部分由脚本自动去取
--input_binar	以二进制格式导出,一定要以二进制格式导出，不然遇到了一些\\r\\0或者一些特殊转义符操作系统会自动编码，导致原数据不对
--output_graph	合并后的pb文件路径
--output_node_names	模型输出层节点名称，这个需要与神经网络有关，在训练时定义的，需要获取输出层里的loss函数相关算子以及输出结构，因为这一层的作用就是拿到全连接层输出的结构进行loss运算最终在按对应结构输出。 需要注意的是PB只能有一个输出层，如果你有多个输出层请选择针对当前模型情况下需要用到的那个。

运行后报了个错：

ImportError: cannot import name 'py_checkpoint_reader' from 'tensorflow.python.training' (/home/yongde1x/.local/lib/python3.7/site-packages/tensorflow/python/training/__init__.py)

解决方案请参考我的这篇文章：ImportError: cannot import name ‘py_checkpoint_reader‘ from ‘tensorflow.pyth_17岁boy的博客-CSDN博客

然后在执行就生成了对应的frozen.pb文件