Skip to main content

Model Convertation

Pytorch to ONNX​

前置步骤​

  1. 按照 get_started.md/Prepare environment 配置环境。
  2. 为 ONNX Runtime 构建自定义操作符，并手动安装 MMCV，具体步骤访问链接
  3. 按照 get_started.md/Install MMdetection 中的步骤 2-3 手动安装 MMdetection。

使用​

python tools/deployment/pytorch2onnx.py \
    ${CONFIG_FILE} \
    ${CHECKPOINT_FILE} \
    --output-file ${OUTPUT_FILE} \
    --input-img ${INPUT_IMAGE_PATH} \
    --shape ${IMAGE_SHAPE} \
    --test-img ${TEST_IMAGE_PATH} \
    --opset-version ${OPSET_VERSION} \
    --cfg-options ${CFG_OPTIONS}
    --dynamic-export \
    --show \
    --verify \
    --simplify \

所有参数的描述

  • config : 模型配置文件的路径。
  • checkpoint : 模型检查点文件的路径。
  • --output-file：输出 ONNX 模型的路径。如果不指定，它将被设置为 tmp.onnx。
  • --input-img：用于跟踪和转换的输入图像的路径。默认情况下，它将被设置为 test/data/color.jpg。
  • --shape：输入到模型的张量的高度和宽度。如果没有指定，它将被设置为 800 1216。
  • --test-img：验证输出的 ONNX 模型的图像路径。默认情况下，它将被设置为无，这意味着它将使用 --input-img 进行验证。
  • --opset-version : ONNX 的 opset 版本。如果没有指定，它将被设置为11。
  • --dynamic-export：决定是否导出具有动态输入和输出形状的 ONNX 模型。如果不指定，它将被设置为 "False"。
  • --show：决定是否打印导出模型的结构，以及当 --verify 被设置为 True 时是否显示检测输出。如果不指定，它将被设置为 False。
  • --verify：决定是否验证导出模型的正确性。如果不指定，它将被设置为False。
  • --simplify：决定是否简化导出的 ONNX 模型。如果不指定，它将被设置为 False。
  • --cfg-options：覆盖所使用的配置文件中的一些设置，xxx=yyy 格式的键值对将被合并到配置文件中。
  • --skip-postprocess：决定是否在没有后期处理的情况下导出模型。如果没有指定，它将被设置为False。注意。这是一个实验性选项。只对某些单阶段模型起作用。用户需要自己实现后处理。我们不保证导出的模型的正确性。

样例：

python tools/deployment/pytorch2onnx.py \
    configs/yolo/yolov3_d53_mstrain-608_273e_coco.py \
    checkpoints/yolo/yolov3_d53_mstrain-608_273e_coco.pth \
    --output-file checkpoints/yolo/yolov3_d53_mstrain-608_273e_coco.onnx \
    --input-img demo/demo.jpg \
    --test-img tests/data/color.jpg \
    --shape 608 608 \
    --show \
    --verify \
    --dynamic-export \
    --cfg-options \
      model.test_cfg.deploy_nms_pre=-1 \

测试导出的模型​

我们准备了一个 tools/deplopyment/test.py，用 ONNXRuntime 和 TensorRT 评估 ONNX 模型。

首先需要安装 ONNXRuntime

  • CPU

    pip install onnx onnxruntime

  • GPU

    pip install onnxruntime-gpu
注意
  • onnxruntime-gpu 依赖于 CUDA 和 CUDNN 的版本，请确保你的环境符合要求。
  • 为 ONNXRuntime 建立自定义操作符如下链接
  • 安装 TensorRT 链接

使用方式：

python tools/deployment/test.py \
    ${CONFIG_FILE} \
    ${MODEL_FILE} \
    --out ${OUTPUT_FILE} \
    --backend ${BACKEND} \
    --format-only ${FORMAT_ONLY} \
    --eval ${EVALUATION_METRICS} \
    --show-dir ${SHOW_DIRECTORY} \
    ----show-score-thr ${SHOW_SCORE_THRESHOLD} \
    ----cfg-options ${CFG_OPTIONS} \
    ----eval-options ${EVALUATION_OPTIONS} \

所有参数的描述

  • config：模型配置文件的路径。
  • model：输入模型文件的路径。
  • --out：pickle 格式的输出结果文件的路径。
  • --backend：输入模型运行的后端，应该是 onnxruntime 或 tensorrt。
  • --format-only：格式化输出结果，不进行评估。当你想将结果格式化为特定的格式并提交给测试服务器时，它很有用。如果不指定，它将被设置为 False。
  • --eval：评价指标，这取决于数据集，例如，COCO 的 "bbox"、"segm"、"proposal"，PASCAL VOC 的 "mAP"、"recall"。
  • --show-dir：绘制的图像将被保存的目录
  • --show-score-thr：分数阈值。默认设置为0.3。
  • --cfg-options：覆盖所使用的配置文件中的一些设置，xxx=yyy 格式的键值对将被合并到配置文件中。
  • --eval-options：用于评估的自定义选项，xxx=yyy 格式的键值对将成为 dataset.evalue() 函数的 kwargs。
注意

如果部署的后端平台是TensorRT，请在运行该文件前添加环境变量。

export ONNX_BACKEND=MMCVTensorRT

如果你想在 TensorRT 后端使用 --dynamic-export 参数来导出 ONNX，请删除 --simplify 参数，反之亦然。

ONNX to TensorRT​

前置条件​

  1. 请参考 get_started.md，安装 MMCV 和 MMDetection。
  2. 请参考 mmcv 中的 ONNXRuntime 和 mmcv 中的 TensorRT 插件 来安装带有 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.rt。
  • --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 \

测试导出的模型​

见上面