Skip to content

Latest commit

 

History

History
51 lines (37 loc) · 3.23 KB

README.md

File metadata and controls

51 lines (37 loc) · 3.23 KB
Raw

How to convert the original PyTorch MODNet model to ONNX

The original pre-trained PyTorch MODNet model comes from ZHKKKe/MODNet. Note that this pre-trained model is under Creative Commons Attribution NonCommercial ShareAlike 4.0 license.

You could use the script in this repository to convert the original PyTorch model to ONNX. I recommend to do such conversion within a python3 virtual environment, since you'd need to use some specific versions of pip3 packages. Below is a step-by-step guide about how to build the python3 virtual environment and then convert the PyTorch MODNet model to ONNX.

  1. Make sure python3 "venv" module is installed.

    $ sudo apt install python3-venv

  2. Create a virtual environment named "venv-onnx" and activate it.

    $ cd ${HOME}/project/tensorrt_demos/modnet
$ python3 -m venv venv-onnx
$ source venv-onnx/bin/activate

    At this point, you should have entered the virtual environment and would see shell prompt proceeded with "(venv-onnx) ". You could do deactivate to quit the virtual environment when you are done using it.

    Download "torch-1.7.0-cp36-cp36m-linux_aarch64.whl" from here: PyTorch for Jetson. Then install all required packages into the virtual environment. (Note the following should be done inside the "venv-onnx" virtual environment.)

    ### update pip to the latest version in the virtual env
$ curl https://bootstrap.pypa.io/get-pip.py | python
### udpate these essential packages
$ python -m pip install -U setuptools Cython
### I recommend numpy 1.16.x on Jetson
$ python -m pip install "numpy<1.17.0"
### install cv2 into the virtual env
$ cp -r /usr/lib/python3.6/dist-packages/cv2 venv-onnx/lib/python3.6/site-packages/
### install PyImage, onnx and onnxruntime
$ python -m pip install PyImage onnx==1.8.1 onnxruntime==1.6.0
### install PyTorch v1.7.0
$ sudo apt install libopenblas-base libopenmpi-dev
$ python -m pip install ${HOME}/Downloads/torch-1.7.0-cp36-cp36m-linux_aarch64.whl

    In addition, you might also install onnx-graphsurgeon and polygraphy for debugging. Otherwise, you could do some simple testing to make sure "onnx" and "torch" are working OK in the virtual env.

  3. Download the pre-trained MODNet model (PyTorch checkpoint file) from the link on this page: /ZHKKKe/MODNet/pretrained. I recommend using "modnet_webcam_portrait_matting.ckpt". Just put the file in the current directory.

  4. Do the conversion using the following command. The ouput "modnet.onnx" would be generated.

    $ python -m torch2onnx.export modnet_webcam_portrait_matting.ckpt modnet.onnx

    By default, the "torch2onnx.export" script sets input image width and height to 512x288. They could be modified by the "--width" and "--height" command-line options. In addition, the "-v" command-line option could be used to enable verbose logs of torch.onnx.export().