windows环境下使用mmdetection+mmdeploy训练自定义数据集并转成onnx格式部署

实验环境

• windows10
• python：3.8
• pytorch :1.8.1
• cuda：11.1
• mmdet：3.1.0
• mmcv：2.0.1
• onnxruntime-gpu：1.9.0

安装conda

conda下载地址：
https://mirrors.tuna.tsinghua.edu.cn/anaconda/archive/Anaconda3-5.2.0-Windows-x86_64.exe

创建虚拟环境

conda create -n openmmlab python=3.8
conda activate openmmlab

安装pytorch

pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

使用 MIM 安装 MMEngine 和 MMCV

pip install -U openmim
mim install mmengine
mim install "mmcv>=2.0.0"

安装 MMDetection

在github上下载MMDetection：

git clone git@github.com:open-mmlab/mmdetection.git

如果网络不好也可到gitee下载：

git clone git@gitee.com:open-mmlab/mmdetection.git

或者可以直接下载zip压缩包：

准备自定义数据集

此处使用VOC数据集进行演示，使用的框架为Faster RCNN。VOC数据集的格式如下：

└─VOCdevkit
    └─VOC2007
        ├─Annotations
        ├─ImageSets
        │  └─Main
                 └─test.txt
                 └─train.txt
                 └─val.txt
        └─JPEGImages

其中Annotations存放xml文件，JPEGImages存放jpg文件。

修改配置信息

复制configs/_base_/datasets/voc0712.py，命名为configs/_base_/datasets/vocengine.py。（该名称可自定义）
修改数据集路径：

复制configs/faster_rcnn/faster-rcnn_r50_fpn_1x_coco.py，命名为configs/faster_rcnn/faster-rcnn_r50_fpn_1x_voc.py，修改信息如下：

修改mmdet/evaluation/functional/class_names.py和mmdet/datasets/voc.py的类别信息：

修改configs/_base_/models/faster-rcnn_r50_fpn.py的类别数：

修改epochs：
configs/_base_/schedules/schedule_1x.py

开始训练

python tools/train.py configs/faster_rcnn/faster-rcnn_r50_fpn_1x_voc.py

训练结束后将在work_dir下得到训练的相关信息以及模型权重文件。

模型转换与推理

下载mmdeploy：

git clone [email protected]:open-mmlab/mmdeploy.git
or
git clone [email protected]:open-mmlab/mmdeploy.git

安装相关库：

pip install mmdeploy==1.0.0
pip install mmdeploy-runtime-gpu==1.0.0
pip install onnxruntime-gpu==1.9.0

将训练得到的模型以及py文件复制到mmdeploy/faster_rcnn_configs/目录下，并设置结果路径为：faster_rcnn_onnx/。（这两个目录可以自定义）

新建inference.py，用以下代码进行模型转换以及推理：

from mmdeploy.apis import torch2onnx
from mmdeploy.apis.utils import build_task_processor
from mmdeploy.utils import get_input_shape, load_config
import torch


# 假设在安装步骤中，mmdetection 和 mmdeploy 代码库在同级目录下，并且当前的工作目录为 mmdetection 的根目录，
# 那么以 Faster R-CNN 模型为例，你可以从此处下载对应的 checkpoint，并使用以下代码将之转换为 onnx 模型：
def pth_to_onnx(args):
    """
    将pth模型转换为onnx模型
    Args:
        args: 

    Returns:

    """
    # 1. convert model to onnx
    torch2onnx(args.image, args.work_dir, args.save_file, args.deploy_cfg, args.model_cfg,
               args.model_checkpoint, args.device)


def onnx_inference(args):
    """
    使用onnx模型进行推理
    Args:
        args: 

    Returns:

    """
    # read deploy_cfg and model_cfg
    deploy_cfg, model_cfg = load_config(args.deploy_cfg, args.model_cfg)

    # build task and backend model
    task_processor = build_task_processor(model_cfg, deploy_cfg, args.device)
    model = task_processor.build_backend_model(args.backend_model)

    # process input image
    input_shape = get_input_shape(deploy_cfg)
    model_inputs, _ = task_processor.create_input(args.image, input_shape)

    # do model inference
    with torch.no_grad():
        result = model.test_step(model_inputs)

    # visualize results
    task_processor.visualize(
        image=args.image,
        model=model,
        result=result[0],
        window_name='visualize',
        output_file='output_detection.png')


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
        description=__doc__)

    # 训练设备类型
    parser.add_argument('--device', default='cuda', help='device')
    # 图像
    parser.add_argument('--image', default='test.jpg')
    # 转换模型后的输出文件夹
    parser.add_argument('--work-dir', default='mmdeploy_models/mmdet/onnx')
    # 保存的模型名称
    parser.add_argument('--save-file', default='end2end.onnx')
    # mmdeploy的配置文件，使用目标检测应选择detection
    parser.add_argument('--deploy-cfg', default='configs/mmdet/detection/detection_onnxruntime_static.py')
    # 训练得到的配置文件
    parser.add_argument('--model-cfg', default='faster_rcnn_configs/faster-rcnn_r50_fpn_1x_voc.py')
    # 训练得到的配置文件
    parser.add_argument('--model-checkpoint', default='')
    # 转换后的onnx模型
    parser.add_argument('--backend_model', default=['mmdeploy_models/mmdet/onnx/end2end.onnx'], type=list)
    args = parser.parse_args()
    print(args)
    pth_to_onnx(args)
    onnx_inference(args)

其中的参数有：

• –device:使用设备，可选择cuda or cpu
• –image:测试使用的图片，可随意选择一张
• –work-dir:转换模型后的输出文件夹
• –save-file:保存的模型名称
• –deploy-cfg:mmdeploy的配置文件，使用目标检测应选择detection
• –model-checkpoint:训练得到的配置文件
• –backend_model:转换后的onnx模型

然后，输入如下命令行（可根据实际需求修改）：

python inference.py --model-checkpoint faster_rcnn_configs/epoch_1.pth --image test.jpg --device cuda

最后，将在根目录下得到如下onnx模型：

以及output_detection.png预测结果文件：