• 基于深度学习的监控场景下是否佩戴安全帽检测（yolov5）

    内容：基于目标检测对图像中的人员是否佩戴安全帽进行检测

    具体要求：1) 使用Python编程语言，建议使用深度学习框架PyTorch；

    2) 完成自定义数据集的制作，基于目标检测方法在数据集上完成训练和验证，可使用开源框架；

    3) 使用数据集外的图像数据进行验证，对图像中的行人是否佩戴安全帽进行检测，检测到有人员未佩戴安全帽时改变人员检测框颜色，起到告警作用；

    环境 

    PyTorch  1.8.1

    Python  3.8

    Cuda  11.1

    RTX 2080 Ti * 1

    显存:11GB

    首先先给大家推荐一个GPU租用平台，关于为什么笔者本人会知道这种类型的平台，那必然是笔者我见多识广（说多了都是泪），你看我这破电脑运行的有多慢！！！（要我连跑60小时，达咩）

    6c76cea5a21247e5bca82f07e1f1b1e6.png

    好了，言归正传，今天给大家推荐的平台是AutoDL

    附上链接：AutoDL-品质GPU租用平台-租GPU就上AutoDL 

    价钱低，关键是注册就有10元体验券，像笔者租了个RTX 2080的卡，一小时0.85元，跑个简单的深度学习还是绰绰有余的。

    eba4440a22d445de99a9e605985b1a0c.png

    还可在其后方进行环境的配置，如图（我用的是pytorch1.8.1，python3.8，cuda11.1）

    072952aa752a48d19807dece8affbd89.png

    进入控制台后，选择jupyterlab进入

    c5a70ddcf1654352bd13927e951060b3.png

    可以先看看pytorch是否安装成功，进入终端打开输入python

    1. python
    2. import torch
    3. torch.__version__

    fdc8bc5d825a45188a391d460b71150e.png

    bea4b81748c1435db58bf8ce0d612ace.png

    基于Yolov5训练模型

    先把yolov5的框架下载出来，可以去github直接下，也可以用我上传的网盘资源

    链接：https://pan.baidu.com/s/1KIIjNOf449ueD8OdRJUH9g 
    提取码：0329

    通过jupyter内的文件上传，即可将下载好的zip文件上传，在终端找到对应压缩包，执行下述指令，即可解压，随后进入yolov5文件夹

     unzip yolov5-5.0.zip
    cd yolov5-5.0

    790e99537bf447cea411be22bf7a352b.png

    右击新建文件夹VOCdevkit，随后进入文件夹

    cf9cbdfc88f7429e9a18e2e951154c2d.png

    老方法，将数据集上传至文件包内

    196788b7b8a64e5d9d1f0049577b85b4.png

    数据集已全部标注，可直接食用

    链接：https://pan.baidu.com/s/1L4t9MQj48pT71mGaQlO4Tw 
    提取码：0329

    标注方法（可跳过）：

    （本机）使用labelimg进行标注

    pip install -i https://pypi.tuna.tsinghua.edu.cn/simple labelimg

    8e4e5aad7cd24af4814b6914bc77037c.png

    输入labelimg，打开界面

    2f5d1ea381e24b728abb68d44729d68c.png

    点击 Open Dir 选择数据集的保存路径，即VOCdevkit文件夹下的JPEGImages文件

    点击Change Save Dir选择保存标注好的文件路径，选择VOCdevkit文件夹下的Annotations文件

    在View中选中Auto Save mode，即可自动保存

    选中Create RectBox，框选出要识别的物体（如：hat），标注相应类别

    快捷键：A上一张，D下一张

    5cc0bfad5643409fbf9195f15bd232c5.png

    数据集处理（接着上两步标注过程前）

    同刚开始处理zip压缩包，unzip数据集即可

    unzip VOC2007.zip

    回退到yolov5主界面，进行环境配置

    pip install -r requirements.txt

    配置预训练权重，用yolov5s.pt即可

    链接：https://pan.baidu.com/s/1oopsiHUEclmO2U47FiuSJQ 
    提取码：0329

    下载完后，老方法上传至jupyter

    d1e2324545934aad8d050b816f21518c.png

    划分训练集和测试集

    新建prepare.py文件

    yolov5-5.0/prepare.py将VOC标签格式.xml转yolo格式.txt进行训练集和测试集的划分

    代码如下：

    1. import xml.etree.ElementTree as ET
    2. import pickle
    3. import os
    4. from os import listdir, getcwd
    5. from os.path import join
    6. import random
    7. from shutil import copyfile
    8.  
    9. classes = ["hat","person"]
    10.  
    11. TRAIN_RATIO = 80
    12.  
    13. def clear_hidden_files(path):
    14.     dir_list = os.listdir(path)
    15.     for i in dir_list:
    16.         abspath = os.path.join(os.path.abspath(path), i)
    17.         if os.path.isfile(abspath):
    18.             if i.startswith("._"):
    19.                 os.remove(abspath)
    20.         else:
    21.             clear_hidden_files(abspath)
    22.  
    23. def convert(size, box):
    24.     dw = 1./size[0]
    25.     dh = 1./size[1]
    26.     x = (box[0] + box[1])/2.0
    27.     y = (box[2] + box[3])/2.0
    28.     w = box[1] - box[0]
    29.     h = box[3] - box[2]
    30.     x = x*dw
    31.     w = w*dw
    32.     y = y*dh
    33.     h = h*dh
    34.     return (x,y,w,h)
    35.  
    36. def convert_annotation(image_id):
    37.     in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' %image_id)
    38.     out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' %image_id, 'w')
    39.     tree=ET.parse(in_file)
    40.     root = tree.getroot()
    41.     size = root.find('size')
    42.     w = int(size.find('width').text)
    43.     h = int(size.find('height').text)
    44.  
    45.     for obj in root.iter('object'):
    46.         difficult = obj.find('difficult').text
    47.         cls = obj.find('name').text
    48.         if cls not in classes or int(difficult) == 1:
    49.             continue
    50.         cls_id = classes.index(cls)
    51.         xmlbox = obj.find('bndbox')
    52.         b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
    53.         bb = convert((w,h), b)
    54.         out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    55.     in_file.close()
    56.     out_file.close()
    57.  
    58. wd = os.getcwd()
    59. wd = os.getcwd()
    60. data_base_dir = os.path.join(wd, "VOCdevkit/")
    61. if not os.path.isdir(data_base_dir):
    62.     os.mkdir(data_base_dir)
    63. work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
    64. if not os.path.isdir(work_sapce_dir):
    65.     os.mkdir(work_sapce_dir)
    66. annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
    67. if not os.path.isdir(annotation_dir):
    68.         os.mkdir(annotation_dir)
    69. clear_hidden_files(annotation_dir)
    70. image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
    71. if not os.path.isdir(image_dir):
    72.         os.mkdir(image_dir)
    73. clear_hidden_files(image_dir)
    74. yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
    75. if not os.path.isdir(yolo_labels_dir):
    76.         os.mkdir(yolo_labels_dir)
    77. clear_hidden_files(yolo_labels_dir)
    78. yolov5_images_dir = os.path.join(data_base_dir, "images/")
    79. if not os.path.isdir(yolov5_images_dir):
    80.         os.mkdir(yolov5_images_dir)
    81. clear_hidden_files(yolov5_images_dir)
    82. yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
    83. if not os.path.isdir(yolov5_labels_dir):
    84.         os.mkdir(yolov5_labels_dir)
    85. clear_hidden_files(yolov5_labels_dir)
    86. yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
    87. if not os.path.isdir(yolov5_images_train_dir):
    88.         os.mkdir(yolov5_images_train_dir)
    89. clear_hidden_files(yolov5_images_train_dir)
    90. yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
    91. if not os.path.isdir(yolov5_images_test_dir):
    92.         os.mkdir(yolov5_images_test_dir)
    93. clear_hidden_files(yolov5_images_test_dir)
    94. yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
    95. if not os.path.isdir(yolov5_labels_train_dir):
    96.         os.mkdir(yolov5_labels_train_dir)
    97. clear_hidden_files(yolov5_labels_train_dir)
    98. yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
    99. if not os.path.isdir(yolov5_labels_test_dir):
    100.         os.mkdir(yolov5_labels_test_dir)
    101. clear_hidden_files(yolov5_labels_test_dir)
    102.  
    103. train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
    104. test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
    105. train_file.close()
    106. test_file.close()
    107. train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
    108. test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
    109. list_imgs = os.listdir(image_dir) # list image files
    110. prob = random.randint(1, 100)
    111. print("Probability: %d" % prob)
    112. for i in range(0,len(list_imgs)):
    113.     path = os.path.join(image_dir,list_imgs[i])
    114.     if os.path.isfile(path):
    115.         image_path = image_dir + list_imgs[i]
    116.         voc_path = list_imgs[i]
    117.         (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
    118.         (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
    119.         annotation_name = nameWithoutExtention + '.xml'
    120.         annotation_path = os.path.join(annotation_dir, annotation_name)
    121.         label_name = nameWithoutExtention + '.txt'
    122.         label_path = os.path.join(yolo_labels_dir, label_name)
    123.     prob = random.randint(1, 100)
    124.     print("Probability: %d" % prob)
    125.     if(prob < TRAIN_RATIO): # train dataset
    126.         if os.path.exists(annotation_path):
    127.             train_file.write(image_path + '\n')
    128.             convert_annotation(nameWithoutExtention) # convert label
    129.             copyfile(image_path, yolov5_images_train_dir + voc_path)
    130.             copyfile(label_path, yolov5_labels_train_dir + label_name)
    131.     else: # test dataset
    132.         if os.path.exists(annotation_path):
    133.             test_file.write(image_path + '\n')
    134.             convert_annotation(nameWithoutExtention) # convert label
    135.             copyfile(image_path, yolov5_images_test_dir + voc_path)
    136.             copyfile(label_path, yolov5_labels_test_dir + label_name)
    137. train_file.close()
    138. test_file.close()
    139.  
    140.

    运行后产生文件夹：
    yolov5-5.0/VOCdevkit/images: 分为train和val两个文件夹，存放图片
    yolov5-5.0/VOCdevkit/labels: 分为train和val两个文件夹，存放.txt label

    4b7d4e3c95094c29a4de6c84e59db397.png

    配置文件的编写

    数据配置文件data/hat.yaml

    在data文件夹内新建一个hat.yaml 

    1. train: /root/yolov5-5.0/VOCdevkit/images/train/  # 16551 images
    2. val: /root/yolov5-5.0/VOCdevkit/images/val/  # 4952 images
    3.  
    4. # Classes
    5. nc: 2  # number of classes
    6. names: ["hat","person"]  # class names
    7.

    模型配置文件models/yolov5s_hat.yaml

    在models文件夹内，找到yolov5x.yaml文件夹，复制后重命名为yolov5s_hat.yaml

    修改nc即可

    1. # parameters
    2. nc: 2  # number of classes
    3. depth_multiple: 0.33  # model depth multiple
    4. width_multiple: 0.50  # layer channel multiple
    5.  
    6. # anchors
    7. anchors:
    8.   - [10,13, 16,30, 33,23]  # P3/8
    9.   - [30,61, 62,45, 59,119]  # P4/16
    10.   - [116,90, 156,198, 373,326]  # P5/32
    11.  
    12. # YOLOv5 backbone
    13. backbone:
    14.   # [from, number, module, args]
    15.   [[-1, 1, Focus, [64, 3]],  # 0-P1/2
    16.    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
    17.    [-1, 3, C3, [128]],
    18.    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
    19.    [-1, 9, C3, [256]],
    20.    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
    21.    [-1, 9, C3, [512]],
    22.    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
    23.    [-1, 1, SPP, [1024, [5, 9, 13]]],
    24.    [-1, 3, C3, [1024, False]],  # 9
    25.   ]
    26.  
    27. # YOLOv5 head
    28. head:
    29.   [[-1, 1, Conv, [512, 1, 1]],
    30.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
    31.    [[-1, 6], 1, Concat, [1]],  # cat backbone P4
    32.    [-1, 3, C3, [512, False]],  # 13
    33.  
    34.    [-1, 1, Conv, [256, 1, 1]],
    35.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
    36.    [[-1, 4], 1, Concat, [1]],  # cat backbone P3
    37.    [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
    38.  
    39.    [-1, 1, Conv, [256, 3, 2]],
    40.    [[-1, 14], 1, Concat, [1]],  # cat head P4
    41.    [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
    42.  
    43.    [-1, 1, Conv, [512, 3, 2]],
    44.    [[-1, 10], 1, Concat, [1]],  # cat head P5
    45.    [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
    46.  
    47.    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
    48.   ]

    训练模型

    先将yolov5-5.0/train.py（456行开始）中的部分内容配置

    一般是修改这些，也可直接复制下方代码替换

    d54201b018904af9b9e35dae8aa723fe.png

    1. import argparse
    2. import logging
    3. import math
    4. import os
    5. import random
    6. import time
    7. from copy import deepcopy
    8. from pathlib import Path
    9. from threading import Thread
    10.  
    11. import numpy as np
    12. import torch.distributed as dist
    13. import torch.nn as nn
    14. import torch.nn.functional as F
    15. import torch.optim as optim
    16. import torch.optim.lr_scheduler as lr_scheduler
    17. import torch.utils.data
    18. import yaml
    19. from torch.cuda import amp
    20. from torch.nn.parallel import DistributedDataParallel as DDP
    21. from torch.utils.tensorboard import SummaryWriter
    22. from tqdm import tqdm
    23.  
    24. import test  # import test.py to get mAP after each epoch
    25. from models.experimental import attempt_load
    26. from models.yolo import Model
    27. from utils.autoanchor import check_anchors
    28. from utils.datasets import create_dataloader
    29. from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \
    30.     fitness, strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \
    31.     check_requirements, print_mutation, set_logging, one_cycle, colorstr
    32. from utils.google_utils import attempt_download
    33. from utils.loss import ComputeLoss
    34. from utils.plots import plot_images, plot_labels, plot_results, plot_evolution
    35. from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, is_parallel
    36. from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume
    37.  
    38. logger = logging.getLogger(__name__)
    39.  
    40.  
    41. def train(hyp, opt, device, tb_writer=None):
    42.     logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
    43.     save_dir, epochs, batch_size, total_batch_size, weights, rank = \
    44.         Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
    45.  
    46.     # Directories
    47.     wdir = save_dir / 'weights'
    48.     wdir.mkdir(parents=True, exist_ok=True)  # make dir
    49.     last = wdir / 'last.pt'
    50.     best = wdir / 'best.pt'
    51.     results_file = save_dir / 'results.txt'
    52.  
    53.     # Save run settings
    54.     with open(save_dir / 'hyp.yaml', 'w') as f:
    55.         yaml.dump(hyp, f, sort_keys=False)
    56.     with open(save_dir / 'opt.yaml', 'w') as f:
    57.         yaml.dump(vars(opt), f, sort_keys=False)
    58.  
    59.     # Configure
    60.     plots = not opt.evolve  # create plots
    61.     cuda = device.type != 'cpu'
    62.     init_seeds(2 + rank)
    63.     with open(opt.data) as f:
    64.         data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
    65.     is_coco = opt.data.endswith('coco.yaml')
    66.  
    67.     # Logging- Doing this before checking the dataset. Might update data_dict
    68.     loggers = {'wandb': None}  # loggers dict
    69.     if rank in [-1, 0]:
    70.         opt.hyp = hyp  # add hyperparameters
    71.         run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
    72.         wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict)
    73.         loggers['wandb'] = wandb_logger.wandb
    74.         data_dict = wandb_logger.data_dict
    75.         if wandb_logger.wandb:
    76.             weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp  # WandbLogger might update weights, epochs if resuming
    77.  
    78.     nc = 1 if opt.single_cls else int(data_dict['nc'])  # number of classes
    79.     names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names']  # class names
    80.     assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
    81.  
    82.     # Model
    83.     pretrained = weights.endswith('.pt')
    84.     if pretrained:
    85.         with torch_distributed_zero_first(rank):
    86.             attempt_download(weights)  # download if not found locally
    87.         ckpt = torch.load(weights, map_location=device)  # load checkpoint
    88.         model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
    89.         exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else []  # exclude keys
    90.         state_dict = ckpt['model'].float().state_dict()  # to FP32
    91.         state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)  # intersect
    92.         model.load_state_dict(state_dict, strict=False)  # load
    93.         logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights))  # report
    94.     else:
    95.         model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
    96.     with torch_distributed_zero_first(rank):
    97.         check_dataset(data_dict)  # check
    98.     train_path = data_dict['train']
    99.     test_path = data_dict['val']
    100.  
    101.     # Freeze
    102.     freeze = []  # parameter names to freeze (full or partial)
    103.     for k, v in model.named_parameters():
    104.         v.requires_grad = True  # train all layers
    105.         if any(x in k for x in freeze):
    106.             print('freezing %s' % k)
    107.             v.requires_grad = False
    108.  
    109.     # Optimizer
    110.     nbs = 64  # nominal batch size
    111.     accumulate = max(round(nbs / total_batch_size), 1)  # accumulate loss before optimizing
    112.     hyp['weight_decay'] *= total_batch_size * accumulate / nbs  # scale weight_decay
    113.     logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")
    114.  
    115.     pg0, pg1, pg2 = [], [], []  # optimizer parameter groups
    116.     for k, v in model.named_modules():
    117.         if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
    118.             pg2.append(v.bias)  # biases
    119.         if isinstance(v, nn.BatchNorm2d):
    120.             pg0.append(v.weight)  # no decay
    121.         elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
    122.             pg1.append(v.weight)  # apply decay
    123.  
    124.     if opt.adam:
    125.         optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentum
    126.     else:
    127.         optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
    128.  
    129.     optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']})  # add pg1 with weight_decay
    130.     optimizer.add_param_group({'params': pg2})  # add pg2 (biases)
    131.     logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
    132.     del pg0, pg1, pg2
    133.  
    134.     # Scheduler https://arxiv.org/pdf/1812.01187.pdf
    135.     # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
    136.     if opt.linear_lr:
    137.         lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf']  # linear
    138.     else:
    139.         lf = one_cycle(1, hyp['lrf'], epochs)  # cosine 1->hyp['lrf']
    140.     scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
    141.     # plot_lr_scheduler(optimizer, scheduler, epochs)
    142.  
    143.     # EMA
    144.     ema = ModelEMA(model) if rank in [-1, 0] else None
    145.  
    146.     # Resume
    147.     start_epoch, best_fitness = 0, 0.0
    148.     if pretrained:
    149.         # Optimizer
    150.         if ckpt['optimizer'] is not None:
    151.             optimizer.load_state_dict(ckpt['optimizer'])
    152.             best_fitness = ckpt['best_fitness']
    153.  
    154.         # EMA
    155.         if ema and ckpt.get('ema'):
    156.             ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
    157.             ema.updates = ckpt['updates']
    158.  
    159.         # Results
    160.         if ckpt.get('training_results') is not None:
    161.             results_file.write_text(ckpt['training_results'])  # write results.txt
    162.  
    163.         # Epochs
    164.         start_epoch = ckpt['epoch'] + 1
    165.         if opt.resume:
    166.             assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs)
    167.         if epochs < start_epoch:
    168.             logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
    169.                         (weights, ckpt['epoch'], epochs))
    170.             epochs += ckpt['epoch']  # finetune additional epochs
    171.  
    172.         del ckpt, state_dict
    173.  
    174.     # Image sizes
    175.     gs = max(int(model.stride.max()), 32)  # grid size (max stride)
    176.     nl = model.model[-1].nl  # number of detection layers (used for scaling hyp['obj'])
    177.     imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size]  # verify imgsz are gs-multiples
    178.  
    179.     # DP mode
    180.     if cuda and rank == -1 and torch.cuda.device_count() > 1:
    181.         model = torch.nn.DataParallel(model)
    182.  
    183.     # SyncBatchNorm
    184.     if opt.sync_bn and cuda and rank != -1:
    185.         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
    186.         logger.info('Using SyncBatchNorm()')
    187.  
    188.     # Trainloader
    189.     dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
    190.                                             hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
    191.                                             world_size=opt.world_size, workers=opt.workers,
    192.                                             image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
    193.     mlc = np.concatenate(dataset.labels, 0)[:, 0].max()  # max label class
    194.     nb = len(dataloader)  # number of batches
    195.     assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
    196.  
    197.     # Process 0
    198.     if rank in [-1, 0]:
    199.         testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt,  # testloader
    200.                                        hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1,
    201.                                        world_size=opt.world_size, workers=opt.workers,
    202.                                        pad=0.5, prefix=colorstr('val: '))[0]
    203.  
    204.         if not opt.resume:
    205.             labels = np.concatenate(dataset.labels, 0)
    206.             c = torch.tensor(labels[:, 0])  # classes
    207.             # cf = torch.bincount(c.long(), minlength=nc) + 1.  # frequency
    208.             # model._initialize_biases(cf.to(device))
    209.             if plots:
    210.                 plot_labels(labels, names, save_dir, loggers)
    211.                 if tb_writer:
    212.                     tb_writer.add_histogram('classes', c, 0)
    213.  
    214.             # Anchors
    215.             if not opt.noautoanchor:
    216.                 check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
    217.             model.half().float()  # pre-reduce anchor precision
    218.  
    219.     # DDP mode
    220.     if cuda and rank != -1:
    221.         model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank,
    222.                     # nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698
    223.                     find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules()))
    224.  
    225.     # Model parameters
    226.     hyp['box'] *= 3. / nl  # scale to layers
    227.     hyp['cls'] *= nc / 80. * 3. / nl  # scale to classes and layers
    228.     hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl  # scale to image size and layers
    229.     hyp['label_smoothing'] = opt.label_smoothing
    230.     model.nc = nc  # attach number of classes to model
    231.     model.hyp = hyp  # attach hyperparameters to model
    232.     model.gr = 1.0  # iou loss ratio (obj_loss = 1.0 or iou)
    233.     model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
    234.     model.names = names
    235.  
    236.     # Start training
    237.     t0 = time.time()
    238.     nw = max(round(hyp['warmup_epochs'] * nb), 1000)  # number of warmup iterations, max(3 epochs, 1k iterations)
    239.     # nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
    240.     maps = np.zeros(nc)  # mAP per class
    241.     results = (0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
    242.     scheduler.last_epoch = start_epoch - 1  # do not move
    243.     scaler = amp.GradScaler(enabled=cuda)
    244.     compute_loss = ComputeLoss(model)  # init loss class
    245.     logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
    246.                 f'Using {dataloader.num_workers} dataloader workers\n'
    247.                 f'Logging results to {save_dir}\n'
    248.                 f'Starting training for {epochs} epochs...')
    249.     for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
    250.         model.train()
    251.  
    252.         # Update image weights (optional)
    253.         if opt.image_weights:
    254.             # Generate indices
    255.             if rank in [-1, 0]:
    256.                 cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc  # class weights
    257.                 iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  # image weights
    258.                 dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  # rand weighted idx
    259.             # Broadcast if DDP
    260.             if rank != -1:
    261.                 indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int()
    262.                 dist.broadcast(indices, 0)
    263.                 if rank != 0:
    264.                     dataset.indices = indices.cpu().numpy()
    265.  
    266.         # Update mosaic border
    267.         # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
    268.         # dataset.mosaic_border = [b - imgsz, -b]  # height, width borders
    269.  
    270.         mloss = torch.zeros(4, device=device)  # mean losses
    271.         if rank != -1:
    272.             dataloader.sampler.set_epoch(epoch)
    273.         pbar = enumerate(dataloader)
    274.         logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'labels', 'img_size'))
    275.         if rank in [-1, 0]:
    276.             pbar = tqdm(pbar, total=nb)  # progress bar
    277.         optimizer.zero_grad()
    278.         for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
    279.             ni = i + nb * epoch  # number integrated batches (since train start)
    280.             imgs = imgs.to(device, non_blocking=True).float() / 255.0  # uint8 to float32, 0-255 to 0.0-1.0
    281.  
    282.             # Warmup
    283.             if ni <= nw:
    284.                 xi = [0, nw]  # x interp
    285.                 # model.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
    286.                 accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())
    287.                 for j, x in enumerate(optimizer.param_groups):
    288.                     # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
    289.                     x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
    290.                     if 'momentum' in x:
    291.                         x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
    292.  
    293.             # Multi-scale
    294.             if opt.multi_scale:
    295.                 sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs  # size
    296.                 sf = sz / max(imgs.shape[2:])  # scale factor
    297.                 if sf != 1:
    298.                     ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
    299.                     imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
    300.  
    301.             # Forward
    302.             with amp.autocast(enabled=cuda):
    303.                 pred = model(imgs)  # forward
    304.                 loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
    305.                 if rank != -1:
    306.                     loss *= opt.world_size  # gradient averaged between devices in DDP mode
    307.                 if opt.quad:
    308.                     loss *= 4.
    309.  
    310.             # Backward
    311.             scaler.scale(loss).backward()
    312.  
    313.             # Optimize
    314.             if ni % accumulate == 0:
    315.                 scaler.step(optimizer)  # optimizer.step
    316.                 scaler.update()
    317.                 optimizer.zero_grad()
    318.                 if ema:
    319.                     ema.update(model)
    320.  
    321.             # Print
    322.             if rank in [-1, 0]:
    323.                 mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
    324.                 mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0)  # (GB)
    325.                 s = ('%10s' * 2 + '%10.4g' * 6) % (
    326.                     '%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
    327.                 pbar.set_description(s)
    328.  
    329.                 # Plot
    330.                 if plots and ni < 3:
    331.                     f = save_dir / f'train_batch{ni}.jpg'  # filename
    332.                     Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
    333.                     # if tb_writer:
    334.                     #     tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
    335.                     #     tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), [])  # add model graph
    336.                 elif plots and ni == 10 and wandb_logger.wandb:
    337.                     wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
    338.                                                   save_dir.glob('train*.jpg') if x.exists()]})
    339.  
    340.             # end batch ------------------------------------------------------------------------------------------------
    341.         # end epoch ----------------------------------------------------------------------------------------------------
    342.  
    343.         # Scheduler
    344.         lr = [x['lr'] for x in optimizer.param_groups]  # for tensorboard
    345.         scheduler.step()
    346.  
    347.         # DDP process 0 or single-GPU
    348.         if rank in [-1, 0]:
    349.             # mAP
    350.             ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights'])
    351.             final_epoch = epoch + 1 == epochs
    352.             if not opt.notest or final_epoch:  # Calculate mAP
    353.                 wandb_logger.current_epoch = epoch + 1
    354.                 results, maps, times = test.test(data_dict,
    355.                                                  batch_size=batch_size * 2,
    356.                                                  imgsz=imgsz_test,
    357.                                                  model=ema.ema,
    358.                                                  single_cls=opt.single_cls,
    359.                                                  dataloader=testloader,
    360.                                                  save_dir=save_dir,
    361.                                                  verbose=nc < 50 and final_epoch,
    362.                                                  plots=plots and final_epoch,
    363.                                                  wandb_logger=wandb_logger,
    364.                                                  compute_loss=compute_loss,
    365.                                                  is_coco=is_coco)
    366.  
    367.             # Write
    368.             with open(results_file, 'a') as f:
    369.                 f.write(s + '%10.4g' * 7 % results + '\n')  # append metrics, val_loss
    370.             if len(opt.name) and opt.bucket:
    371.                 os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
    372.  
    373.             # Log
    374.             tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss',  # train loss
    375.                     'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
    376.                     'val/box_loss', 'val/obj_loss', 'val/cls_loss',  # val loss
    377.                     'x/lr0', 'x/lr1', 'x/lr2']  # params
    378.             for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
    379.                 if tb_writer:
    380.                     tb_writer.add_scalar(tag, x, epoch)  # tensorboard
    381.                 if wandb_logger.wandb:
    382.                     wandb_logger.log({tag: x})  # W&B
    383.  
    384.             # Update best mAP
    385.             fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
    386.             if fi > best_fitness:
    387.                 best_fitness = fi
    388.             wandb_logger.end_epoch(best_result=best_fitness == fi)
    389.  
    390.             # Save model
    391.             if (not opt.nosave) or (final_epoch and not opt.evolve):  # if save
    392.                 ckpt = {'epoch': epoch,
    393.                         'best_fitness': best_fitness,
    394.                         'training_results': results_file.read_text(),
    395.                         'model': deepcopy(model.module if is_parallel(model) else model).half(),
    396.                         'ema': deepcopy(ema.ema).half(),
    397.                         'updates': ema.updates,
    398.                         'optimizer': optimizer.state_dict(),
    399.                         'wandb_id': wandb_logger.wandb_run.id if wandb_logger.wandb else None}
    400.  
    401.                 # Save last, best and delete
    402.                 torch.save(ckpt, last)
    403.                 if best_fitness == fi:
    404.                     torch.save(ckpt, best)
    405.                 if wandb_logger.wandb:
    406.                     if ((epoch + 1) % opt.save_period == 0 and not final_epoch) and opt.save_period != -1:
    407.                         wandb_logger.log_model(
    408.                             last.parent, opt, epoch, fi, best_model=best_fitness == fi)
    409.                 del ckpt
    410.  
    411.         # end epoch ----------------------------------------------------------------------------------------------------
    412.     # end training
    413.     if rank in [-1, 0]:
    414.         # Plots
    415.         if plots:
    416.             plot_results(save_dir=save_dir)  # save as results.png
    417.             if wandb_logger.wandb:
    418.                 files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
    419.                 wandb_logger.log({"Results": [wandb_logger.wandb.Image(str(save_dir / f), caption=f) for f in files
    420.                                               if (save_dir / f).exists()]})
    421.         # Test best.pt
    422.         logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
    423.         if opt.data.endswith('coco.yaml') and nc == 80:  # if COCO
    424.             for m in (last, best) if best.exists() else (last):  # speed, mAP tests
    425.                 results, _, _ = test.test(opt.data,
    426.                                           batch_size=batch_size * 2,
    427.                                           imgsz=imgsz_test,
    428.                                           conf_thres=0.001,
    429.                                           iou_thres=0.7,
    430.                                           model=attempt_load(m, device).half(),
    431.                                           single_cls=opt.single_cls,
    432.                                           dataloader=testloader,
    433.                                           save_dir=save_dir,
    434.                                           save_json=True,
    435.                                           plots=False,
    436.                                           is_coco=is_coco)
    437.  
    438.         # Strip optimizers
    439.         final = best if best.exists() else last  # final model
    440.         for f in last, best:
    441.             if f.exists():
    442.                 strip_optimizer(f)  # strip optimizers
    443.         if opt.bucket:
    444.             os.system(f'gsutil cp {final} gs://{opt.bucket}/weights')  # upload
    445.         if wandb_logger.wandb and not opt.evolve:  # Log the stripped model
    446.             wandb_logger.wandb.log_artifact(str(final), type='model',
    447.                                             name='run_' + wandb_logger.wandb_run.id + '_model',
    448.                                             aliases=['last', 'best', 'stripped'])
    449.         wandb_logger.finish_run()
    450.     else:
    451.         dist.destroy_process_group()
    452.     torch.cuda.empty_cache()
    453.     return results
    454.  
    455.  
    456. if __name__ == '__main__':
    457.     parser = argparse.ArgumentParser()
    458.     parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
    459.     parser.add_argument('--cfg', type=str, default='models/yolov5s_hat.yaml', help='model.yaml path')
    460.     parser.add_argument('--data', type=str, default='data/hat.yaml', help='data.yaml path')
    461.     parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
    462.     parser.add_argument('--epochs', type=int, default=50)
    463.     parser.add_argument('--batch-size', type=int, default=32, help='total batch size for all GPUs')
    464.     parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
    465.     parser.add_argument('--rect', action='store_true', help='rectangular training')
    466.     parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    467.     parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    468.     parser.add_argument('--notest', action='store_true', help='only test final epoch')
    469.     parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    470.     parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
    471.     parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    472.     parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
    473.     parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    474.     parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    475.     parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    476.     parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    477.     parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
    478.     parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    479.     parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
    480.     parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
    481.     parser.add_argument('--project', default='runs/train', help='save to project/name')
    482.     parser.add_argument('--entity', default=None, help='W&B entity')
    483.     parser.add_argument('--name', default='exp', help='save to project/name')
    484.     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    485.     parser.add_argument('--quad', action='store_true', help='quad dataloader')
    486.     parser.add_argument('--linear-lr', action='store_true', help='linear LR')
    487.     parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    488.     parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
    489.     parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
    490.     parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
    491.     parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
    492.     opt = parser.parse_args()
    493.  
    494.     # Set DDP variables
    495.     opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
    496.     opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1
    497.     set_logging(opt.global_rank)
    498.     if opt.global_rank in [-1, 0]:
    499.         check_git_status()
    500.         check_requirements()
    501.  
    502.     # Resume
    503.     wandb_run = check_wandb_resume(opt)
    504.     if opt.resume and not wandb_run:  # resume an interrupted run
    505.         ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run()  # specified or most recent path
    506.         assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
    507.         apriori = opt.global_rank, opt.local_rank
    508.         with open(Path(ckpt).parent.parent / 'opt.yaml') as f:
    509.             opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader))  # replace
    510.         opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = '', ckpt, True, opt.total_batch_size, *apriori  # reinstate
    511.         logger.info('Resuming training from %s' % ckpt)
    512.     else:
    513.         # opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
    514.         opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(opt.cfg), check_file(opt.hyp)  # check files
    515.         assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
    516.         opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size)))  # extend to 2 sizes (train, test)
    517.         opt.name = 'evolve' if opt.evolve else opt.name
    518.         opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok | opt.evolve)  # increment run
    519.  
    520.     # DDP mode
    521.     opt.total_batch_size = opt.batch_size
    522.     device = select_device(opt.device, batch_size=opt.batch_size)
    523.     if opt.local_rank != -1:
    524.         assert torch.cuda.device_count() > opt.local_rank
    525.         torch.cuda.set_device(opt.local_rank)
    526.         device = torch.device('cuda', opt.local_rank)
    527.         dist.init_process_group(backend='nccl', init_method='env://')  # distributed backend
    528.         assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count'
    529.         opt.batch_size = opt.total_batch_size // opt.world_size
    530.  
    531.     # Hyperparameters
    532.     with open(opt.hyp) as f:
    533.         hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps
    534.  
    535.     # Train
    536.     logger.info(opt)
    537.     if not opt.evolve:
    538.         tb_writer = None  # init loggers
    539.         if opt.global_rank in [-1, 0]:
    540.             prefix = colorstr('tensorboard: ')
    541.             logger.info(f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/")
    542.             tb_writer = SummaryWriter(opt.save_dir)  # Tensorboard
    543.         train(hyp, opt, device, tb_writer)
    544.  
    545.     # Evolve hyperparameters (optional)
    546.     else:
    547.         # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
    548.         meta = {'lr0': (1, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
    549.                 'lrf': (1, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
    550.                 'momentum': (0.3, 0.6, 0.98),  # SGD momentum/Adam beta1
    551.                 'weight_decay': (1, 0.0, 0.001),  # optimizer weight decay
    552.                 'warmup_epochs': (1, 0.0, 5.0),  # warmup epochs (fractions ok)
    553.                 'warmup_momentum': (1, 0.0, 0.95),  # warmup initial momentum
    554.                 'warmup_bias_lr': (1, 0.0, 0.2),  # warmup initial bias lr
    555.                 'box': (1, 0.02, 0.2),  # box loss gain
    556.                 'cls': (1, 0.2, 4.0),  # cls loss gain
    557.                 'cls_pw': (1, 0.5, 2.0),  # cls BCELoss positive_weight
    558.                 'obj': (1, 0.2, 4.0),  # obj loss gain (scale with pixels)
    559.                 'obj_pw': (1, 0.5, 2.0),  # obj BCELoss positive_weight
    560.                 'iou_t': (0, 0.1, 0.7),  # IoU training threshold
    561.                 'anchor_t': (1, 2.0, 8.0),  # anchor-multiple threshold
    562.                 'anchors': (2, 2.0, 10.0),  # anchors per output grid (0 to ignore)
    563.                 'fl_gamma': (0, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
    564.                 'hsv_h': (1, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
    565.                 'hsv_s': (1, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
    566.                 'hsv_v': (1, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
    567.                 'degrees': (1, 0.0, 45.0),  # image rotation (+/- deg)
    568.                 'translate': (1, 0.0, 0.9),  # image translation (+/- fraction)
    569.                 'scale': (1, 0.0, 0.9),  # image scale (+/- gain)
    570.                 'shear': (1, 0.0, 10.0),  # image shear (+/- deg)
    571.                 'perspective': (0, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
    572.                 'flipud': (1, 0.0, 1.0),  # image flip up-down (probability)
    573.                 'fliplr': (0, 0.0, 1.0),  # image flip left-right (probability)
    574.                 'mosaic': (1, 0.0, 1.0),  # image mixup (probability)
    575.                 'mixup': (1, 0.0, 1.0)}  # image mixup (probability)
    576.  
    577.         assert opt.local_rank == -1, 'DDP mode not implemented for --evolve'
    578.         opt.notest, opt.nosave = True, True  # only test/save final epoch
    579.         # ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indices
    580.         yaml_file = Path(opt.save_dir) / 'hyp_evolved.yaml'  # save best result here
    581.         if opt.bucket:
    582.             os.system('gsutil cp gs://%s/evolve.txt .' % opt.bucket)  # download evolve.txt if exists
    583.  
    584.         for _ in range(300):  # generations to evolve
    585.             if Path('evolve.txt').exists():  # if evolve.txt exists: select best hyps and mutate
    586.                 # Select parent(s)
    587.                 parent = 'single'  # parent selection method: 'single' or 'weighted'
    588.                 x = np.loadtxt('evolve.txt', ndmin=2)
    589.                 n = min(5, len(x))  # number of previous results to consider
    590.                 x = x[np.argsort(-fitness(x))][:n]  # top n mutations
    591.                 w = fitness(x) - fitness(x).min()  # weights
    592.                 if parent == 'single' or len(x) == 1:
    593.                     # x = x[random.randint(0, n - 1)]  # random selection
    594.                     x = x[random.choices(range(n), weights=w)[0]]  # weighted selection
    595.                 elif parent == 'weighted':
    596.                     x = (x * w.reshape(n, 1)).sum(0) / w.sum()  # weighted combination
    597.  
    598.                 # Mutate
    599.                 mp, s = 0.8, 0.2  # mutation probability, sigma
    600.                 npr = np.random
    601.                 npr.seed(int(time.time()))
    602.                 g = np.array([x[0] for x in meta.values()])  # gains 0-1
    603.                 ng = len(meta)
    604.                 v = np.ones(ng)
    605.                 while all(v == 1):  # mutate until a change occurs (prevent duplicates)
    606.                     v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
    607.                 for i, k in enumerate(hyp.keys()):  # plt.hist(v.ravel(), 300)
    608.                     hyp[k] = float(x[i + 7] * v[i])  # mutate
    609.  
    610.             # Constrain to limits
    611.             for k, v in meta.items():
    612.                 hyp[k] = max(hyp[k], v[1])  # lower limit
    613.                 hyp[k] = min(hyp[k], v[2])  # upper limit
    614.                 hyp[k] = round(hyp[k], 5)  # significant digits
    615.  
    616.             # Train mutation
    617.             results = train(hyp.copy(), opt, device)
    618.  
    619.             # Write mutation results
    620.             print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
    621.  
    622.         # Plot results
    623.         plot_evolution(yaml_file)
    624.         print(f'Hyperparameter evolution complete. Best results saved as: {yaml_file}\n'
    625.               f'Command to train a new model with these hyperparameters: $ python train.py --hyp {yaml_file}')

    终端执行train.py即可

    python train.py

    75b02cf52b334d0ea9fa3f8dd13e66d4.png

    训练好的模型是runs/train/exp3/best.pt

    最后一次模型是runs/train/exp3/last.pt

    训练好后检测图片

    先修改detect.py（149行开始）内的部分内容

    主要修改以下参数，也可直接复制下方代码替换

    6a900f4797854b218308a2121dbbdcc2.png

    1. import argparse
    2. import time
    3. from pathlib import Path
    4.  
    5. import cv2
    6. import torch
    7. import torch.backends.cudnn as cudnn
    8. from numpy import random
    9.  
    10. from models.experimental import attempt_load
    11. from utils.datasets import LoadStreams, LoadImages
    12. from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
    13.     scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
    14. from utils.plots import plot_one_box
    15. from utils.torch_utils import select_device, load_classifier, time_synchronized
    16.  
    17.  
    18. def detect(save_img=False):
    19.     source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
    20.     save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
    21.     webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
    22.         ('rtsp://', 'rtmp://', 'http://', 'https://'))
    23.  
    24.     # Directories
    25.     save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run
    26.     (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
    27.  
    28.     # Initialize
    29.     set_logging()
    30.     device = select_device(opt.device)
    31.     half = device.type != 'cpu'  # half precision only supported on CUDA
    32.  
    33.     # Load model
    34.     model = attempt_load(weights, map_location=device)  # load FP32 model
    35.     stride = int(model.stride.max())  # model stride
    36.     imgsz = check_img_size(imgsz, s=stride)  # check img_size
    37.     if half:
    38.         model.half()  # to FP16
    39.  
    40.     # Second-stage classifier
    41.     classify = False
    42.     if classify:
    43.         modelc = load_classifier(name='resnet101', n=2)  # initialize
    44.         modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
    45.  
    46.     # Set Dataloader
    47.     vid_path, vid_writer = None, None
    48.     if webcam:
    49.         view_img = check_imshow()
    50.         cudnn.benchmark = True  # set True to speed up constant image size inference
    51.         dataset = LoadStreams(source, img_size=imgsz, stride=stride)
    52.     else:
    53.         dataset = LoadImages(source, img_size=imgsz, stride=stride)
    54.  
    55.     # Get names and colors
    56.     names = model.module.names if hasattr(model, 'module') else model.names
    57.     colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
    58.  
    59.     # Run inference
    60.     if device.type != 'cpu':
    61.         model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    62.     t0 = time.time()
    63.     for path, img, im0s, vid_cap in dataset:
    64.         img = torch.from_numpy(img).to(device)
    65.         img = img.half() if half else img.float()  # uint8 to fp16/32
    66.         img /= 255.0  # 0 - 255 to 0.0 - 1.0
    67.         if img.ndimension() == 3:
    68.             img = img.unsqueeze(0)
    69.  
    70.         # Inference
    71.         t1 = time_synchronized()
    72.         pred = model(img, augment=opt.augment)[0]
    73.  
    74.         # Apply NMS
    75.         pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
    76.         t2 = time_synchronized()
    77.  
    78.         # Apply Classifier
    79.         if classify:
    80.             pred = apply_classifier(pred, modelc, img, im0s)
    81.  
    82.         # Process detections
    83.         for i, det in enumerate(pred):  # detections per image
    84.             if webcam:  # batch_size >= 1
    85.                 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
    86.             else:
    87.                 p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
    88.  
    89.             p = Path(p)  # to Path
    90.             save_path = str(save_dir / p.name)  # img.jpg
    91.             txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
    92.             s += '%gx%g ' % img.shape[2:]  # print string
    93.             gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
    94.             if len(det):
    95.                 # Rescale boxes from img_size to im0 size
    96.                 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
    97.  
    98.                 # Print results
    99.                 for c in det[:, -1].unique():
    100.                     n = (det[:, -1] == c).sum()  # detections per class
    101.                     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
    102.  
    103.                 # Write results
    104.                 for *xyxy, conf, cls in reversed(det):
    105.                     if save_txt:  # Write to file
    106.                         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
    107.                         line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
    108.                         with open(txt_path + '.txt', 'a') as f:
    109.                             f.write(('%g ' * len(line)).rstrip() % line + '\n')
    110.  
    111.                     if save_img or view_img:  # Add bbox to image
    112.                         label = f'{names[int(cls)]} {conf:.2f}'
    113.                         plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
    114.  
    115.             # Print time (inference + NMS)
    116.             print(f'{s}Done. ({t2 - t1:.3f}s)')
    117.  
    118.             # Stream results
    119.             if view_img:
    120.                 cv2.imshow(str(p), im0)
    121.                 cv2.waitKey(1)  # 1 millisecond
    122.  
    123.             # Save results (image with detections)
    124.             if save_img:
    125.                 if dataset.mode == 'image':
    126.                     cv2.imwrite(save_path, im0)
    127.                 else:  # 'video' or 'stream'
    128.                     if vid_path != save_path:  # new video
    129.                         vid_path = save_path
    130.                         if isinstance(vid_writer, cv2.VideoWriter):
    131.                             vid_writer.release()  # release previous video writer
    132.                         if vid_cap:  # video
    133.                             fps = vid_cap.get(cv2.CAP_PROP_FPS)
    134.                             w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    135.                             h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    136.                         else:  # stream
    137.                             fps, w, h = 30, im0.shape[1], im0.shape[0]
    138.                             save_path += '.mp4'
    139.                         vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
    140.                     vid_writer.write(im0)
    141.  
    142.     if save_txt or save_img:
    143.         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
    144.         print(f"Results saved to {save_dir}{s}")
    145.  
    146.     print(f'Done. ({time.time() - t0:.3f}s)')
    147.  
    148.  
    149. if __name__ == '__main__':
    150.     parser = argparse.ArgumentParser()
    151.     parser.add_argument('--weights', nargs='+', type=str, default='best.pt', help='model.pt path(s)')
    152.     parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
    153.     parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    154.     parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    155.     parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    156.     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    157.     parser.add_argument('--view-img', action='store_true', help='display results')
    158.     parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    159.     parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    160.     parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    161.     parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    162.     parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    163.     parser.add_argument('--augment', action='store_true', help='augmented inference')
    164.     parser.add_argument('--update', action='store_true', help='update all models')
    165.     parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    166.     parser.add_argument('--name', default='exp', help='save results to project/name')
    167.     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    168.     opt = parser.parse_args()
    169.     print(opt)
    170.     check_requirements(exclude=('pycocotools', 'thop'))
    171.  
    172.     with torch.no_grad():
    173.         if opt.update:  # update all models (to fix SourceChangeWarning)
    174.             for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
    175.                 detect()
    176.                 strip_optimizer(opt.weights)
    177.         else:
    178.             detect()

    运行detect.py

    python detect.py --weight runs/train/exp3/weights/best.pt   --source  VOCdevkit/images/train/000002.jpg

    或者

    python detect.py

    755c2a7d34484dae8a60ab57d8cfe4fa.png

    识别出佩戴安全帽和未佩戴安全帽，完成简单案例的流程（第一次用yolo框架，如有不妥请见谅）！ 

     

     

  • 相关阅读:
    六、回归与聚类算法 - K-means算法
    Oracle的listagg的用法和例子
    jvm 内存泄露、内存溢出、栈溢出区别
    Redis基础架构
    【四旋翼飞行器】模拟四旋翼飞行器的平移和旋转动力学（Simulink仿真实现）
    MySQL 8.0 新特性
    C++ 学习笔记（五）（泛型算法篇）
    Redis基础
    【ROOTFS】1-构建rootfs与nfs调试
    【通信工程笔记】【终端与业务-第十二章】市场营销计划、实施和控制
  • 原文地址：https://blog.csdn.net/yyfloveqcw/article/details/125628085