How to convert YOLOv5 model to ncnn model

ncnn
2022.11.12

この記事では、YOLOv5のPyTorchモデルをncnnモデルに変換する方法について解説します。一度PyTorchモデル(.pt)をONNXモデルに変換し、ONNXモデルからncnnモデル(.paramと.bin)に変換します。またYOLOv5での注意点を説明します。

環境

セットアップ

こちらの記事では、ncnnをJetson Nanoにインストールしましたが、今回はホストPCで動作確認をするために、ホストPCにncnnをインストールします。

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 install
Code language: Bash (bash)

なお詳細なインストール方法はこちらとこちらの記事を参照して下さい。

動作確認

以下を実行します。レポジトリshaoshengsong/yolov5_62_export_ncnnはncnnのyolov5サンプルコードが参照する yolov5s_6.2.paramとyolov5s_6.2.binを持っています。

git clone https://github.com/shaoshengsong/yolov5_62_export_ncnn
cd yolov5_62_export_ncnn
../examples/yolov5 bus.jpg
Code language: Bash (bash)

以下の画像が表示されます。

bus.jpg with ncnn yolov5

PyTorchモデルからONNXモデルへの変換

PyTorchモデル(.pt)からONNXモデルへの変換は、先程クローンしたshaoshengsong/yolov5_62_export_ncnnレポジトリを使用します。YOLOv5オリジナルのレポジトリにもONNXに変換するスクリプトがありますが、出力結果が異なります。

$ pwd
yolov5_62_export_ncnn
Code language: Bash (bash)

必要なパッケージをインストールします。

sudo apt install -y python3-pip
python3 -m pip install -U pip
python3 -m pip install -r requirements.txt
Code language: Bash (bash)

YOLOv5のPyTorchモデルを取得し、ONNXモデルに変換します。

$ 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
Code language: Bash (bash)

ONNXモデルからncnnモデルへの変換

ONNXモデルからncnnモデルへの変換は、onnx2ncnnを使用します。onnx2ncnnはncnnをビルド＆インストールしたときにbuild/install/bin/onnx2ncnnに作成されます。

../install/bin/onnx2ncnn yolov5s.onnx yolov5s.param yolov5s.bin
Code language: Bash (bash)

さて、ncnnのYOLOv5のサンプルコードを見てみるとyolov5s_6.2.paramとyolov5s_6.2.binを直接指定してモデルを読み込んでいることが分かります。今回作成したファイルはyolov5s.paramとyolov5s.binですのでファイル名を以下のように変更します。

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)

ビルドして再度実行します。

$ 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)

先程と異なり、エラーが表示されbus.jpgが表示されませんでした。

ここでは簡単に修正方法について解説します。

まず上のエラーについてyolov5s.paramを調べてみます。

$ 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
Code language: Bash (bash)

output、354、366という数字が含まれていました。yolov5s.onnxをhttps://netron.appで可視化してみます。

yolov5s.onnxはoutput、354、366という３つの出力があるようです。

一方、shaoshengsong/yolov5_62_export_ncnnにあるyolov5s_6.2.paramを調べてみます。

$ 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
Code language: Bash (bash)

表示は少し異なりますが、output、353、367が出力となっているようです。これはyolov5s.paramを調べた値(354、366)と異なります。

ncnnのYOLOv5のサンプルコードを見てみると以下のコードが存在します。

$ grep -w -e 353 -e 367 ../../examples/yolov5.cpp 
        ex.extract("353", out);
        ex.extract("367", out);
Code language: Bash (bash)

ここでoutputの出力を指定しているようです。yolov5.cppを以下のように修正します。

@@ -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);
 #else
Code language: Diff (diff)

ビルドして再度実行します。

$ 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
...
Code language: Bash (bash)

エラーはでなくなりましたが、以下の画像が表示されます。

詳細はこちらの記事にありますが、Reshapeのパラメータを修正します。なおパラメータのフォーマットはこちらを参照してください。

$ 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
Code language: Bash (bash)

yolov5s.paramは0=6400、0=1600、0=400となっています。

一方yolov5s_6.2.paramは0=-1、0=-1、0=-1となっています。

yolov5s.paramを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=1
Code language: Diff (diff)

yolov5s.paramの修正ですのでビルドは不要です。修正したら再度実行します。

../examples/yolov5 bus.jpg
Code language: Bash (bash)

うまく表示されました！

自分のデータセットで学習した結果を使ってJetson Nano上でncnnモデルを使用しYOLOv5を実行した動画が以下になります。

ncnn YOLOv5 with own dataset on Jetson Nano

以上です。

参考

GitHub - Tencent/ncnn: ncnn is a high-performance neural network inference framework optimized for the mobile platform
ncnn is a high-performance neural network inference framework optimized for the mobile platform - GitHub - Tencent/ncnn: ncnn is a high-performance neural netwo...
github.com
how to build
ncnn is a high-performance neural network inference framework optimized for the mobile platform - Tencent/ncnn
github.com
Tutorial for compiling NCNN library - Zekun Wang's Blog | Data Science
Across the Great Wall, we can reach every corner in the world
waittim.github.io
详细记录u版YOLOv5目标检测ncnn实现
详细记录u版YOLOv5目标检测ncnn实现 允许在不修改内容前提下转载本文 0x0 u版YOLOv5众所周知，原版YOLO系列是 darknet 框架训练的，而广泛使用的是 YOLOv4 作者 AlexeyAB 的版本 AlexeyAB 首字母是a，于是也被叫做…
zhuanlan.zhihu.com
operation param weight table
ncnn is a high-performance neural network inference framework optimized for the mobile platform - Tencent/ncnn
github.com
Netron
Visualizer for neural network, deep learning and machine learning models.
netron.app