手把手教你如何將 yolov3/yolov4 轉為 caffe 模型

紅色石頭發表於2021-03-15
YOLO模型

實際工作中，目標檢測 yolov3 或者 yolov4 模型移植到 AI 晶片中，經常需要將其先轉換為 caffe1.x 模型，大家可能或多或少也有這方面的需求。例如華為海思 NNIE 只支援caffe1.x 模型，所以 yolov3/yolov4 模型要想在海思晶片上部署，轉換為 caffe1.x 模型是必須的。

今天，專門給大家寫一下將 yolov3/yolov4 模型轉為 caffe 模型的詳細步驟。

0. 系統環境

  • Ubuntu 16.04
  • cuda 9.x

1. 部署 caffe 環境

常規的方法是下載 caffe 的原始碼，進行編譯安裝。但是，這種方法比較複雜，不太容易成功。更簡單的方法是直接下載 caffe1.x 的docker 映象。這種方法需要你提前在 Ubuntu 裡安裝了 nvidia-docker。

開啟 docker hub 網站：

https://registry.hub.docker.com/

搜尋 caffe，看到這個映象：

點選進去，看到該映象的下拉命令：

sudo docker pull bvlc/caffe

就可以看到已下載的 caffe 映象了。

為映象建立容器：

docker run --runtime=nvidia --name caffe -i -t bvlc/caffe /bin/bash

這裡我們給該容器起的名字是 caffe，大家可以自由設定，根據 sudo docker ps -a 這條命令來檢視各個容器。

這樣，我們就進入了 caffe1.x 的 docker 容器之內了。

該容器已部署好了 caffe1.x，caffe1.x 路徑為：

2. caffe 原始碼修改

因為官方 caffe1.x 框架不支援 yolo3/yolov4 的 upsample 層，所以需要手動增加 upsample 層，對 caffe 原始碼進行修改。

克隆 GitHub 上的轉換工具專案：

git clone https://github.com/ChenYingpeng/darknet2caffe.git

  • 將 darknet2caffe/caffe_layers/mish_layer 下的 mish_layer.hpp 檔案和 darknet2caffe/tree/master/caffe_layers/upsample_layer 下的 upsample_layer.hpp 拷貝到容器的路徑：/opt/caffe/include/caffe/layers 下。

  • 將 darknet2caffe/caffe_layers/mish_layer 下的 mish_layer.cpp、mish_layer.cu 檔案和 darknet2caffe/tree/master/caffe_layers/upsample_layer 下的 upsample_layer.cpp、upsample_layer.cu 拷貝到容器的路徑：/opt/caffe/src/caffe/layers/ 下。

  • 將 darknet2caffe/caffe_layers/pooling_layer 下的 pooling_layer.cpp 拷貝到容器的路徑：/opt/caffe/src/caffe/layers/ 下。

然後，開啟容器內的 caffe 檔案：/opt/caffe/src/caffe/proto/caffe.proto。按照如下說明修改相應欄位的程式。

// LayerParameter next available layer-specific ID: 147 (last added: recurrent_param)
message LayerParameter {
  optional TileParameter tile_param = 138;
  optional VideoDataParameter video_data_param = 207;
  optional WindowDataParameter window_data_param = 129;
++optional UpsampleParameter upsample_param = 149; //added by chen for Yolov3, make sure this id 149 not the same as before.
++optional MishParameter mish_param = 150; //added by chen for yolov4,make sure this id 150 not the same as before.
}

// added by chen for YoloV3
++message UpsampleParameter{
++  optional int32 scale = 1 [default = 1];
++}

// Message that stores parameters used by MishLayer
++message MishParameter {
++  enum Engine {
++    DEFAULT = 0;
++    CAFFE = 1;
++    CUDNN = 2;
++  }
++  optional Engine engine = 2 [default = DEFAULT];
++}

其中，++ 表示該行是增加的內容。

3. caffe 重新編譯

修改完 caffe 的一些原始碼之後，需要對 caffe 重新編譯。

進入 /opt/caffe/build 目錄，輸入以下命令：

make clean
make all -j8
make pycaffe -j8

caffe 重新編譯之後，就可以對 yolov3/yolov4 模型進行 caffe 轉換了。

4. 模型轉換

準備好我們已有的 yolov3 模型的配置檔案和權重檔案，例如：yolov3.cfg 和 yolov3.weights。在 darknet2caffe 目錄下，輸入以下命令：

python darknet2caffe.py ./yolov3.cfg ./yolov3.weights ./yolov3.prototxt ./yolov3.caffemodel

如果輸出類似下面的語句，則證明轉換成功！

I0522 10:19:19.015708 25251 net.cpp:228] layer1-act does not need backward computation.
I0522 10:19:19.015712 25251 net.cpp:228] layer1-scale does not need backward computation.
I0522 10:19:19.015714 25251 net.cpp:228] layer1-bn does not need backward computation.
I0522 10:19:19.015718 25251 net.cpp:228] layer1-conv does not need backward computation.
I0522 10:19:19.015722 25251 net.cpp:228] input does not need backward computation.
I0522 10:19:19.015725 25251 net.cpp:270] This network produces output layer139-conv
I0522 10:19:19.015731 25251 net.cpp:270] This network produces output layer150-conv
I0522 10:19:19.015736 25251 net.cpp:270] This network produces output layer161-conv
I0522 10:19:19.015911 25251 net.cpp:283] Network initialization done.
unknow layer type yolo 
unknow layer type yolo 
save prototxt to ./yolov3.prototxt
save caffemodel to ./yolov3.caffemodel

其中，yolov3.prototxt 和 yolov3.caffemodel 為轉換後的 caffe 模型。

至此，yolov3/yolov4 轉換為 caffe 模型完成！

相關文章