tensorflow冻结模型为pb文件的各种方法

笔者最近因为工作需要将TensorFlow训练模型迁移到晟腾芯片平台上离线推理，所以需要将ckpt或者h5模型冻结成pb，再利用ATC模型转换工具将pb转为离线模型om文件，期间遇到一些问题和坑，总结一下，供大家参考。

1.Tensorflow1.x

训练好的模型Ckpt文件：

DB_resnet_v1_50_adam_model.ckpt-16801.index
DB_resnet_v1_50_adam_model.ckpt-16801.data-00000-of-00001
DB_resnet_v1_50_adam_model.ckpt-16801.meta

model.ckpt.meta文件保存了TensorFlow计算图的结构，可以理解为神经网络的网络结构，该文件可以被 tf.train.import_meta_graph 加载到当前默认的图来使用。
ckpt.data ： 保存模型中每个变量的取值

对于这种文件，转换成pb的方法，大家应该很熟悉了，在csdn上面可以搜到很多现成的转换代码，例如：
tf1.x ckpt转pb文件
核心思想就是找到输出节点，然后冻结pb。
本文介绍一种简单的方法，可以利用tf自带的冻结脚本，这里以tf1.15为例（不同版本py文件路径有可能不能）：

python3.7.5 /usr/local/python3.7.5/lib/python3.7/site-packages/tensorflow_core/python/tools/freeze_graph.py  \
--input_checkpoint=./iwslt2016_E19L2.64-29146 \
--output_graph=./transformer.pb \
--output_node_names="transformer/strided_silice_8" \
--input_meta_graph=./iwslt2016_E19L2.64-29146.meta  \
--input_binary=true

参考以上命令，根据你自己的ckpt的命名和模型输出节点修改一下，执行命令即可得到冻结的pb文件。

2.Tensorflow2.x

Keras现在是一个非常流行的工具库，包括tensorflow已经把keras合并到了自己的主代码当中了，大家可以直接tf.keras就可以直接调用其中的工具库。

如果是用Keras训练得到的模型，我们想移植到昇腾芯片上运行，那么就需要先把模型固化成TF1.x的pb格式，然后才能使用ATC模型转换工具转成om离线模型。
开始转换之前，需要确认Keras保存的模型文件（hdf5或者h5）是完整的模型文件还是仅保存权重文件，二者使用的接口不一样：

model.save(“xxx.h5”) ------保存的是完整的模型文件，即模型结构+权重文件

model.save_weights(“xxx.h5”)------保存的仅仅是权重文件，还需要调用model.to_json()来保存模型结构到一个json文件中

对于这样的h5文件转pb模型，我们可以参考：
https://github.com/amir-abdi/keras_to_tensorflow

2.1、对于保存完整模型的h5文件：

python3 keras_to_tensorflow.py 
    --input_model="path/to/keras/model.h5" 
    --output_model="path/to/save/model.pb"

2.2、对于模型权重和结构分开保存的文件：

python3 keras_to_tensorflow.py 
    --input_model="path/to/keras/model.h5" 
    --input_model_json="path/to/keras/model.json" 
    --output_model="path/to/save/model.pb"

另外还有一种情况就是模型保存的时候传入的是目录：

model.save("./checkpoint")

这时候保存的就是checkpoint文件，不同于tf1.x的checkpoint，这个目录下保存的是pb+若干子目录，不能使用tf1.x固化checkpoint的方法来生成pb文件。

对于这种ckpt文件，可以参考我的方法来生成pb：

import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
import pdb

new_model = tf.keras.models.load_model('./Wide_ResNet/')

full_model = tf.function(lambda x: new_model(x))
full_model = full_model.get_concrete_function(x=tf.TensorSpec((None,32,32,3),'float32'))

forzen_func = convert_variables_to_constants_v2(full_model)
forzen_func.graph.as_graph_def()

layers = [op.name for op in forzen_func.graph.get_operations()]
print("-"*50)
print("Frozen model layers:")
for layer in layers:
    print(layer)

print("*"*50)
print("Frozen model input:")
print(forzen_func.inputs)
print("Frozen model output:")
print(forzen_func.outputs)

tf.io.write_graph(
    graph_or_graph_def=forzen_func.graph,
    logdir="./",
    name="WRN.pb",
    as_text=False
)

full_model = full_model.get_concrete_function(x=tf.TensorSpec((None,32,32,3),‘float32’))--------输入的shape请根据自己的网络进行修改。

3.常见问题：

1、Keras加载模型时报错找不到自定义的层？

解决方法：找到训练代码中对应的自定义层代码，在加载模型时引入自定义层即可。

例如这样：

from keras.models import load_model
model = load_model('model.h5', custom_objects={'SincConv1D': SincConv1D})

2、一些特殊的网络，比如Retinanet，单独定义了第三方库，使用load_model加载模型失败？

解决方法：使用对应第三方库的模型加载接口来加载模型，然后再转pb

keras_retinanet.models.loadmodel("xxx.h5", backbone_name='resnet50')