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

时光怂恿深爱的人放手 提交于 2020-03-17 01:45:59

概述

关于使用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())

最终得到的三种模型map为:

Model mode dataset MAP MAP (0.5) MAP(0.75)
Yolov3-416 raw COCOval2014 0.315 0.561 0.319
Yolov3-416 fp32 COCOval2014 0.315 0.561 0.319
Yolov3-416 int8 COCOval2014 0.304 0.551 0.295

所有代码都已经开源,具体推理使用方法见Yolov3-TensorRT-py

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!