Slim生成的ckpt文件转PB文件

Posted 17岁boy想当攻城狮

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了Slim生成的ckpt文件转PB文件相关的知识,希望对你有一定的参考价值。

前言

搭建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_graphgraph pb文件路径,就是刚刚生成的pb文件
--input_checkpoint要导出的ckpt路径,这里文件名给到ckpt-0就可以了后面不用给,因为ckpt文件是多个文件存储不同的网格数据,所以给前半部分,后半部分由脚本自动去取
--input_binar以二进制格式导出,一定要以二进制格式导出,不然遇到了一些\\r\\0或者一些特殊转义符操作系统会自动编码,导致原数据不对
--output_graph合并后的pb文件路径
--output_node_names

模型输出层节点名称,这个需要与神经网络有关,在训练时定义的,需要获取输出层里的loss函数相关算子以及输出结构,这一层其实就是你在训练时生成全连接层的那一个Layer名称

运行后报了个错:

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文件

以上是关于Slim生成的ckpt文件转PB文件的主要内容,如果未能解决你的问题,请参考以下文章

TensorFlow 自定义模型导出:将 .ckpt 格式转化为 .pb 格式

tensorflow模型ckpt转pb以及其遇到的问题

Tensor Flow PB文件量化到TFLITE

ckpt pb

TensorFlow 生成 .ckpt 和 .pb

何时在 Tensorflow 模型保存中使用 .ckpt、.hdf5 和 .pb 文件扩展名?