yolov5训练时 gpu mem 0g

本文以训练NWPU VHR-10数据集为例，NWPU VHR-10遥感数据集是由西北工业大学公布的用于遥感图像目标检测的公开数据集，包含10类地物目标共800张遥感图像，具体有airplane、ship 、storage tank 、baseball diamond、tennis court、basketball court、ground track field、harbor、bridge、vehicle等十种类别。

一，在官方的开源的下载YOLOV5，一般下载zip压缩包即可，下载链接为GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

二，接下来则是配置环境，按照下载的文件中的requirements.txt里面提到的就行，这一块可以去线上搜搜配置的流程，我就不过多赘述了，如果实在不想自己动手，就去咸鱼上十来块钱找人配置一下，保证又快又好。有个小建议，很多时候下载一些库函数很慢，因此大家可以使用一些镜像源。

三，数据集的制作，对于大多数的数据集来说一般是VOC格式的，在这里也就逐步的给大家具体介绍一下如何去做符合官方代码训练的数据集格式。一般来讲，我们下载的数据集的分布大概如下图所示：

其中Annotations里面存的是.xml的标签文件，而JPEGImages则是存放的图片，并且标签名和图片名一一对应，接下来，我们在打开的文件夹中先建立一个新的文件夹data1，如下所示：

之后直接将下载好的VOC格式的数据集放在此文件夹下（其中ImageSets可有可无，不影响）。

接下来我们需要做两个步骤，1：将VOC数据集转为YOLO格式的数据集。2：对转换好的数据集进行划分。

这里我们使用两个脚本，分别是用来转换标签的类别和对数据集进行划分。

1.标签类别转换，这里要记住classes里面写的是数据集中的类别。

import xml.etree.ElementTree as ET
import os, cv2
import numpy as np
from os import listdir
from os.path import join

classes = ['airplane','baseball diamond','basketball court','bridge','ship','storage tank','ground track field','harbor','vehicle','tennis court']#这个地方是数据集的类别名称

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(xmlpath, xmlname):
    with open(xmlpath, "r", encoding='utf-8') as in_file:
        txtname = xmlname[:-4] + '.txt'
        txtfile = os.path.join(txtpath, txtname)
        tree = ET.parse(in_file)
        root = tree.getroot()
        filename = root.find('filename')
        img = cv2.imdecode(np.fromfile('{}/{}.{}'.format(imgpath, xmlname[:-4], postfix), np.uint8), cv2.IMREAD_COLOR)
        h, w = img.shape[:2]
        res = []
        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes:
                classes.append(cls)
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            bb = convert((w, h), b)
            res.append(str(cls_id) + " " + " ".join([str(a) for a in bb]))
        if len(res) != 0:
            with open(txtfile, 'w+') as f:
                f.write('\n'.join(res))


if __name__ == "__main__":
    postfix = 'jpg'
    imgpath = 'VOCdevkit/JPEGImages'
    xmlpath = 'VOCdevkit/Annotations'
    txtpath = 'VOCdevkit/txt'
    
    if not os.path.exists(txtpath):
        os.makedirs(txtpath, exist_ok=True)
    
    list = os.listdir(xmlpath)
    error_file_list = []
    for i in range(0, len(list)):
        try:
            path = os.path.join(xmlpath, list[i])
            if ('.xml' in path) or ('.XML' in path):
                convert_annotation(path, list[i])
                print(f'file {list[i]} convert success.')
            else:
                print(f'file {list[i]} is not xml format.')
        except Exception as e:
            print(f'file {list[i]} convert error.')
            print(f'error message:\n{e}')
            error_file_list.append(list[i])
    print(f'this file convert failure\n{error_file_list}')
    print(f'Dataset Classes:{classes}')

2，数据集比例划分

import os, shutil
from sklearn.model_selection import train_test_split

val_size = 0.1
test_size = 0.2
postfix = 'jpg'
imgpath = 'VOCdevkit/JPEGImages'
txtpath = 'VOCdevkit/txt'

os.makedirs('images/train', exist_ok=True)
os.makedirs('images/val', exist_ok=True)
os.makedirs('images/test', exist_ok=True)
os.makedirs('labels/train', exist_ok=True)
os.makedirs('labels/val', exist_ok=True)
os.makedirs('labels/test', exist_ok=True)

listdir = os.listdir(txtpath)
train, test = train_test_split(listdir, test_size=test_size, shuffle=True, random_state=0)
train, val = train_test_split(train, test_size=val_size, shuffle=True, random_state=0)

for i in train:
    shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'images/train/{}.{}'.format(i[:-4], postfix))
    shutil.copy('{}/{}'.format(txtpath, i), 'labels/train/{}'.format(i))

for i in val:
    shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'images/val/{}.{}'.format(i[:-4], postfix))
    shutil.copy('{}/{}'.format(txtpath, i), 'labels/val/{}'.format(i))

for i in test:
    shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'images/test/{}.{}'.format(i[:-4], postfix))
    shutil.copy('{}/{}'.format(txtpath, i), 'labels/test/{}'.format(i))

之后把这两个脚本放在data1文件夹下，并依次运行，记住不要放错目录级别。

依次运行后，data1文件夹下的格式便有了变化，其中新生成的images是划分后的图片，labels则是划分后并转换好的标签文件。

四，文件中的一些配置修改

1，在我们下载的文件夹中，有一个是data，其中保存着一些数据集的信息就是不同的yaml文件，然后我们可以复制其中的一个，并进行重命名操作，并且在其中进行配置。我这里直接复制并重命名为NWPU.yaml。

这里面需要注意的是，有三处需要进行修改，第一就是train和val的图片文件路径，因为我存放在

train: data1/images/train val: data1/images/val

所以修改为这样，第二就是nc是数据集的类别数，我使用的数据集是十个类别，因此这里是10，最后names是数据集不同类别的名字，并且这里要和前面对标签类型转换的classes中的类别顺序一致。

2，在models下模型的配置文件进行修改，因为我训练的是yolov5s，所以这里只需要对这一文件中的类别数进行修改，修改为10即可。

3，对train.py中的一些训练参数进行配置，首先是文件的权重，这里需要自己下载权重文件，并且放到根目录下。其次是--cfg，这里放的是模型的参数，我们使用的是s版本，最后是数据集的信息，因为我们训练的是NWPU数据集，因此这里放的文件路径为之前配置的

data/NWPU.yaml。

5，在个人进行使用的时候出现了一些小问题，例如报错：Can‘t get attribute ‘SPPF‘ on ＜module ‘models.common‘ from ‘/yolov5-5.0这种，错误的原因实际上是下载的权重文件和模型不是完全配套的，因此一个快速的解决方法就是在models/common.py中增加一个类：

import warnings

class SPPF(nn.Module):
    def __init__(self, c1, c2, k=5): 
        super().__init__()
        c_ = c1 // 2  
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * 4, c2, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore') 
            y1 = self.m(x)
            y2 = self.m(y1)
            return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))

此外还遇到了一个问题就是RuntimeError: Given groups=1, weight of size [512, 1024, 1, 1], expected input[1, 512, 8, 8] to have 1024 channels, but got 512 channels instead，这个其实也好解决，只需要在train.py中修改--cfg这一行，将默认权重改了即可。

这就是我的全部分享了，下一期可能会出yolo其他系列的训练方法，也可以出一些具体的教大家如何改一些网络结构的方法，大家有什么想看的，可以在评论区说一声。