yolov5-5.0训练自己的数据集并用ncnn部署在安卓设备（一）

qq_39226865

已于 2023-09-04 16:35:56 修改

阅读量179

点赞数 1

文章标签： yolov5 pytorch 安卓

于 2023-09-04 16:26:01 首次发布

本文链接：https://blog.csdn.net/qq_39226865/article/details/132669106

版权

本篇文章梳理如何在pytorch上进行训练测试和导出onnx模型。

模型训练

参考文章：
https://blog.csdn.net/jiaoty19/article/details/125614783

总结其主要修改：
1）在models/yolov5s_mydataset.yaml中修改类别数量：

# parameters
nc: 1  # number of classes

2）在data/pupil.yaml中修改:

# 路径
train: /home/....../yolov5-5.0/mydataset/images/train/  
val: /home/......../yolov5-5.0/mydataset/images/val/  

#类别数
nc: 1

# 类别名称
names: [ 'pupil']

遇到的问题：
yolov5 没有SPPF：
https://blog.csdn.net/Steven_Cary/article/details/120886696
Float can‘t be cast to the desired output type long int：
https://blog.csdn.net/Thebest_jack/article/details/125649451
module ‘numpy‘ has no attribute ‘int‘ : 改成numpy1.21.6版本了
https://blog.csdn.net/jacke121/article/details/128373220
新的需求：修改图片大小：
为了加速模型推理速度，可以选择缩小输入图片尺寸。比如我的数据图片大小是640×480（W×H），缩小五倍，变成128*96，在train.py中修改如下代码：（要注意h,w的顺序）
```
parser.add_argument('--img-size', nargs='+', type=int, default=[96,128], help='[train, test] image sizes')  # h,w
```
推理的时候，在detect.py中的设置如下：
```
parser.add_argument('--img-size', type=int, default=128, help='inference size (pixels)')
```
这里128是在长和宽中取较大的一个。

parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp5/weights/best.pt', help='model.pt path(s)') # this

    parser.add_argument('--source', type=str, default='pupil-0.5-aug/images/val', help='source')  # file/folder, 0 for webcam  this  val里面
    # parser.add_argument('--source', type=str, default='/home/new/project/work/EYE/data-model/classfication-test/eyel', help='source')  #  this 所有的eye data
    # parser.add_argument('--source', type=str, default='/home/new/project/work/EYE/pic_wrong', help='source')  #  this 最新的wrong
    
    parser.add_argument('--img-size', type=int, default=128, help='inference size (pixels)')  # this 640  64--64,48
    parser.add_argument('--conf-thres', type=float, default=0.8, help='object confidence threshold')  # 0.25
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') # this 
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') # this default=True,
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')  # this
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')

修改完之后就可以运行查看自己模型的训练效果了

模型导出

在models/export.py中设置：

parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./runs/train/exp/weights/best.pt', help='weights path')  # 设置自己的参数路径
parser.add_argument('--img-size', nargs='+', type=int, default=[96,128], help='image size')  # h,w 设置输入图像大小
parser.add_argument('--batch-size', type=int, default=1, help='batch size')   # 设置为1
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')  # 虽然是用GPU训练的，但是用CPU推理就改成CPU
opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand

在models/yolo.py的**Detect(nn.Module)**类中加上self.training = True，如下：

def forward(self, x):
        # x = x.copy()  # for profiling
        z = []  # inference output
        self.training |= self.export
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
            self.training = True # 加上这句，非常重要 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            if not self.training:  # inference
                if self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)

                y = x[i].sigmoid()
                y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                z.append(y.view(bs, -1, self.no))

运行models/export.py生成onnx文件，后面需要用它转成ncnn文件
记录用自己数据集训练好的pt中保存的anchor大小，后面需要修改安卓程序中的jni代码

import torch
from models.experimental import attempt_load

model = attempt_load('runs/train/exp/weights/best.pt', map_location=torch.device('cpu')) # 模型路径
m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1]
print(m.anchor_grid)

本代码参考自：
https://blog.csdn.net/flyfish1986/article/details/117594265
文中详细介绍了anchor的设置和用法，可以学习一下~

至此pytorch部分就结束啦~~~