深度学习之格式转换笔记(三)：keras(.hdf5)模型转TensorFlow(.pb) 转TensorRT(.uff)格式

环境：
tensorflow1.15，cuda10.0，cudnn7.6.4

将keras训练好保存的.hdf5格式模型转为tensorflow的.pb模型，然后转为tensorrt支持的uff格式。

keras(.hdf5)模型转TensorFlow(.pb)

# h5_to_pb.py

from keras.models import load_model
import tensorflow as tf
import os
import os.path as osp
from keras import backend as K

# 路径参数
input_path = 'D:/pycharm/facenet/models/'
weight_file = 'resnet50.hdf5'
weight_file_path = osp.join(input_path, weight_file)
output_graph_name = weight_file[:-3] + '.pb'


# 转换函数
def h5_to_pb(h5_model, output_dir, model_name, out_prefix="output_", log_tensorboard=True):
    if osp.exists(output_dir) == False:
        os.mkdir(output_dir)
    out_nodes = []
    for i in range(len(h5_model.outputs)):
        out_nodes.append(out_prefix + str(i + 1))
        tf.identity(h5_model.output[i], out_prefix + str(i + 1))
    sess = K.get_session()
    from tensorflow.python.framework import graph_util, graph_io
    init_graph = sess.graph.as_graph_def()
    main_graph = graph_util.convert_variables_to_constants(sess, init_graph, out_nodes)
    graph_io.write_graph(main_graph, output_dir, name=model_name, as_text=False)
    if log_tensorboard:
        from tensorflow.python.tools import import_pb_to_tensorboard
        import_pb_to_tensorboard.import_to_tensorboard(osp.join(output_dir, model_name), output_dir)


# 输出路径
output_dir = osp.join(os.getcwd(), "trans_model")
# 加载模型
print(weight_file_path)
h5_model = load_model(weight_file_path)
# model.load_weights
h5_to_pb(h5_model, output_dir=output_dir, model_name=output_graph_name)
print('model saved')

TensorFlow(.pb) 转TensorRT(.uff)

只需要安装好TensorRT之后通过终端运行convert-to-uff resnet50.pb即可得到resnet50.uff 。

测试：如果你想实验一下可以根据tensorrt自带的例子来进行实验，过程如下：

• 在jetson nano上测试tensorRT，tensorrt是当时安装镜像的时候就自带有得，通过import tensorrt和import
  uff都能够成功，进行测试得时候我用得是在/usr/src/tensorrt/samples/python文件夹下有很多python的例子，我们以第一个end_to_end_tensorflow_mnist的例子为例，

  mkdir models python model.py

  然后下载数据mnist.npz，并开始训练，完成后在models文件夹下生成lenet5.pb文件
  
  然后在进行格式转换，要将tensorflow的pb文件转化为uff格式的文件，首先找到convert_to_uff文件，看自己用的是哪个版本的python，如果是python3，则在/usr/lib/python3.5/dist-packages/uff/bin文件夹下，如果是python2，则在/usr/lib/python2.7/dist-packages/uff/bin文件夹下，我们在终端中进入end_to_end_tensorflow_mnist，运行以下指令

python3 /usr/lib/python3.6/dist-packages/uff/bin/convert_to_uff.py --input_file models/lenet5.pb

则会在models文件夹中生成lenet5.uff文件，完成转换，在通过sample.py进行测试，出现报错没有安装pycuda，后发现这个库是用于GPU加速的。

如果在jetson nano上安装pycuda不会的同学可以看这篇文章：
https://blog.csdn.net/m0_51004308/article/details/115280360

再次运行python3 sample.py 报错Could not find 9.pgm。

解决：在/usr/src/tensorrt/data/mnist下找到download_pgms.py文件并通过python3运行即可。
后面就通过了，运行sample.py之后出现
Test Case：8
Prediction：8