yolo系列的ONNX部署(C++)【适用于v4\v5\v5-6.1\v7】

前言

  首先感谢bubbliiiing博主提供的练习代码！ 整个布置程序分为四步（C++）：

1. 下载bubbliiiing博主项目更改predict 为 export_onnx\
2. 更改文件夹nets下的yolo.py代码\
3. 生成onnx文件
4. 下载此项目进行cmake&make 经测验，可以对(pytorch版本的)yolov4 — yolov5 — yolov5-6.1 yolov7 进行转换测验 因为bubbliiiing博主已经将网络结构和怎么练习，我这边就不再过多的讲述怎么练习。

怎么运用？

  这儿我以yolov7为例子向咱们展现怎么得到onnx以及运转onnx进行布置(bubbliiiing博主v4>>v5>>v5-6.1同理)

Step1:

git clone github.com/bubbliiiing… cd yolov7-pytorch
pip install -r requirements.txt
vim predict.py

line24: mode = "export_onnx" ## predict >> 'export_onnx'

Step2:

cd /nets
vim yolo.py

line300 ----line316:
        #---------------------------------------------------#
        #   第三个特征层
        #   y3=(batch_size, 75, 80, 80)
        #---------------------------------------------------#
        out2 = self.yolo_head_P3(P3)
        #---------------------------------------------------#
        #   第二个特征层
        #   y2=(batch_size, 75, 40, 40)
        #---------------------------------------------------#
        out1 = self.yolo_head_P4(P4)
        #---------------------------------------------------#
        #   第一个特征层
        #   y1=(batch_size, 75, 20, 20)
        #---------------------------------------------------#
        out0 = self.yolo_head_P5(P5)
        return [out0, out1, out2]

如下图为bubbliiiing博主源码生成的onnx输入输出信息截取

  bubbliiiing 这儿的return[out0, out1, out2] 在经过”export_onnx”后并没有torch.cat，咱们需求将输出的成果进行torch.cat操作， 在进行torch.cat操作之前需求将输出成果进行转换一下维度，详见如下代码：

        out2 = self.yolo_head_P3(P3)
        bs, _, ny, nx = out2.shape  
        # 这儿的 no = 85[类别数nc（80） + 目标坐标和置信度（4+1）]
        out2 = out2.view(bs, 3, 85, ny, nx).permute(0, 1, 3, 4, 2).contiguous() # no = 85
        out2 = out2.view(bs * 3 * ny * nx, 85, 1).contiguous() #  no = 85
        # ---------------------------------------------------#
        #   第二个特征层
        # ---------------------------------------------------#
        out1 = self.yolo_head_P4(P4)
        bs, _, ny, nx = out1.shape
        out1 = out1.view(bs, 3, 85, ny, nx).permute(0, 1, 3, 4, 2).contiguous() #  no = 85
        out1 = out1.view(bs * 3 * ny * nx, 85, 1).contiguous() #  no = 85
        # ---------------------------------------------------#
        #   第一个特征层
        # ---------------------------------------------------#
        out0 = self.yolo_head_P5(P5)
        bs, _, ny, nx = out0.shape
        out0 = out0.view(bs, 3, 85, ny, nx).permute(0, 1, 3, 4, 2).contiguous() #  no = 85
        out0 = out0.view(bs * 3 * ny * nx, 85, 1).contiguous() #  no = 85
        output = torch.cat((out2, out1, out0))
        output = output.permute(2, 0, 1)
        # return [out0, out1, out2]
        return output

  运用上述的代码部分替代源代码的line300 至 line316: 替换后得到的信息下图所示

  咱们在替换自己的数据集进行练习得到的权重进行pth2onnx时需求注意：上述代码注释部分的替换85（coco数据集的类别数+5）， 在这儿需求替换为自己的数据集的类别数再加上5。例如自己的数据集是10个类别(nc=10)，那么这儿的no=10+5 ;也便是85需求更改为15（no = nc + 5）

Step3:

cd ..
python predict.py
这样咱们就可以得到C++调用onnx需求的权重

step4:

git clone github.com/kivenyangmi…
cd yolo-bubbliiiing-onnx
mkdir build
cd build
cmake ..
cmake –build . (或许:make)