MMDeploy学习笔记（三）：MMDetection模型部署

MMDetection模型部署

本博客以YOLOX模型为例，演示了如何利用MMDeploy进行模型的部署。首先需要下载并配置MMDetection和MMDeploy的环境，使用预先训练好的模型参数，使用MMDeploy将MMDetection中的Pytorch模型转换成ONNX模型，最后使用ONNXRunTime进行推理，简要过程如下图所示：

更多关于模型部署和Pytorch_to_ONNX的相关知识点，可以参考博客MMDeploy学习笔记（一）：模型部署基础知识

环境安装与配置

博主自己的硬件配置为（供参考）：

• Linux（Ubuntu 20.04、X86_64）
• NVIDIA GeForce RTX 3090 显卡
• NVIDIA显卡驱动版本：12.0
• CUDA version：11.3

关于CUDA、Cudnn等安装方法可参考博主的另一篇博客Ubuntu18.04从零开始安装显卡驱动、配置MMDetection3D环境

MMDetection环境配置与安装：

# 创建conda环境
conda create -n mmdetection python=3.8 -y
conda activate mmdetection
# 安装pytorch
conda install pytorch==1.10.0 torchvision==0.11.0 torchaudio==0.10.0 cudatoolkit=11.3 -c pytorch -c conda-forge
# 安装mmcv
pip install mmcv-full==2.1.0 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.10.0/index.html
# 安装mmdetection
git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection
pip install -r requirements/build.txt
pip install -v -e .  # or "python setup.py develop"

关于MMDeploy的安装方法可以参考博主的该系列的上一篇博客MMDeploy学习笔记（二）：MMDeploy简介与环境配置

模型转换与推理

此节将介绍如何用MMDeploy将Pytorch模型转换为ONNX模型，并使用ONNXRunTime进行推理，步骤如下：

1. 下载模型权重文件

mkdir mmdetection/checkpoint
cd mmdetection/checkpoint/
# 下载 yolox 模型权重文件
wget https://download.openmmlab.com/mmdetection/v2.0/yolox/yolox_tiny_8x8_300e_coco/yolox_tiny_8x8_300e_coco_20211124_171234-b4047906.pth 

# 下载 retinanet 模型权重文件
# wget https://download.openmmlab.com/mmdetection/v2.0/retinanet/retinanet_r18_fpn_1x_coco/retinanet_r18_fpn_1x_coco_20220407_171055-614fd399.pth

# 下载 Faster R-CNN 模型权重文件
# wget https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth

2. Pytorch模型转ONNX模型

python ./tools/deploy.py \
    ${DEPLOY_CFG_PATH} \
    ${MODEL_CFG_PATH} \
    ${MODEL_CHECKPOINT_PATH} \
    ${INPUT_IMG} \
    --test-img ${TEST_IMG} \
    --work-dir ${WORK_DIR} \
    --calib-dataset-cfg ${CALIB_DATA_CFG} \
    --device ${DEVICE} \
    --log-level INFO \
    --show \
    --dump-info

参数描述：

• deploy_cfg : mmdeploy 针对此模型的部署配置，包含推理框架类型、是否量化、输入 shape 是否动态等。配置文件之间可能有引用关系，configs/mmpretrain/classification_ncnn_static.py 是一个示例。
• model_cfg : mm 算法库的模型配置，例如 mmpretrain/configs/vision_transformer/vit-base-p32_ft-64xb64_in1k-384.py，与 mmdeploy 的路径无关。
• checkpoint : torch 模型路径。可以 http/https 开头，详见 mmcv.FileClient 的实现。
• img : 模型转换时，用做测试的图像或点云文件路径。
• –test-img : 用于测试模型的图像文件路径。默认设置成None。
• –work-dir : 工作目录，用来保存日志和模型文件。
• –calib-dataset-cfg : 此参数只有int8模式下生效，用于校准数据集配置文件。若在int8模式下未传入参数，则会自动使用模型配置文件中的’val’数据集进行校准。
• –device : 用于模型转换的设备。 默认是cpu，对于 trt 可使用 cuda:0 这种形式。
• –log-level : 设置日记的等级，选项包括’CRITICAL’， ‘FATAL’， ‘ERROR’， ‘WARN’， ‘WARNING’， ‘INFO’， ‘DEBUG’， ‘NOTSET’。 默认是INFO。
• –show : 是否显示检测的结果。
• –dump-info : 是否输出 SDK 信息。

对YOLOX模型进行转换：

python mmdeploy/tools/deploy.py \
        mmdeploy/configs/mmdet/detection/detection_onnxruntime_dynamic.py \
        mmdetection/configs/yolox/yolox_tiny_8x8_300e_coco.py \
        mmdetection/checkpoint/yolox_tiny_8x8_300e_coco_20211124_171234-b4047906.pth \
        mmdetection/demo/demo.jpg \
        --work-dir mmdeploy_models/mmdetection/yolox \
        --device cpu

3. 使用ONNXRunTime推理

方式一：在命令行shell中运行

python
import cv2
from mmdeploy_python import Detector
img = cv2.imread('mmdetection/demo/demo.jpg')
onnx_path = 'mmdeploy_models/mmdetection/yolox'
detector = Detector(onnx_path, 'cpu')
result = detector(img)

在命令行可以直接查看结果为bbox、类别以及对应的得分。

方式二：直接运行命令行推理脚本

SPDLOG_LEVEL=error python mmdeploy/demo/python/object_detection.py cpu mmdeploy_models/mmdetection/yolox mmdetection/demo/demo.jpg

运行成功后，在本地会生成output_detection.png目标检测预测结果图像文件。