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版,其他的版本对应网络模型不一样。
Opencv、dnn部署自己的Yolov5模型记录_第1张图片
Opencv、dnn部署自己的Yolov5模型记录_第2张图片

三、训练自己的模型

先插一句话:一定要配置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根目录下。Opencv、dnn部署自己的Yolov5模型记录_第3张图片
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

Opencv、dnn部署自己的Yolov5模型记录_第4张图片
在电脑上运行,正常速度大概是100ms。但是树莓派居然要1.5s。看了一下别人做的,好像大家都是这样子,兴许yolov5真的不太适合移植到移动端和树莓派上做实时检测吧。
Opencv、dnn部署自己的Yolov5模型记录_第5张图片
前前后后折腾了好几天,接下来准备试试YOLO-Fastest。

你可能感兴趣的:(yolo系列算法学习,pytorch,深度学习,神经网络)