深度学习系列六：将网络迁移到TensorRT7.0平台

最新推荐文章于 2024-09-13 22:12:42 发布

m米咔00

最新推荐文章于 2024-09-13 22:12:42 发布

阅读量3.3k

点赞数 2

分类专栏：深度学习

本文链接：https://blog.csdn.net/liu3612162/article/details/104855691

版权

深度学习专栏收录该内容

10 篇文章 4 订阅

订阅专栏

概述

关于使用TensorRT对网络进行加速的基本工作在深度学习系列五：使用TensorRT对网络进行加速中描述很详细了，这里主要记录下在迁移到7.0时，顺便解决了之前ResizeNearestNeighbor这个层在uff模型中无法直接支持必须外挂plugin的问题。

Keras模型转ONNX模型

之前使用6.0时，采用的方案是Keras模型转uff模型，再创建engine，其实从官方文档可以看出来，目前TensorRT对onnx支持更好，它对onnx支持的层是最多的，这样容易理解，毕竟onnx可以由各种模型转化得到，因此本次使用onnx模型建立engine，转化可以使用keras-onnx库，安装也很简单：

pip3 install keras2onnx

安装完成后使用api进行转换，非常简单，当然还必须安装tensorflow或者keras：

import keras2onnx
import onnx
# load keras model
model = tf.compat.v1.keras.models.load_model(model_path)

# convert to onnx model
onnx_model = keras2onnx.convert_keras(model, model.name, target_opset=10)
# keras model has dynamic batch size, which will add extra work in generating engine.
# Therefore, convert it to static batch size=1
onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1
# # runtime prediction
output_path = os.path.join(output_path, 'converted.onnx')
onnx.save_model(onnx_model, output_path)

需要注意得是tensorrt和onnx版本必须兼容，这里试了好多版本，最终onnx==1.6和opset==10生成的模型基本没有问题。
使用onnx生成模型输入跟keras模型是一致的，为NHWC，而在uff模型中输入会变成NCHW格式，在之前的工作中input总需要进行额外的transpose操作。

ONNX模型生成engine

有了ONNX模型后，生成engine非常简单：

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(
        network, TRT_LOGGER) as parser:
    builder.max_batch_size = 1
    builder.max_workspace_size = 1 << 28
    if self.infer_mode == 'int8':
        calibration_cache = "model_data/tensorrt_model/yolo_calibration.cache"
        calib = YOLOEntropyCalibrator(self.calib_img_path, calibration_cache, self.model_image_size,
                                      batch_size=8)
        builder.int8_mode = True
        builder.int8_calibrator = calib

    # profile = builder.create_optimization_profile()
    # profile.set_shape('input_1', (1, 416, 416, 3), (1, 416, 416, 3), (1, 416, 416, 3))
    # config.add_optimization_profile(profile)

    onnx_path = common.model_path_to_onnx_path(model_path)
    with open(onnx_path, 'rb') as model:
        if not parser.parse(model.read()):
            raise TypeError("Parser parse failed.")

    engine = builder.build_cuda_engine(network)

    if not engine:
        raise TypeError("Build engine failed.")
    with open(engine_path, 'wb') as f:
        f.write(engine.serialize())

得到engine文件后，加在它进行推理即可：

if os.path.isfile(engine_path):
    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())