【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程

写在前面:Triton镜像加速与生成TensorRT模型所采用的的TensorRT版本需要一致。本实验中,MMDeploy版本为0.8.0,TensorRT版本为8.2.3.0,triton-server版本为22.02-py3,NVIDIA驱动为515,CUDA11.1。

特别注意:TensorRT版本需要与CUDA版本、NVIDIA 驱动版本匹配。

一、MMRotate 模型训练

MMRotate是基于PyTorch的旋转目标检测工具箱,被广泛用在遥感影像的目标检测任务中。官方指路☞:https://github.com/open-mmlab/mmrotate

1. MMRotate安装

MMRotate依赖于PyTorch、MMCV和MMDetection,快速安装如下。具体可参考官方安装教程。

conda create -n open-mmlab python=3.7 pytorch==1.7.0 cudatoolkit=10.1 torchvision -c 
## 可替换自己的pytorch和cuda版本
pytorch -y
conda activate open-mmlab
pip install openmim
mim install mmcv-full
mim install mmdet
git clone https://github.com/open-mmlab/mmrotate.git
cd mmrotate
pip install -r requirements/build.txt
pip install -v -e .

2. MMRotate训练

(1)数据准备

MMRotate支持多种格式的数据,包括DOTA、SSDD、HRSC、HRSID等,也支持自定义数据格式。最便捷的方式就是将数据按已有的格式类型进行组织整理,比如DOTA数据集的格式为:

mmrotate
├── data
│   ├── DOTA
│   │   ├── train
│   │   │   ├──images
│   │   │   │   ├──P00001.png   # 影像文件名与标签文件名一一对应
│   │   │   │   ├──P00002.png
│   │   │   │   ├──...
│   │   │   ├──labelTxt
│   │   │   │   ├──P00001.txt  # 标签文件中一行表示一个目标,由四个角点的坐标值和类别(和difficult字段)组成
│   │   │   │   ├──P00002.txt
│   │   │   │   ├──...
│   │   ├── val
│   │   │   ├──images
│   │   │   │   ├──P10001.png
│   │   │   │   ├──P10002.png
│   │   │   │   ├──...
│   │   │   ├──labelTxt
│   │   │   │   ├──P10001.txt
│   │   │   │   ├──P10002.txt
│   │   │   │   ├──...
│   │   ├── test
│   │   │   ├──images
│   │   │   │   ├──P20001.png
│   │   │   │   ├──P20002.png
│   │   │   │   ├──...
│   │   │   ├──labelTxt
│   │   │   │   ├──P20001.txt
│   │   │   │   ├──P20002.txt
│   │   │   │   ├──...

 (2)config文件准备

config配置文件学习:教程 1: 学习配置文件 — MMDetection 2.25.1 文档

config文件以字典的形式存储模型配置参数,主要包括4个基本组件,即数据集(dataset)、模型(model),训练策略(schedule)和运行时的默认设置(default runtime),可通过修改不同组件实现config文件的设置。

(3)模型训练

## 单卡GPU
python tools/train.py ${CONFIG_FILE} [optional arguments]

## 多卡GPU
./tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [optional arguments]

'''
optional arguments 可设置的参数有:
--work-dir ${WORK_DIR} 权重文件保存路径
--resume-from ${CHECKPOINT_FILE} 从指定的权重文件继续训练
'''

训练过程中会记录模型在验证集上的召回率recall和mAP,可用于评估模型的效果。训练结束后,设置的work_dirs目录下会生成权重文件.pth及对应的配置文件.py,可用于模型的测试和推理。

3. MMRotate推理

MMRotate在demo文件夹中提供了模型推理脚本image_demo.py,可通过输入待预测影像img、配置文件config(即work_dirs目录下的.py)、训练后的权重文件checkpoint(即work_dirs目录下的.pth)等,实现影像的预测推理。

在遥感影像中,通常会进行大幅影像的预测,MMRotate也贴心的提供了可用于大图推理的huge_image_demo.py,通过滑窗以切片的形式实现大图推理。

在实际应用中,mmcv读取几个G的影像还是有一定压力的,而且也无法获取影像的地理信息,因此可以考虑用gdal读取并滑窗处理。

img = gdal.Open(image)
        width = img.RasterXSize
        height = img.RasterYSize
        channel = img.RasterCount
        # prj = img.GetProjection()
        # trans = img.GetGeoTransform()
        # # img_data=img.ReadAsArray(0,0,width,height)
        print('width: ', width, 'height: ', height, 'channel: ', channel)

二、MMDeploy将Torch模型转换为TensorRT

MMDeploy 是 OpenMMLab 的模型部署工具箱,可以为各类算法库提供统一的部署体验,包括MMDetection,MMSegmentation,MMClassification,MMRotate等。官方指路:GitHub - open-mmlab/mmdeploy: OpenMMLab Model Deployment Framework

【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程_第1张图片

1. MMDeploy安装

MMDeploy的安装教程可参考:mmdeploy/get_started.md at master · open-mmlab/mmdeploy · GitHub

 如果只需要实现PyTorch-->ONNX-->TensorRT,安装对应的ONNX和TensorRT部分即可。

重点注意:安装MMDeploy和TensorRT时,MMDeploy的TensorRT、本地安装的TensorRT、Triton Server中的TensorRT需保持一致,且TensorRT与CUDA驱动也有对应关系,务必匹配!

2. 生成TensorRT

在完成MMDeploy的安装后,可以将模型训练时的配置文件和结果权重文件拷贝至MMDeploy中对应文件下来(不拷贝也没关系,设置对应的路径即可)。MMDeploy中config文件夹下的配置文件指定了模型转换时的输入输出、动态参数等,通过调用tools文件夹下的deploy.py就可以实现PyTorch-->ONNX-->TensorRT的转换。

# run the command to start model conversion
python mmdeploy/tools/deploy.py \
    mmdeploy/configs/mmrotate/rotated-detection_tensorrt-fp16_dynamic-320x320-1024x1024.py \   ## 部署的配置文件
    mmrotate/configs/oriented_rcnn/Oriented_rcnn_r101_fpn_1x_dota.py \   # 模型训练时的配置文件
    checkpoints/oriented_rcnn_r101_fpn_1x_dota_epoch_36.pth \    # 模型训练后的权重文件
    mmrotate/demo/demo.jpg \     # 一组待推理的输入数据
    --work-dir mmdeploy_model/oriented-rcnn \    # 结果保存路径
    --device cuda \
    --dump-info

 在模型转换过程中需要指定一组输入的原因: PyTorch 模型到 ONNX 模型本质上就像编译器一样彻底解析原模型的代码,记录所有控制流,但是通常只用 ONNX 记录不考虑控制流的静态图。因此,PyTorch 提供了一种叫做追踪(trace)的模型转换方法:给定一组输入,再实际执行一遍模型,即把这组输入对应的计算图记录下来,保存为 ONNX 格式。export 函数用的就是追踪导出方法,需要给任意一组输入,让模型跑起来。

 转换过程部分代码如下:

【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程_第2张图片

 【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程_第3张图片

 ① TraceWaring:表示在追踪转换的过程中没有记录流信息,大概率是因为该参数不是Tensor类型;

② Successfully created plugins:PyTorch-->ONNX-->TensroRT的转换过程是由一个个算子完成的,当TensorRT工具包中缺少某个算子时,需要进行自定义。这里MMDeploy自定义生成了MMCVMultiLevelRotatedRoiAlign和TRTBatchRotatedNMS两个算子,存储在MMDeploy/build/lib下。

完成TensorRT转换的最终提示信息:

【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程_第4张图片

 最终,在设置的work_dirs文件夹下,会生成七个文件,其中end2end.onnx和end2end.engine文件分别为转换后的ONNX文件和TensorRT文件。

注:.engine文件和.plan文件完全一样,只是后缀不同。

三、Triton部署及推理

NVIDIA Triton推理服务器在生产中提供快速且可拓展的AI。开源推理服务TritonInferenceServer能够从任何框架(TensorFlow、TensorRT、PyTorch、ONNX、自定义等)、在任何基础设备(GPU、CPU)上部署训练后的AI模型,从而简化AI推理过程。

1. Triton Inference Server安装

官方指路:GitHub - triton-inference-server/server: The Triton Inference Server provides an optimized cloud and edge inferencing solution.

 本教程使用的是Triton Docker镜像,因此首先必须安装Nvidia Docker。

(1)拉取镜像

# 为Triton的版本
docker pull nvcr.io/nvidia/tritonserver:-py3

# 例如,拉取 22.02
docker pull nvcr.io/nvidia/tritonserver:22.02-py3

 重点注意:不同Triton版本对cuda驱动的版本需求,以及对应的TensorRT版本。具体可参考Frameworks Support Matrix :: NVIDIA Deep Learning Frameworks Documentation

(2)创建模型存储库

将代码克隆到本地并创建示例模型的存储库。

## 克隆代码仓库并创建示例模型存储器
git clone  https://github.com/triton-inference-server/server.git

cd server/docs/examples

./fetch_models.sh    # 下载示例模型

(3)运行Triton 的示例服务

利用下载好的示例模型运行Triton服务,这里-v 挂在的就是下载本地的示例模型文件夹。

docker run --gpus all --rm --net=host -p8000:8000 -p8001:8001 -p8002:8002 -v full/path/to/triton_server/model_repository:/models nvcr.io/nvidia/tritonserver:22.02-py3 tritonserver --model-repository=/models

# --rm表示关闭容器时删除容器

成功部署的话,会看到模型的运行状态为READY,对外提供服务的端口已开启,如下所示:

【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程_第5张图片

  (4)验证Triton 服务

可另起一个终端,通过curl命令验证Triton服务是否正常运行。

curl -v localhost:8000/v2/health/ready

 【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程_第6张图片

2. 生成自己的模型

上述运行Triton示例服务时,采用的是下载好的预设模型,这里我们需要根据MMDeploy生成的TenorRT配置自己的模型文件。

(1)准备TensorRT文件

新建一个models文件夹用于存储自己的模型,文件夹结构如下,可将MMDeploy生成的engine文件复制到'1'文件夹下并重命名为model.engine。

小Tips:可以在models文件夹外面新建一个plugins文件夹,把MMDeploy新生成的算子.so文件放进去,后面会用上。

  /     # 新建的存储模型的文件夹,如../triton_server/models
    /              # 模型名称,如OrientedRCNN_trt
      config.pbtxt             # 模型的配置文件
      1/                       # 版本号,默认为1
        model.engine             # MMDeploy生成的Triton文件,即end2ned.engine,可将其重命名为model.engine
 

(2)模型配置文件

 文件夹下的配置文件config.pbtxt如下:

name: "OrientedRCNN_trt"     # 模型名,也是文件夹的目录名
platform: "tensorrt_plan"    # 模型对应的平台,参考文章下面给出的表格
max_batch_size : 8           # 一次送入模型的最大batch_size。
input [
  {
    name: "input"
    data_type: TYPE_FP32
    dims: [ 3,-1,-1 ]             # 第一个维度默认是batch size,不用配置,直接从第二个维度开始。如果是可变维度,就用 -1
  }
]
output [
  {
    name: "dets"                  # 输出目标框的节点名称
    data_type: TYPE_FP32
    dims: [-1,6]
  },
  {
    name: "labels"                # 输出类别的节点名称
    data_type: TYPE_INT32
    dims: [ -1 ]
  }
]

default_model_filename: "model.engine"               # TensorRT的文件名

其中:

name为模型名,需要与文件夹名称一致;

platform是模型对应的平台,不同模型的存储格式都有对应的值。具体如下表。

input、output包括输入输出的节点名称、数据类型、数据维度。

不同框架对应的platform
框架 platform
TensorRT tensorrt_plan
TensorFlow SavedModel tensorflow_savedmodel
TensorFlow GraphDef tensorflow_graphdef
ONNX onnxruntime_onnx
Torch pytorch_libtorch

(3)启动服务

启动自己的服务时,多了两个参数。一个是-v中的plugins,这个对应的是MMDeploy自定义的算子文件夹,另一个是--env,这个是把自定义的算子文件libmmdeploy_tensorrt_ops.so添加到环境变量中,使容器能够调用这些算子,否则会因为缺失算子而报错。

启动服务命令如下:

docker run --gpus all --rm --net=host -p8000:8000 -p8001:8001 -p8002:8002 -v full/path/to/triton_server/models:/models -v full/path/to/triton_server/plugins:/plugins --env LD_PRELOAD=/plugins/libmmdeploy_tensorrt_ops.so nvcr.io/nvidia/tritonserver:22.02-py3 tritonserver --model-repository=/models

 模型正常运行的READY:

【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程_第7张图片

再一次看到服务端口已开启:

3. Triton Client推理

Triton服务启动后,可通过Triton Client来进行模型的推理请求,并且提供了http和grpc两种协议。下面是Python的请求示例。

import os
import time

import gevent.ssl
import numpy as np
import tritonclient.http as httpclient
from PIL import Image
from postprocess import *


def client_init(url="localhost:8000",
                ssl=False, key_file=None, cert_file=None, ca_certs=None, insecure=False,
                verbose=False):

    if ssl:
        ssl_options = {}
        if key_file is not None:
            ssl_options['keyfile'] = key_file
        if cert_file is not None:
            ssl_options['certfile'] = cert_file
        if ca_certs is not None:
            ssl_options['ca_certs'] = ca_certs
        ssl_context_factory = None
        if insecure:
            ssl_context_factory = gevent.ssl._create_unverified_context
        triton_client = httpclient.InferenceServerClient(
            url=url,
            verbose=verbose,
            ssl=True,
            ssl_options=ssl_options,
            insecure=insecure,
            ssl_context_factory=ssl_context_factory)
    else:
        triton_client = httpclient.InferenceServerClient(
            url=url, verbose=verbose)

    return triton_client


def infer(triton_client, model_name='OrientedRCNN_trt',
          input='input', output0='dets', output1='labels',
          request_compression_algorithm=None,
          response_compression_algorithm=None):
    """
    :param triton_client:启动的服务
    :param model_name:模型名称
    :param input:输入节点名
    :param output0:输出节点名1
    :param output1:输出节点名2
    :return:推理结果
    """
    inputs = []
    outputs = []
    # batch_size=8
    # 如果batch_size超过配置文件的max_batch_size,infer则会报错

    inputs.append(httpclient.InferInput(input, [1, 3, 1024, 1024], "FP32"))

    # 给输入的数据赋值
    root_dir = os.getcwd()
    img_path = os.path.join(root_dir, 'data/test.tif')  # 待推理的文件路径
    img = np.array(Image.open(img_path))
    img = img.astype(np.float32)
    img = img.transpose((2, 0, 1))

    img = np.expand_dims(img, axis=0)  # (1, 3, 1024, 1024)
    inputs[0].set_data_from_numpy(img)

    outputs.append(httpclient.InferRequestedOutput(output0, binary_data=False))
    outputs.append(httpclient.InferRequestedOutput(output1,
                                                   binary_data=False))
    results = triton_client.infer(
        model_name=model_name,
        inputs=inputs,
        outputs=outputs,
        request_compression_algorithm=request_compression_algorithm,
        response_compression_algorithm=response_compression_algorithm)

    # 将结果转化为numpy格式
    result_dets = results.as_numpy(output0)
    result_labels = results.as_numpy(output1)

    print(result_dets)
    print(result_labels)


if __name__ == '__main__':
    client = client_init()
    infer(triton_client=client, model_name='OrientedRCNN_trt')

输出的result_dets和result_labels分别存储着目标的位置信息和类别信息。至此,即完成了Triton的部署及推理。

四、踩的一些坑

1. 驱动太low

最开始安装MMDeploy的时候采用了8.2.3.0版本的TensorRT,因此本地环境配置也采用了TensorRT  8.2.3.0,triton用的最新22.09(随便装的。。。)。然后发现示例模型可以启动,但自己的模型就怎么都搞不定,广读度娘,终于发现了bug:①我的工作站驱动是450以下,但是Triton-22.09要求驱动高于450,卒!② Triton-22.09对应的TensorRT版本是8.5.0.12,模型文件对应的TensorRT版本是8.2.3.0,Double Kill❗

2. 自定义算子

MMDeploy在生成TensorRT文件时会自定义算子,但是Triton Inference Server并不存在该算子,因此需要将生成的自定义算子作为环境变量导入镜像中,也就是启动服务时--env中的libmmdeploy_tensorrt_ops.so。否则会有如下报错:

【Triton部署目标检测模型】Linux环境 MMRotate+MMDeploy+Triton 目标检测模型训练、部署及推理最全教程_第8张图片

 小Tips:启动服务时的界面中(如上图所示),I开头的是提示信息,E开头的就是报错信息。

五、总结

MMRotate+MMdeploy+Triton这一套流程前后整了十来天,从安装到部署也踩了不少坑,有的没有记录,还在折腾这么久总算可用了。BUT!!!我的ONNX还没启动成功,继续跳坑中。。。。

你可能感兴趣的:(pytorch,深度学习,目标检测)