This article explains how to convert a YOLOv5 PyTorch model to an ncnn model. Once the PyTorch model (.pt) is converted to an ONNX model, and then from the ONNX model to the ncnn model (.param and .bin). In addition, some notes on YOLOv5 are explained.
Environment
- HostPC: Ubuntu20.04
- YOLOv5: https://github.com/ultralytics/yolov5/tree/v6.2
- version v6.2
- ncnn: https://github.com/Tencent/ncnn
Setup
In this article, ncnn was installed on Jetson Nano, but this time ncnn is installed on the host PC to check the operation on the host PC.
sudo apt update
sudo apt install -y cmake build-essential libprotobuf-dev protobuf-compiler libopencv-dev
git clone https://github.com/Tencent/ncnn
cd ncnn
mkdir build
cd build
cmake -DCMAKE_BUILD_TYPE=Release -DNCNN_VULKAN=OFF -DNCNN_SYSTEM_GLSLANG=ON -DNCNN_BUILD_EXAMPLES=ON ..
make -j
make installCode language: Bash (bash)For detailed installation instructions, please refer to this article and this article.
operation check
Execute the following. Repository shaoshengsong/yolov5_62_export_ncnn has yolov5s_6.2.param and yolov5s_6.2.bin which are referenced by the yolov5 sample code in ncnn.
git clone https://github.com/shaoshengsong/yolov5_62_export_ncnn
cd yolov5_62_export_ncnn
../examples/yolov5 bus.jpgCode language: Bash (bash)The following images are shown.
Conversion from PyTorch model to ONNX model
To convert a PyTorch model (.pt) to an ONNX model, use the shaoshengsong/yolov5_62_export_ncnn repository that was cloned above. the original YOLOv5 repository also has a script to convert to ONNX, but the output results differently.
$ pwd
yolov5_62_export_ncnnCode language: Bash (bash)Install the necessary packages.
Code language: Bash (bash)sudo apt install -y python3-pip python3 -m pip install -U pip python3 -m pip install -r requirements.txt
Retrieve a PyTorch model of YOLOv5 and convert it to an ONNX model.
Code language: Bash (bash)$ get https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s.pt $ python3 export.py --weights yolov5s.pt --img 640 --batch 1 --include onnx --simplify $ ls yolov5s.onnx yolov5s.onnx
Conversion from ONNX model to ncnn model
The conversion from ONNX models to ncnn models is done using onnx2ncnn, which is created in build/install/bin/onnx2ncnn when ncnn is built & installed.
Code language: Bash (bash)../install/bin/onnx2ncnn yolov5s.onnx yolov5s.param yolov5s.bin
Now, if we look at ncnn’s YOLOv5 sample code, we can see that the model is loaded by directly specifying yolov5s_6.2.param and yolov5s_6.2.bin. The files created this time are yolov5s.param and yolov5s.bin, so the file names should be changed as follows
diff --git a/examples/yolov5.cpp b/examples/yolov5.cpp
index 88f6db2..ef5b2d5 100644
--- a/examples/yolov5.cpp
+++ b/examples/yolov5.cpp
@@ -280,9 +280,9 @@ static int detect_yolov5(const cv::Mat& bgr, std::vector<Object>& objects)
// original pretrained model from https://github.com/ultralytics/yolov5
// the ncnn model https://github.com/nihui/ncnn-assets/tree/master/models
#if YOLOV5_V62
- if (yolov5.load_param("yolov5s_6.2.param"))
+ if (yolov5.load_param("yolov5s.param"))
exit(-1);
- if (yolov5.load_model("yolov5s_6.2.bin"))
+ if (yolov5.load_model("yolov5s.bin"))
exit(-1);
#elif YOLOV5_V60
if (yolov5.load_param("yolov5s_6.0.param"))Code language: Diff (diff)Build and run again.
$ make -C ..
$ ../examples/yolov5 bus.jpg
find_blob_index_by_name 353 failed
Try
ex.extract("output", out0);
ex.extract("354", out1);
ex.extract("366", out2);
Floating point exception (core dumped)Code language: Bash (bash)Unlike before, an error was displayed and bus.jpg did not show up.
Here is a brief explanation of how to fix it.
First, let’s examine yolov5s.param for the above error.
Code language: Bash (bash)$ grep -w -e output -e 354 -e 366 yolov5s.param Permute /model.24/Transpose 1 1 /model.24/Reshape_output_0 output 0=1 Permute /model.24/Transpose_1 1 1 /model.24/Reshape_1_output_0 354 0=1 Permute /model.24/Transpose_2 1 1 /model.24/Reshape_2_output_0 366 0=1
It contained the numbers output, 354 and 366. let’s visualize yolov5s.onnx at https://netron.app.
The yolov5s.onnx seems to have three outputs: output, 354, and 366.
On the other hand, examine yolov5s_6.2.param in shaoshengsong/yolov5_62_export_ncnn.
Code language: Bash (bash)$ grep Permute yolov5s_6.2.param Permute Transpose_200 1 1 338 output 0=1 Permute Transpose_203 1 1 352 353 0=1 Permute Transpose_206 1 1 366 367 0=1
The display is slightly different, but output, 353 and 367 seem to be the outputs. This differs from the values (354, 366) examined yolov5s.param.
If we look at the sample code for YOLOv5 in ncnn, the following code exists.
$ grep -w -e 353 -e 367 ../../examples/yolov5.cpp
ex.extract("353", out);
ex.extract("367", out);Code language: Bash (bash)It seems that output is specified here. yolov5.cpp should be modified as follows.
@@ -366,7 +366,7 @@ static int detect_yolov5(const cv::Mat& bgr, std::vector<Object>& objects)
ncnn::Mat out;
#if YOLOV5_V62
- ex.extract("353", out);
+ ex.extract("354", out);
#elif YOLOV5_V60
ex.extract("376", out);
#else
@@ -391,7 +391,7 @@ static int detect_yolov5(const cv::Mat& bgr, std::vector<Object>& objects)
{
ncnn::Mat out;
#if YOLOV5_V62
- ex.extract("367", out);
+ ex.extract("366", out);
#elif YOLOV5_V60
ex.extract("401", out);
#elseCode language: Diff (diff)Build and run again.
Code language: Bash (bash)$ make -C .. $ ../examples/yolov5 bus.jpg 67 = 0.99632 at 656.76 509.26 152.24 x 0.00 67 = 0.99476 at 644.50 499.47 164.50 x 0.00 67 = 0.99158 at 809.00 464.90 0.00 x 0.00 67 = 0.98984 at 809.00 452.97 0.00 x 0.00 67 = 0.98339 at 195.63 596.97 222.76 x 0.00 67 = 0.98278 at 749.81 304.04 59.19 x 0.00 67 = 0.98148 at 809.00 294.05 0.00 x 0.00 67 = 0.97992 at 0.00 523.36 195.36 x 0.00 67 = 0.97671 at 726.73 307.79 82.27 x 0.00 67 = 0.97668 at 45.13 341.54 48.62 x 0.00 ...
The error no longer appears, but the following image is shown.
The details can be found in this article, but we will modify the parameters of Reshape. The format of the parameters can be found here.
Code language: Bash (bash)$ grep ^Reshape yolov5s.param Reshape /model.24/Reshape 1 1 /model.24/m.0/Conv_output_0 /model.24/Reshape_output_0 0=6400 1=85 2=3 Reshape /model.24/Reshape_1 1 1 /model.24/m.1/Conv_output_0 /model.24/Reshape_1_output_0 0=1600 1=85 2=3 Reshape /model.24/Reshape_2 1 1 /model.24/m.2/Conv_output_0 /model.24/Reshape_2_output_0 0=400 1=85 2=3 $ $ grep ^Reshape yolov5s_6.2.param Reshape Reshape_199 1 1 326 338 0=-1 1=85 2=3 Reshape Reshape_202 1 1 340 352 0=-1 1=85 2=3 Reshape Reshape_205 1 1 354 366 0=-1 1=85 2=3
The yolov5s.param has 0=6400, 0=1600, and 0=400.
On the other hand, yolov5s_6.2.param has 0=-1, 0=-1, and 0=-1.
Modify yolov5s.param to 0=-1.
@@ -165,11 +165,11 @@
Convolution /model.23/cv3/conv/Conv 1 1 /model.23/Concat_output_0 /model.23/cv3/conv/Conv_output_0 0=512 1=1 11=1 2=1 12=1 3=1 13=1 4=0 14=0 15=0 16=0 5=1 6=262144
Swish /model.23/cv3/act/Mul 1 1 /model.23/cv3/conv/Conv_output_0 /model.23/cv3/act/Mul_output_0
Convolution /model.24/m.0/Conv 1 1 /model.17/cv3/act/Mul_output_0_splitncnn_0 /model.24/m.0/Conv_output_0 0=255 1=1 11=1 2=1 12=1 3=1 13=1 4=0 14=0 15=0 16=0 5=1 6=32640
-Reshape /model.24/Reshape 1 1 /model.24/m.0/Conv_output_0 /model.24/Reshape_output_0 0=6400 1=85 2=3
+Reshape /model.24/Reshape 1 1 /model.24/m.0/Conv_output_0 /model.24/Reshape_output_0 0=-1 1=85 2=3
Permute /model.24/Transpose 1 1 /model.24/Reshape_output_0 output 0=1
Convolution /model.24/m.1/Conv 1 1 /model.20/cv3/act/Mul_output_0_splitncnn_0 /model.24/m.1/Conv_output_0 0=255 1=1 11=1 2=1 12=1 3=1 13=1 4=0 14=0 15=0 16=0 5=1 6=65280
-Reshape /model.24/Reshape_1 1 1 /model.24/m.1/Conv_output_0 /model.24/Reshape_1_output_0 0=1600 1=85 2=3
+Reshape /model.24/Reshape_1 1 1 /model.24/m.1/Conv_output_0 /model.24/Reshape_1_output_0 0=-1 1=85 2=3
Permute /model.24/Transpose_1 1 1 /model.24/Reshape_1_output_0 354 0=1
Convolution /model.24/m.2/Conv 1 1 /model.23/cv3/act/Mul_output_0 /model.24/m.2/Conv_output_0 0=255 1=1 11=1 2=1 12=1 3=1 13=1 4=0 14=0 15=0 16=0 5=1 6=130560
-Reshape /model.24/Reshape_2 1 1 /model.24/m.2/Conv_output_0 /model.24/Reshape_2_output_0 0=400 1=85 2=3
+Reshape /model.24/Reshape_2 1 1 /model.24/m.2/Conv_output_0 /model.24/Reshape_2_output_0 0=-1 1=85 2=3
Permute /model.24/Transpose_2 1 1 /model.24/Reshape_2_output_0 366 0=1Code language: Diff (diff)This is a modification of yolov5s.param, so no build is required. Once corrected, run it again.
Code language: Bash (bash)../examples/yolov5 bus.jpg
It showed up well!
Below is a video of YOLOv5 running on Jetson Nano using the ncnn model with the results trained with own dataset.
That’s all.
Reference
