MMDetection系列之 ONNX TO TENSORRT
如何转换模型从ONNX到TensorRT
usage
python tools/deployment/onnx2tensorrt.py \
${CONFIG} \
${MODEL} \
--trt-file ${TRT_FILE} \
--input-img ${INPUT_IMAGE_PATH} \
--shape ${INPUT_IMAGE_SHAPE} \
--min-shape ${MIN_IMAGE_SHAPE} \
--max-shape ${MAX_IMAGE_SHAPE} \
--workspace-size {WORKSPACE_SIZE} \
--show \
--verify \
Description of all arguments:
config : The path of a model config file.
model : The path of an ONNX model file.
–trt-file: The Path of output TensorRT engine file. If not specified, it will be set to tmp.trt.
–input-img : The path of an input image for tracing and conversion. By default, it will be set to demo/demo.jpg.
–shape: The height and width of model input. If not specified, it will be set to 400 600.
–min-shape: The minimum height and width of model input. If not specified, it will be set to the same as --shape.
–max-shape: The maximum height and width of model input. If not specified, it will be set to the same as --shape.
–workspace-size : The required GPU workspace size in GiB to build TensorRT engine. If not specified, it will be set to 1 GiB.
–show: Determines whether to show the outputs of the model. If not specified, it will be set to False.
–verify: Determines whether to verify the correctness of models between ONNXRuntime and TensorRT. If not specified, it will be set to False.
–verbose: Determines whether to print logging messages. It’s useful for debugging. If not specified, it will be set to False.
Example
python tools/deployment/onnx2tensorrt.py \
configs/retinanet/retinanet_r50_fpn_1x_coco.py \
checkpoints/retinanet_r50_fpn_1x_coco.onnx \
--trt-file checkpoints/retinanet_r50_fpn_1x_coco.trt \
--input-img demo/demo.jpg \
--shape 400 600 \
--show \
--verify \
如何评价导出的模型
我们准备了一个工具tools/ deploy /test.py来评估TensorRT模型。
可转换为TensorRT的支持模型列表
Model Config Dynamic Shape Batch Inference Note
SSD configs/ssd/ssd300_coco.py Y Y
FSAF configs/fsaf/fsaf_r50_fpn_1x_coco.py Y Y
FCOS configs/fcos/fcos_r50_caffe_fpn_4x4_1x_coco.py Y Y
YOLOv3 configs/yolo/yolov3_d53_mstrain-608_273e_coco.py Y Y
RetinaNet configs/retinanet/retinanet_r50_fpn_1x_coco.py Y Y
Faster R-CNN configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py Y Y
Cascade R-CNN configs/cascade_rcnn/cascade_rcnn_r50_fpn_1x_coco.py Y Y
Mask R-CNN configs/mask_rcnn/mask_rcnn_r50_fpn_1x_coco.py Y Y
Cascade Mask R-CNN configs/cascade_rcnn/cascade_mask_rcnn_r50_fpn_1x_coco.py Y Y
PointRend configs/point_rend/point_rend_r50_caffe_fpn_mstrain_1x_coco.py Y Y