Opencv、dnn部署自己的Yolov5模型记录

Opencv、dnn部署自己的Yolov5模型记录

一、环境配置

1.opencv == 4.5.1+dnn模块
2.pytorch == 1.8
3.ubuntu18.04

二、代码来源

1.https://github.com/hpc203/yolov5-dnn-cpp-python-v2
2.https://github.com/ultralytics/yolov5
注：选用第四版。一定是第4版，其他的版本对应网络模型不一样。

三、训练自己的模型

先插一句话：一定要配置pytorch>=1.7.1,python版本可以不用担心（我的python3.6没问题）
1.本次实验训练了自己的二维码模型，具体步骤可参考我的博客训练之前先参考下一步修改common.py
2.修改common.py文件中的两个类

class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
        self.contract = Contract(gain=2)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        #return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
        return self.conv(self.contract(x))


class Contract(nn.Module):
    # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
    def __init__(self, gain=2):
        super().__init__()
        self.gain = gain

注：这一步是为了修改切片。
3.开始训练，获得自己的模型。将自己的模型（best.pt或者last.pt）从命名为yolov5s.pt。并放置于yolov5根目录下。
4.新建一个py文件，我的是pth.py，并写入如下代码

import torch
from collections import OrderedDict
import pickle
import os

device = 'cuda' if torch.cuda.is_available() else 'cpu'

if __name__ == '__main__':
    choices = ['yolov5s', 'yolov5l', 'yolov5m', 'yolov5x']
    modelfile = choices[0] + '.pt'
    utl_model = torch.load(modelfile, map_location=device)
    utl_param = utl_model['model'].model
    torch.save(utl_param.state_dict(), os.path.splitext(modelfile)[0] + '_param.pth')
    own_state = utl_param.state_dict()
    print(len(own_state))

    numpy_param = OrderedDict()
    for name in own_state:
        numpy_param[name] = own_state[name].data.cpu().numpy()
    print(len(numpy_param))
    with open(os.path.splitext(modelfile)[0] + '_numpy_param.pkl', 'wb') as fw:
        pickle.dump(numpy_param, fw)

运行，生成yolov5s_param.pth文件

四、pth转化为onnx

1.下载代码
---------------https://github.com/hpc203/yolov5-dnn-cpp-python-v2
2.复制文件
---------------将yolov5s_param.pth文件复制到convert-onnx文件夹下
3.修改参数
1）修改convert_onnx.py中的参数

output_onnx = args.net_type+'.onnx'
inputs = torch.randn(1, 3, 416, 416).to(device)

将尺寸修改为自己训练的尺寸即可

parser.add_argument('--num_classes', default=1, type=int)

修改为自己的类别数量，原始为80
2）修改coco.names文件，修改成自己的类别名称就可以
4.运行convert_onnx.py
生成onnx文件

五、使用onnx做推理

打开main_yolov5.py
1.修改输入尺寸

        self.inpWidth = 416
        self.inpHeight = 416

运行即可，效果如下。

总结一下，opencv、dnn模块对于部署来说比较方便，但是我这边速度有点慢，需要0.2s一张图像，比原始yolov5来说，慢了一点，继续找原因，提高速度！
问题记录：
1.输出结果全部为(nan,nan,nan…)
2.dnn模块无法读入onnx模型
在排除环境影响的前提下，基本上都是因为没有选对yolov5的代码，yolov5每天都在对代码进行更新，第三版和第四版，网络结构做了修改，网络层数都不一致了。所以一定要选用第四版的yolov5.

六、移植到树莓派4b

在电脑上运行，正常速度大概是100ms。但是树莓派居然要1.5s。看了一下别人做的，好像大家都是这样子，兴许yolov5真的不太适合移植到移动端和树莓派上做实时检测吧。

前前后后折腾了好几天，接下来准备试试YOLO-Fastest。