利用GPU（CUDA）跑YOLO V5（Windows环境）（二）

最新推荐文章于 2024-06-23 11:30:00 发布

星野悠一

最新推荐文章于 2024-06-23 11:30:00 发布

阅读量2.7k

点赞数 2

分类专栏： yolov5-learning 文章标签： python numpy pip conda matplotlib Powered by 金山文档

本文链接：https://blog.csdn.net/HoshinoYuuichi/article/details/128682361

版权

yolov5-learning 专栏收录该内容

2 篇文章 2 订阅

订阅专栏

书接上回：利用GPU（CUDA）跑YOLO V5（Windows环境）（一）

为什么不用CPU跑？->风扇起飞，性能堪忧，CPU跑一遍的时间CUDA能跑几十遍……

五、准备数据集

梳理整体框架

可按照以下方式在yolov5-master目录下配置空文件夹，后面要用。

寻找资源与数据集标注

可以自己准备或者从网上下载准备数据集，例如：猫狗识别之准备数据集，内有25000张图片。

我选择了图片的前200张，复制到yolov5-master\own_datas\images\train文件夹下作为训练集

然后打开PyCharm，在Pycharm终端中分别输入（yolov5环境下，没有的话需要先激活环境）：

pip install pyqt5
pip install labelme

这两个库是用来给数据集打标签的

安装完成后，在终端输入labelme，回车

我们将会在这里完成数据集的标注

点击左上角的Open Dir，选择yolov5-master\own_datas\images\train文件夹，就会出现训练集里的猫猫图片。右键选择Create Rectangle，

框选结束后，输入标签名cat，点击ok，这个标签就保存下来了。如果有多只猫猫，就继续框选

注：多只猫猫时应尽量避免交叉标注，以免之后训练检测时值为负而无法参与训练

整张图片框选完毕后，点击左侧的Next Image，根据提示把标注文件保存到路径yolov5-master\own_datas\labels\json中，文件的格式是.json。全部图片标注完成后可以在yolov5-master\own_datas\labels\json文件夹中看到与训练集图片数量相同且对应的.json文件

但我们需要将.json批量转换为.txt文件，因为yolov5只能识别.txt格式的标签

新建json2txt.py文件，输入以下代码并保存

import json
import os
name2id = {'cat': 0}  # 标签名称
def convert(img_size, box):
    dw = 1. / (img_size[0])
    dh = 1. / (img_size[1])
    x = (box[0] + box[2]) / 2.0 - 1
    y = (box[1] + box[3]) / 2.0 - 1
    w = box[2] - box[0]
    h = box[3] - box[1]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)
def decode_json(json_floder_path, json_name):
    txt_name = '*txt文件夹的路径' + json_name[0:-5] + '.txt'
    txt_file = open(txt_name, 'w')
    json_path = os.path.join(json_floder_path, json_name)
    data = json.load(open(json_path, 'r', encoding='gb2312', errors='ignore'))
    img_w = data['imageWidth']
    img_h = data['imageHeight']
    for i in data['shapes']: 
        label_name = i['label']
        if (i['shape_type'] == 'rectangle'):
            x1 = int(i['points'][0][0])
            y1 = int(i['points'][0][1])
            x2 = int(i['points'][1][0])
            y2 = int(i['points'][1][1]) 
            bb = (x1, y1, x2, y2)
            bbox = convert((img_w, img_h), bb)
            txt_file.write(str(name2id[label_name]) + " " + " ".join([str(a) for a in bbox]) + '\n')
if __name__ == "__main__": 
    json_floder_path = '*json文件夹的路径'
    json_names = os.listdir(json_floder_path)
    for json_name in json_names:
        decode_json(json_floder_path, json_name)

标*地方需要自行调整

运行json2txt.py，结束后可以在yolov5-master\own_datas\labels\txt文件夹中看到对应的.txt文件

最后把txt文件夹中的文件全部复制到yolov5-master\own_datas\labels\train文件夹中

六、修改配置文件

首先需要复制两个文件：

在yolov5-master\data路径下找到coco128.yaml文件，复制到yolov5-master\own_datas路径下，改名为cat.yaml

因为我们用的不是coco128数据集，而是我们自己的猫猫数据集。更改后便于我们理解。

在yolov5-master\models路径下找到yolov5s.yaml文件，同样复制到yolov5-master\own_datas路径下。当然我们也可以选择yolov5l.yaml。实际上，如果在GPU平台，我更推荐用l，因为我个人用s跑出来的效果实在是令人难以接受。

打开cat.yaml文件，需要修改这三行：

# path: ../datasets/coco128  # dataset root dir
train: own_datas/images/train  # train images (relative to 'path') 128 images
val: own_datas/images/train  # val images (relative to 'path') 128 images

然后修改这两行（nc是类别数，names是类别名称，因为这里只有一个类别标签：猫）：

nc: 1  # number of classes
names: ['cat']  # class names

七、训练

打开train.py文件，找到并修改下列代码：

parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')    #修改模型，yolov5s可改成yolov5l
parser.add_argument('--cfg', type=str, default='own_datas/yolov5s.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default='own_datas/cat.yaml', help='dataset.yaml path')

下面这一行的default可改可不改，主要是训练的迭代次数

parser.add_argument('--epochs', type=int, default=300)

最后一个要修改的是这一行的default，这个视个人配置而定：

parser.add_argument('--workers', type=int, default=4, help='maximum number of dataloader workers')

运行train.py即可开始训练

我的是200张图片300次迭代，20分钟基本完成

训练完成后可以在yolov5-master\runs\train\exp文件夹中查看训练相关的信息

但是在此之前你可能会遇到以下问题：

RuntimeError: result type Float can‘t be cast to the desired output type __int64

这时你需要找到5.0版报错的loss.py中最后那段for函数，将其整体替换为yolov5-master版中loss.py最后一段for函数即可正常运行

        for i in range(self.nl):
            anchors, shape = self.anchors[i], p[i].shape
            gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]]  # xyxy gain
 
            # Match targets to anchors
            t = targets * gain  # shape(3,n,7)
            if nt:
                # Matches
                r = t[..., 4:6] / anchors[:, None]  # wh ratio
                j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t']  # compare
                # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
                t = t[j]  # filter
 
                # Offsets
                gxy = t[:, 2:4]  # grid xy
                gxi = gain[[2, 3]] - gxy  # inverse
                j, k = ((gxy % 1 < g) & (gxy > 1)).T
                l, m = ((gxi % 1 < g) & (gxi > 1)).T
                j = torch.stack((torch.ones_like(j), j, k, l, m))
                t = t.repeat((5, 1, 1))[j]
                offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
            else:
                t = targets[0]
                offsets = 0
 
            # Define
            bc, gxy, gwh, a = t.chunk(4, 1)  # (image, class), grid xy, grid wh, anchors
            a, (b, c) = a.long().view(-1), bc.long().T  # anchors, image, class
            gij = (gxy - offsets).long()
            gi, gj = gij.T  # grid indices
 
            # Append
            indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))  # image, anchor, grid
            tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
            anch.append(anchors[a])  # anchors
            tcls.append(c)  # class

还可能有以下问题：

RuntimeError: result type Float can‘t be cast to the desired output type __int64

这时你需要将loss.py中的

gain = torch.ones(7, device=targets.device)

改为

gain = torch.ones(7, device=targets.device).long()

问题实现解决

还可能会有以下问题：

RuntimeError: Unable to find a valid cuDNN algorithm to run convolution

这就是GPU占用过高了，需要降低workers和batch-size的default值，并且调整虚拟内存大小

 parser.add_argument('--batch-size', type=int, default=24, help='total batch size for all GPUs, -1 for autobatch')

个人调试，3060 laptop + 30G虚拟内存时 workers为4，batch-size为24 较为合适

剩下基本上就没问题……了吧

之后你会得到一个best.pt，那个就是用于检测猫猫的模型

回到detect.py，调整以下代码：

parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'runs/train/exp/weights/best.pt', help='model path(s)')
parser.add_argument('--source', type=str, default=ROOT / 'own_datas/images/test', help='file/dir/URL/glob, 0 for webcam')

运行detect.py即可得到检测结果，目录在\runs\detect

附几张比较离谱的图，是用yolov5s跑的，作为本篇的结尾吧