mmdeploy 转换torch为onnx
安装 mmdeploy
以前已经安装了相关的mmsegment、mmcv
和 mmdetection
等库
方式安装 mmdeploy
:
pip install mmdeploy
或者,也可以从源代码安装,注意分支版本:
git clone https://github.com/open-mmlab/mmdeploy.git
cd mmdeploy
pip install -e .
应用
torch转onnx 使用MMDeploy 中的工具 tools/deploy.py
python tools/deploy.py \$DEPLOY_CFG \$MODEL_CFG \$PTH_MODEL_PATH \--work-dir $OUT_PATH \--show --device cuda:0 --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 configs/mmseg/segmentation_onnxruntime_static-640x640.py ~/prjs/mmsegmentation/configs/hrnet/fcn-HR48.py ~/prjs/mmsegmentation/work_dirs/fcn-HR48/iter_32000.pth ~/prjs/mmsegmentation/demo/demo.jpg --work-dir work_dir_hr --show --device cuda:0