pytorch yolov5 训练自定义数据

---

---

前言

环境

python: 3.9.7

torch: 1.10.2

labelimg: 1.8.6

#yolov5 https://github.com/ultralytics/yolov5
#pytorch https://pytorch.org/
#labelimg https://github.com/tzutalin/labelImg
1
2
3

paddleocr 有三种模型 det 检测 cls 方向 rec 识别

---

一、创建环境

1. 安装miniconda

   https://blog.csdn.net/mtl1994/article/details/114968140
   1

2. 创建环境

   #linux 需要先 source
   conda create -n pytorch_yolov5 python=3.9.7  --channel https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
   1
   2

3. 安装完以后进入环境

   conda activate pytorch_yolov5
   1

二、安装环境

1. pytorch

   #选择对应的cuda/cpu版本
   pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
   1
   2

2. yolov5

   #下载源代码
   git clone https://github.com/ultralytics/yolov5  # clone
   cd yolov5
   #安装依赖
   python -m pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple/
   1
   2
   3
   4
   5

3. 测试一下

   import torch
   
   # Model
   model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5n - yolov5x6, custom
   
   # Images
   img = 'https://ultralytics.com/images/zidane.jpg'  # or file, Path, PIL, OpenCV, numpy, list
   
   # Inference
   results = model(img)
   
   # Results
   results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
   1
   2
   3
   4
   5
   6
   7
   8
   9
   10
   11
   12
   13

4. 选择模型

   https://github.com/ultralytics/yolov5/releases

   模型参数
   

   我使用的模型是

   yolov5x
   1

三、标注图片

我使用的labelimg

#安装
pip install labelimg
#打开
labelimg
1
2
3
4

标注完以后会有两个目录、一个存图片、一个存txt

四、训练

1.切分数据

def make_datasets(txt_path, img_path, out="./out", split_rate=0.05):
    txt_dir = Path(txt_path)
    img_dir = Path(img_path)
    out_dir = Path(out)
    dataset = []

    train_label = out_dir / "labels/train2007/"
    train_image = out_dir / "images/train2007/"
    test_label = out_dir / "labels/test2007/"
    test_image = out_dir / "images/test2007/"
    train_label.mkdir(parents=True, exist_ok=True)
    train_image.mkdir(parents=True, exist_ok=True)
    test_label.mkdir(parents=True, exist_ok=True)
    test_image.mkdir(parents=True, exist_ok=True)

    """
    过滤空txt
    """
    for item in txt_dir.rglob("*.txt"):
        if item.read_text() != "":
            dataset.append(item)
    """
    切分训练集，验证集
    """
    tv = random.sample(dataset, int(len(dataset) * split_rate))
    """
    组装数据
    """

    print(len(dataset))
    for item in dataset:
        for jpg in  img_dir.rglob(item.stem + ".jpg"):
            if item in tv:
                print(jpg, test_image / jpg.name)
                # jpg.replace(test_image / jpg.name)
                shutil.copy(str(jpg), test_image / jpg.name)
                shutil.copy(str(item), test_label / item.name)

                # item.replace(test_label / item.name)
            else:
                print(jpg, train_image / jpg.name)
                # jpg.replace(train_image / jpg.name)
                # item.replace(train_label / item.name)
                shutil.copy(str(jpg), train_image / jpg.name)
                shutil.copy(str(item), train_label / item.name)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

执行完以后，目录结构

2.修改训练的模型 yml

3.开始训练

nohup python train.py --img 640 --batch 32 --epochs 600 --data voc.yaml --weights yolov5s.pt --device 0,1,2,3 &
1

训练结果都保存在runs/train/递增的运行目录中，例如runs/train/exp2，runs/train/exp3

五、已经训练的模型，有了新数据需要继续训练

1.使用迁移学习

训练的时候指定 weights 为 上一次训练的输出

六、常用命令

1.训练

nohup python train.py --img 640 --batch 32 --epochs 600 --data wp_voc.yaml --weights runs/train/exp27/weights/best.pt --device 0,1,2,3 &
1

2.识别

python detect.py --weights runs/train/exp6/weights/best.pt --source ../datasets/infer/2022-2-24/
1

3.导出onnx

python export.py --weights yolov5s.pt --img 640 --batch 1  # export at 640x640 with batch size 1
1

总结