YOLOv5 CPU Export and OpenVINO Inference

Documentation on exporting YOLOv5 models for fast CPU inference using Intel's OpenVINO framework (Tested on commits up to June 6, 2022 in docker).

Google Colab Conversion

Convert yolov5 model to IR format with Google Colab. (Recommended)

1. Clone and set up the Official YOLOv5 GitHub repository

Setup

All package installations should be done in a virtualenv or conda env to prevent package conflict errors.

Install required requirements for onnx and openvino Inference

pip install --upgrade pip
pip install -r inf_requirements.txt

Clone and install requirements for yolov5 repository

git clone https://github.com/ultralytics/yolov5                    # clone repo
cd yolov5
pip install -r requirements.txt                                    # base requirements

2. Export a Trained YOLOv5 Model as ONNX

Export

Export a pre-trained or custom trained YOLOv5 model to generate the respective ONNX, TorchScript and CoreML formats of the model. The pre-trained yolov5s.pt is the lightest and fastest model for CPU inference. Other slower but more accurate models include yolov5m.pt, yolov5l.pt and yolov5x.pt. All available model details at Ultralytics YOLOv5 README.

A custom training checkpoint i.e. runs/exp/weights/best.pt can be used for conversion as well.

Export a pre-trained light yolov5s.pt model at 640x640 with batch size 1

python export.py --weights yolov5s.pt --include onnx --img 640 --batch 1

Export a custom checkpoint for dynamic input shape {BATCH_SIZE, 3, HEIGHT, WIDTH}. Note, for CPU inference mode, BATCH_SIZE must be set to 1. Install onnx-simplifier for simplifying onnx exports

pip install onnx-simplifier==0.3.10                                
python export.py --weights runs/exp/weights/best.pt --include onnx  --dynamic --simplify

Cd to yolov5_export_cpu dir and move the onnx model to yolov5_export_cpu/models directory

mv <PATH_TO_ONNX_MODEL> yolov5_export_cpu/models/

3. Test YOLOv5 ONNX model inference

ONNX inference

python detect_onnx.py -m image -i <IMG_FILE_PATH/IMG_DIR_PATH>
python detect_onnx.py -m video -i <VID_PATH_FILE>
# python detect_onnx.py -h for more info

Optional: To convert the all frames in the output directory into a mp4 video using ffmpeg, use ffmpeg -r 25 -start_number 00001 -i output/frame_onnx_%5d.jpg -vcodec libx264 -y -an onnx_result.mp4

4. Export ONNX to OpenVINO

Recommended Option A

Option A. Use OpenVINO's python dev library

A1. Install OpenVINO python dev library

Instructions for setting OpenVINO available here

# install required OpenVINO lib to convert ONNX to OpenVINO IR
pip install openvino-dev[onnx]

A2. Export ONNX to OpenVINO IR

This will create the OpenVINO Intermediate Model Representation (IR) model files (xml and bin) in the directory models/yolov5_openvino.

Important Note: --input_shape must be provided and match the img shape used to export ONNX model. Batching might not supported for CPU inference

# export onnx to OpenVINO IR
mo \
  --progress \
  --input_shape [1,3,640,640] \
  --input_model models/yolov5s.onnx \
  --output_dir models/yolov5_openvino

Full OpenVINO export options

Option B. Use OpenVINO Docker

B1. Download Docker and OpenVINO Docker Image

Install docker in your system if not already installed.

Pass the docker run command below in a terminal which will automatically download the OpenVINO Docker Image and run it. The models directory containing the ONNX model must be in the current working directory.

docker run -it --rm \
            -v $PWD/models:/home/openvino/models \
            openvino/ubuntu18_dev:latest \
            /bin/bash -c "cd /home/openvino/; bash"

B2. Export ONNX model to an OpenVINO IR representation

This will create the OpenVINO Intermediate Model Representation (IR) model files (xml and bin) in the directory models/yolov5_openvino which will be available in the host system outside the docker container.