一步一步教你如何将 yolov3/yolov4 转为 caffe 模型

2022-01-20 20:50:23 浏览数 (1)

实际工作中,目标检测 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,看到这个镜像:

点击进去,看到该镜像的下拉命令:

代码语言:javascript复制
sudo docker pull bvlc/caffe

打开 Ubuntu 终端,输入上面的下拉命令,将 caffe 的 docker 镜像下载下来。

下载完成之后,终端输入命令:

代码语言:javascript复制
sudo docker images

就可以看到已下载的 caffe 镜像了。

为镜像创建容器:

代码语言:javascript复制
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 上的转换工具项目:

代码语言:javascript复制
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。按照如下说明修改相应字段的程序。

代码语言:javascript复制
// 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 目录,输入以下命令:

代码语言:javascript复制
make clean
make all -j8
make pycaffe -j8

caffe 重新编译之后,就可以对 yolov3/yolov4 模型进行 caffe 转换了。

4. 模型转换

准备好我们已有的 yolov3 模型的配置文件和权重文件,例如:yolov3.cfg 和 yolov3.weights。在 darknet2caffe 目录下,输入以下命令:

代码语言:javascript复制
python darknet2caffe.py ./yolov3.cfg ./yolov3.weights ./yolov3.prototxt ./yolov3.caffemodel

如果输出类似下面的语句,则证明转换成功!

代码语言:javascript复制
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 模型完成!

0 人点赞