この記事では、YOLOv5のPyTorchモデルをncnnモデルに変換する方法について解説します。一度PyTorchモデル(.pt)をONNXモデルに変換し、ONNXモデルからncnnモデル(.paramと.bin)に変換します。またYOLOv5での注意点を説明します。
環境
- ホストPC: Ubuntu20.04
- YOLOv5: https://github.com/ultralytics/yolov5/tree/v6.2
- version v6.2
- ncnn: https://github.com/Tencent/ncnn
セットアップ
こちらの記事では、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 installCode 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.jpgCode language: Bash (bash)以下の画像が表示されます。
PyTorchモデルからONNXモデルへの変換
PyTorchモデル(.pt)からONNXモデルへの変換は、先程クローンしたshaoshengsong/yolov5_62_export_ncnnレポジトリを使用します。YOLOv5オリジナルのレポジトリにもONNXに変換するスクリプトがありますが、出力結果が異なります。
$ pwd
yolov5_62_export_ncnnCode language: Bash (bash)必要なパッケージをインストールします。
sudo apt install -y python3-pip
python3 -m pip install -U pip
python3 -m pip install -r requirements.txtCode 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.onnxCode language: Bash (bash)ONNXモデルからncnnモデルへの変換
ONNXモデルからncnnモデルへの変換は、onnx2ncnnを使用します。onnx2ncnnはncnnをビルド＆インストールしたときにbuild/install/bin/onnx2ncnnに作成されます。
../install/bin/onnx2ncnn yolov5s.onnx yolov5s.param yolov5s.binCode 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=1Code 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=1Code 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);
#elseCode 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=3Code 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=1Code language: Diff (diff)yolov5s.paramの修正ですのでビルドは不要です。修正したら再度実行します。
../examples/yolov5 bus.jpgCode language: Bash (bash)うまく表示されました！
自分のデータセットで学習した結果を使ってJetson Nano上でncnnモデルを使用しYOLOv5を実行した動画が以下になります。
以上です。
参考
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 - Tencent/ncnn
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
https://zhuanlan.zhihu.com/p/275989233
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
