概述
关于使用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
来源:CSDN
作者:m米咔00
链接:https://blog.csdn.net/liu3612162/article/details/104855691