TensorRT推理加速-基於Tensorflow(keras)的uff格式模型(文件准備)

一、引子//Windows

tf（keras）訓練好了模型，想要用Nvidia-TensorRT來重構訓練好的模型為TRT推理引擎加快推理的速度。

二、准備文件

1、訓練好模型以后（keras）可以通過以下方式保存keras模型為h5文件

tf.keras.models.save_model(model, 'keras_model\\classify.h5')

 2、再通過以下代碼來將h5文件轉化為pb文件

import tensorflow.compat.v1 as tf1

tf1.reset_default_graph()
tf1.keras.backend.set_learning_phase(0)  # 調用模型前一定要執行該命令
tf1.disable_v2_behavior()  # 禁止tensorflow2.0的行為
# 加載hdf5模型
hdf5_pb_model = tf1.keras.models.load_model('keras_model\\classify.h5')


def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
    graph = session.graph
    with graph.as_default():
        #         freeze_var_names = list(set(v.op.name for v in tf1.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        #         output_names += [v.op.name for v in tf1.global_variables()]
        print("output_names", output_names)
        input_graph_def = graph.as_graph_def()
        #         for node in input_graph_def.node:
        #             print('node:', node.name)
        print("len node1", len(input_graph_def.node))
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ""
        frozen_graph = tf1.graph_util.convert_variables_to_constants(session, input_graph_def,
                                                                     output_names)

        outgraph = tf1.graph_util.remove_training_nodes(frozen_graph)  # 雲掉與推理無關的內容
        print("##################################################################")
        for node in outgraph.node:
            print('node:', node.name)
        print("len node1", len(outgraph.node))
        return outgraph


output_folder2 = 'keras_model'

frozen_graph = freeze_session(tf1.compat.v1.keras.backend.get_session(),
                              output_names=[out.op.name for out in hdf5_pb_model.outputs])
tf1.train.write_graph(frozen_graph, output_folder2, "classify.pb", as_text=False)

3、注意：以上代碼基於tf2.0運行

4、pb模型文件轉化為uff模型文件（tensorrt解析tf模型只能用uff格式）

首先，先安裝TensorRT自帶的(兩個文件就在trt文件夾里面，cd到路徑)

pip install uff-0.6.5-py2.py3-none-any.whl
pip install graphsurgeon-0.4.1-py2.py3-none-any.whl

 5、執行（cd到路徑，執行以下過程需要tf1.x版本，否則報錯，沒有Graphdef）

轉換

convert-to-uff xxxx.pb -o xxxx.uff

查看模型信息

convert-to-uff xxxx.uff -l

 參考：

【Tensorflow2.0】8、tensorflow2.0_hdf5_savedmodel_pb模型轉換