Shortcuts

教程 9: ONNX 到 TensorRT 的模型转换(实验性支持)

如何将模型从 ONNX 转换为 TensorRT

先决条件

  1. 请参考 get_started.md 从源码安装 MMCV 和 MMDetection。

  2. 请参考 ONNXRuntime in mmcvTensorRT plugin in mmcv 安装支持 ONNXRuntime 自定义操作和 TensorRT 插件的 mmcv-full

  3. 使用工具 pytorch2onnx 将模型从 PyTorch 转换为 ONNX。

用法

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 \

所有参数的说明:

  • config: 模型配置文件的路径。

  • model: ONNX 模型文件的路径。

  • --trt-file: 输出 TensorRT 引擎文件的路径。如果未指定,它将被设置为 tmp.trt

  • --input-img: 用于追踪和转换的输入图像的路径。默认情况下,它将设置为 demo/demo.jpg

  • --shape: 模型输入的高度和宽度。如果未指定,它将设置为 400 600

  • --min-shape: 模型输入的最小高度和宽度。如果未指定,它将被设置为与 --shape 相同。

  • --max-shape: 模型输入的最大高度和宽度。如果未指定,它将被设置为与 --shape 相同。

  • --workspace-size: 构建 TensorRT 引擎所需的 GPU 工作空间大小(以 GiB 为单位)。如果未指定,它将设置为 1 GiB。

  • --show: 确定是否显示模型的输出。如果未指定,它将设置为 False

  • --verify: 确定是否在 ONNXRuntime 和 TensorRT 之间验证模型的正确性。如果未指定,它将设置为 False

  • --verbose: 确定是否打印日志消息。它对调试很有用。如果未指定,它将设置为 False

例子:

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/deplopyment/test.py 来评估 TensorRT 模型。

请参阅以下链接以获取更多信息。

  • 如何评估导出的模型

  • 结果和模型

支持转换为 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

注意:

  • 以上所有模型通过 Pytorch==1.6.0, onnx==1.7.0 与 TensorRT-7.2.1.6.Ubuntu-16.04.x86_64-gnu.cuda-10.2.cudnn8.0 测试

提醒

  • 如果您在上面列出的模型中遇到任何问题,请创建 issue,我们会尽快处理。对于未包含在列表中的模型,由于资源有限,我们可能无法在此提供太多帮助。请尝试深入挖掘并自行调试。

  • 由于此功能是实验性的,并且可能会快速更改,因此请始终尝试使用最新的 mmcvmmdetecion

常见问题

Read the Docs v: stable
Versions
latest
stable
v2.25.0
v2.24.1
v2.24.0
v2.23.0
v2.22.0
v2.21.0
v2.20.0
v2.19.1
v2.19.0
v2.18.1
v2.18.0
v2.17.0
v2.16.0
v2.15.1
v2.15.0
v2.14.0
Downloads
pdf
html
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.