【mmdeploy】mmseg转ONNX/TensorRT,附推理代码
创始人
2024-05-13 03:00:51
0

目录

1.关于mmdeploy 

2.环境安装

2.1预编译安装(Linux-x86_64, CUDA 11.x, TensorRT 8.2.3.0):

2.2预编译安装(Linux-x86_64, CUDA 11.x, ONNX):

3.mmseg工程torch转onnx

4.使用onnx模型文件进行推理


1.关于mmdeploy 

MMDeploy 是 OpenMMLab 模型部署工具箱,为各算法库提供统一的部署体验。基于 MMDeploy,开发者可以轻松从训练 repo 生成指定硬件所需 SDK,省去大量适配时间。MMDeploy 提供了一系列工具,帮助您更轻松的将 OpenMMLab 下的算法部署到各种设备与平台上。

2.环境安装

提示:不要使用mmseg工程里的deploy工具,建议使用openmmlab下的mmdeploy工程(血泪教训)。

既然是要推理部署,大概率是已经训练好模型了,基础的openmmlab环境包(mmsegmentation,mmcv-full......)应该已经安装过了,因此官网中的步骤(安装mmseg、mmcv)可以不用操作了,只需要激活本地的训练时的环境(比如:openmmlab),再安装缺少的包就可以了。

需要安装的包有:mmdeploy\onnx\tensorrt\pycuda。。。

预编译好的平台和设备如下,不在的请自行下载源码编译安装。

2.1预编译安装(Linux-x86_64, CUDA 11.x, TensorRT 8.2.3.0):

wget https://github.com/open-mmlab/mmdeploy/releases/download/v0.12.0/mmdeploy-0.12.0-linux-x86_64-cuda11.1-tensorrt8.2.3.0.tar.gz
tar -zxvf mmdeploy-0.12.0-linux-x86_64-cuda11.1-tensorrt8.2.3.0.tar.gz
cd mmdeploy-0.12.0-linux-x86_64-cuda11.1-tensorrt8.2.3.0
pip install dist/mmdeploy-0.12.0-py3-none-linux_x86_64.whl
pip install sdk/python/mmdeploy_python-0.12.0-cp38-none-linux_x86_64.whl
cd ..
# 安装推理引擎 TensorRT
# !!! 从 NVIDIA 官网下载 TensorRT-8.2.3.0 CUDA 11.x 安装包并解压到当前目录
pip install TensorRT-8.2.3.0/python/tensorrt-8.2.3.0-cp38-none-linux_x86_64.whl
pip install pycuda
export TENSORRT_DIR=$(pwd)/TensorRT-8.2.3.0
export LD_LIBRARY_PATH=${TENSORRT_DIR}/lib:$LD_LIBRARY_PATH
# !!! 从 NVIDIA 官网下载 cuDNN 8.2.1 CUDA 11.x 安装包并解压到当前目录
export CUDNN_DIR=$(pwd)/cuda
export LD_LIBRARY_PATH=$CUDNN_DIR/lib64:$LD_LIBRARY_PATH

2.2预编译安装(Linux-x86_64, CUDA 11.x, ONNX):

wget https://github.com/open-mmlab/mmdeploy/releases/download/v0.12.0/mmdeploy-0.12.0-linux-x86_64-onnxruntime1.8.1.tar.gz
tar -zxvf mmdeploy-0.12.0-linux-x86_64-onnxruntime1.8.1.tar.gz
cd mmdeploy-0.12.0-linux-x86_64-onnxruntime1.8.1
pip install dist/mmdeploy-0.12.0-py3-none-linux_x86_64.whl
pip install sdk/python/mmdeploy_python-0.12.0-cp3X-none-linux_x86_64.whl  #注意自己python版本对应
cd ..
# 安装推理引擎 ONNX
pip install onnxruntime-gpu==1.8.1  #安装GPU版本
pip install onnxruntime==1.8.1       #安装CPU版本wget https://github.com/microsoft/onnxruntime/releases/download/v1.8.1/onnxruntime-linux-x64-1.8.1.tgz
tar -zxvf onnxruntime-linux-x64-1.8.1.tgz
export ONNXRUNTIME_DIR=$(pwd)/onnxruntime-linux-x64-1.8.1
export LD_LIBRARY_PATH=$ONNXRUNTIME_DIR/lib:$LD_LIBRARY_PATH

3.mmseg工程torch转onnx

在准备工作就绪后,我们可以使用 MMDeploy 中的工具 tools/deploy.py,将 OpenMMLab 的 PyTorch 模型转换成推理后端支持的格式。

python tools/deploy.py \$DEPLOY_CFG \$MODEL_CFG \$PTH_MODEL_PATH \--work-dir $OUT_PATH \--show --device cuda --dump-info

其中:

DEPLOY_PATH为mmdeploy工程下./mmedeploy/configs/mmseg/XXX.py的config文件路径

MODEL_CFG为自己训练时的config文件,一般在pth模型同目录下就有

PTH_MODEL_PATH为需要转的pth模型文件地址

OUT_PATH为输出的onnx模型文件地址和对应的json存放地址

举个栗子:

python tools/deploy.py \/root/workspace/mmdeploy/configs/mmseg/segmentation_onnxruntime_dynamic.py \/root/workspace/mmseg/MMSEG_DEPLOY/uper_swin_base.py \/root/workspace/mmseg/MMSEG_DEPLOY/best_mIoU_epoch_500.pth \--work-dir /root/workspace/mmseg/MMSEG_DEPLOY/OUTPUT \--show --device cuda --dump-info

紧接着就等待去结果路径中寻找生成的onnx模型文件吧!

4.使用onnx模型文件进行推理

推理过程比较简单,可以自行花式修改

from mmdeploy.apis import inference_modelresult = inference_model(model_cfg = 'data/project/mmseg/weights/ersi_uper_swin_base.py',deploy_cfg = '/data/project/mmdeploy-master/configs/mmseg/segmentation_onnxruntime_dynamic.py',backend_files = ['/data/project/end2end.onnx'],img ='in_img_path/xxx.png',device='cuda:0')

其中:

model_cfg为训练时的config文件

deploy_cfg为mmdeploy工程下的/configs/mmseg/XXX.py文件

backend_file为刚才转出来的onnx模型地址

img就是要推理的原始影像

注意:此处输出的result为list格式,如果需要保存结果的话,需要先转成numpy再取对应的维度,最后保存成影像格式。或者参考mmseg/tools/test.py的处理方法,把所有影像预测结果append到一个大的list中,最后依次保存。

pred = (np.array(result)[0,:,:]).astype(np.uint8)  #将result保存成一张影像
cv2.imwrite(out_img_path,pred)

至此,torch转ONNX的全流程就走完了,欢迎交流debug!

整理不易,欢迎一键三连!!!

相关内容

热门资讯

【NI Multisim 14...   目录 序言 一、工具栏 🍊1.“标准”工具栏 🍊 2.视图工具...
银河麒麟V10SP1高级服务器... 银河麒麟高级服务器操作系统简介: 银河麒麟高级服务器操作系统V10是针对企业级关键业务...
不能访问光猫的的管理页面 光猫是现代家庭宽带网络的重要组成部分,它可以提供高速稳定的网络连接。但是,有时候我们会遇到不能访问光...
AWSECS:访问外部网络时出... 如果您在AWS ECS中部署了应用程序,并且该应用程序需要访问外部网络,但是无法正常访问,可能是因为...
Android|无法访问或保存... 这个问题可能是由于权限设置不正确导致的。您需要在应用程序清单文件中添加以下代码来请求适当的权限:此外...
北信源内网安全管理卸载 北信源内网安全管理是一款网络安全管理软件,主要用于保护内网安全。在日常使用过程中,卸载该软件是一种常...
AWSElasticBeans... 在Dockerfile中手动配置nginx反向代理。例如,在Dockerfile中添加以下代码:FR...
AsusVivobook无法开... 首先,我们可以尝试重置BIOS(Basic Input/Output System)来解决这个问题。...
ASM贪吃蛇游戏-解决错误的问... 要解决ASM贪吃蛇游戏中的错误问题,你可以按照以下步骤进行:首先,确定错误的具体表现和问题所在。在贪...
月入8000+的steam搬砖... 大家好,我是阿阳 今天要给大家介绍的是 steam 游戏搬砖项目,目前...