一、Yolov5简介
2020年6月25日,Ultralytics发布了YOLOV5 的第一个正式版本,其性能与YOLO V4不相伯仲,同样也是现今最先进的对象检测技术,并在推理速度上是目前最强,yolov5按大小分为四个模型yolov5s、yolov5m、yolov5l、yolov5x。其中的复杂的网络结构、数学基础在这里就不一一介绍(太复杂,笔者也只能看个大概,很难说清楚),在这里,引用另一个博主的Yolov5的网络结构图:Yolov5网络结构图,以及一篇流程图:Yolov5操作流程图
YOLOv5是YOLO系列的一个延申,您也可以看作是基于YOLOv3、YOLOv4的改进作品。YOLOv5没有相应的论文说明,但是作者在Github上积极地开放源代码,通过对源码分析,我们也能很快地了解YOLOv5的网络架构和工作原理。
二、项目背景
当前新冠疫情仍然严重,在公众场合需要佩戴口罩已经成为常识。新型冠状病毒的主要传播途径就是飞沫传播,戴上口罩就可以有效的阻隔病毒的传播。口罩是预防呼吸道传染病的重要防线,可以降低新型冠状病毒感染风险。口罩不仅可以防止病人喷射飞沫,降低飞沫量和喷射速度,还可以阻挡含病毒的飞沫核,防止佩戴者吸入。有研究显示,只要双方都佩戴口罩且间隔1.8米以上,造成感染的几率几乎为0。
但是,在我们周围总有人不喜欢戴口罩,无论是进出商场、教室、街道、地下停车场等公共场所,还是在人员密集的会议室里,他们都不喜欢口罩的“束缚”。运用Yolov5训练出来的口罩检测模型进行检测,就能准确实时的找到哪些人带了口罩、哪些人没带。可以做的定点提醒,或者是阻止他出入公共场所。节省了人力,大幅提高效率。
三、检测效果
因为我是拿CPU运行的,速度很慢,epoch取了5次,即每张图片学习了5次,一共有1200组训练数据,训练了三个半小时,之后在项目的实际应用的时候会考虑修改为GPU运行,这样速度可以提高很多。我们直接看模型的检测效果及视频的检测效果:
通过上述例子可以看到,仅经过五次学习,识别的精度已经很高了,再一次感叹Yolov5的强大!
四、数据集处理 添加标签
训练的数据集总共有1200张戴口罩的和没戴口罩的照片,验证集有400张照片,对应的标签也已经存在相应的文件夹下。这里重点讲解下数据集标签的标注,我觉得这是Yolov5特别亲民的一个地方,也是他的强大之处——你可以标记你任何想标记的地方!
可以在你的虚拟环境中进入labelImg,这是他的界面。左侧open Dir可以打开数据集的文件夹,Change Save Dir是你的保存路径,Next和Prev Image分别是上一张和下一张图片。重点来了:Create RectBox绘制一个矩形框将你需要框选的对象框选出来,并添加标签。我这里框选了红色衣服的女士,并给她添加标签mask,说明她带了口罩。接着标记第二个人,直到所有的人都标记完之后,可以得到一个该图片的txt文件:
我们一行一行看,每一行代表着图片里面的一个人,第一个数字 0说明是第一类,在这里就是带了口罩,后两个数字是矩形框的的中心点坐标,最后两个数字是矩形框的长宽。在训练模型时要将原图片和对应的标签一起传入进去作为一组训练集,这样机器才能够慢慢认识一个人有没有戴口罩。
五、结果分析
从图中可以看出,仅训练了5次(即每张图片机器学习了5遍), 对mask的识别精度可以达到0.564,对face的识别精度可以达到0.904,实在是恐怖!
六、总结
Yolov5真的是一个利器,确实要比CNN强大很多,里面复杂的神经网络函数复杂交错。这个开源项目让我再一次感受到了机器学习的强大,你可以让计算机认出他任何想要认出的东西,且识别的精度很高。比如说
- 火灾检测,一片森林或楼道只要一有火焰的模样,计算机就能分辨出并报警,这样可以有效地减少经济损失甚至挽救生命。
- 安检检测,现在的地铁飞机安检,都是人观察扫描仪扫描出来的图像,看有没有可疑物体,如果用Yolov5训练出识别危险物品的模型,就能减少大量人力,且准确率可能比人还要高。
- 无人驾驶。通过yolov5检测车道和行人车辆,控制整个车子运转,只要有足够多的数据集,这个领域还是很值得探索的。
这次实战让我深深明白了:人工智能 = 人工+智能,先有人工才有智能,人工筛选标注数据集甚至会比搭建整个机器学习框架所用的时间更长,数据集的宝贵一不言而喻。每天我们看似习以为常的图片验证码(选出图片中的红绿灯)其实都在把我们当成他们免费的劳动力,在帮他们给图片添加标签,哈哈。未来的世界很广阔,人工智能的世界依旧很精彩,继续加油!如果你对本篇文章感兴趣也欢迎私信或者评论区交流哦!
想要继续深入研究的小伙伴可以看这几个文章:
手把手教你使用YOLOV5训练自己的目标检测模型-口罩检测-视频教程
七、模型代码(部分)
源于博主肆十二:手把手教你使用YOLOV5训练自己的目标检测模型,链接如上
训练模型:
1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2. """ 3. Train a YOLOv5 model on a custom dataset 4. 5. Usage: 6. $ python path/to/train.py --data coco128.yaml --weights yolov5s.pt --img 640 7. """ 8. import argparse 9. import math 10. import os 11. import random 12. import sys 13. import time 14. from copy import deepcopy 15. from datetime import datetime 16. from pathlib import Path 17. 18. import numpy as np 19. import torch 20. import torch.distributed as dist 21. import torch.nn as nn 22. import yaml 23. from torch.cuda import amp 24. from torch.nn.parallel import DistributedDataParallel as DDP 25. from torch.optim import SGD, Adam, lr_scheduler 26. from tqdm import tqdm 27. 28. FILE = Path(__file__).resolve() 29. ROOT = FILE.parents[0] # YOLOv5 root directory 30. if str(ROOT) not in sys.path: 31. sys.path.append(str(ROOT)) # add ROOT to PATH 32. ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative 33. 34. import val # for end-of-epoch mAP 35. from models.experimental import attempt_load 36. from models.yolo import Model 37. from utils.autoanchor import check_anchors 38. from utils.autobatch import check_train_batch_size 39. from utils.callbacks import Callbacks 40. from utils.datasets import create_dataloader 41. from utils.downloads import attempt_download 42. from utils.general import (LOGGER, NCOLS, check_dataset, check_file, check_git_status, check_img_size, 43. check_requirements, check_suffix, check_yaml, colorstr, get_latest_run, increment_path, 44. init_seeds, intersect_dicts, labels_to_class_weights, labels_to_image_weights, methods, 45. one_cycle, print_args, print_mutation, strip_optimizer) 46. from utils.loggers import Loggers 47. from utils.loggers.wandb.wandb_utils import check_wandb_resume 48. from utils.loss import ComputeLoss 49. from utils.metrics import fitness 50. from utils.plots import plot_evolve, plot_labels 51. from utils.torch_utils import EarlyStopping, ModelEMA, de_parallel, select_device, torch_distributed_zero_first 52. 53. LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html 54. RANK = int(os.getenv('RANK', -1)) 55. WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) 56. 57. 58. def train(hyp, # path/to/hyp.yaml or hyp dictionary 59. opt, 60. device, 61. callbacks 62. ): 63. save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, = \ 64. Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \ 65. opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze 66. 67. # Directories 68. w = save_dir / 'weights' # weights dir 69. (w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir 70. last, best = w / 'last.pt', w / 'best.pt' 71. 72. # Hyperparameters 73. if isinstance(hyp, str): 74. with open(hyp, errors='ignore') as f: 75. hyp = yaml.safe_load(f) # load hyps dict 76. LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items())) 77. 78. # Save run settings 79. with open(save_dir / 'hyp.yaml', 'w') as f: 80. yaml.safe_dump(hyp, f, sort_keys=False) 81. with open(save_dir / 'opt.yaml', 'w') as f: 82. yaml.safe_dump(vars(opt), f, sort_keys=False) 83. data_dict = None 84. 85. # Loggers 86. if RANK in [-1, 0]: 87. loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance 88. if loggers.wandb: 89. data_dict = loggers.wandb.data_dict 90. if resume: 91. weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp 92. 93. # Register actions 94. for k in methods(loggers): 95. callbacks.register_action(k, callback=getattr(loggers, k)) 96. 97. # Config 98. plots = not evolve # create plots 99. cuda = device.type != 'cpu' 100. init_seeds(1 + RANK) 101. with torch_distributed_zero_first(LOCAL_RANK): 102. data_dict = data_dict or check_dataset(data) # check if None 103. train_path, val_path = data_dict['train'], data_dict['val'] 104. nc = 1 if single_cls else int(data_dict['nc']) # number of classes 105. names = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names 106. assert len(names) == nc, f'{len(names)} names found for nc={nc} dataset in {data}' # check 107. is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset 108. 109. # Model 110. check_suffix(weights, '.pt') # check weights 111. pretrained = weights.endswith('.pt') 112. if pretrained: 113. with torch_distributed_zero_first(LOCAL_RANK): 114. weights = attempt_download(weights) # download if not found locally 115. ckpt = torch.load(weights, map_location=device) # load checkpoint 116. model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create 117. exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys 118. csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 119. csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect 120. model.load_state_dict(csd, strict=False) # load 121. LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report 122. else: 123. model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create 124. 125. # Freeze 126. freeze = [f'model.{x}.' for x in range(freeze)] # layers to freeze 127. for k, v in model.named_parameters(): 128. v.requires_grad = True # train all layers 129. if any(x in k for x in freeze): 130. LOGGER.info(f'freezing {k}') 131. v.requires_grad = False 132. 133. # Image size 134. gs = max(int(model.stride.max()), 32) # grid size (max stride) 135. imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple 136. 137. # Batch size 138. if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size 139. batch_size = check_train_batch_size(model, imgsz) 140. 141. # Optimizer 142. nbs = 64 # nominal batch size 143. accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing 144. hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay 145. LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}") 146. 147. g0, g1, g2 = [], [], [] # optimizer parameter groups 148. for v in model.modules(): 149. if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias 150. g2.append(v.bias) 151. if isinstance(v, nn.BatchNorm2d): # weight (no decay) 152. g0.append(v.weight) 153. elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): # weight (with decay) 154. g1.append(v.weight) 155. 156. if opt.adam: 157. optimizer = Adam(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum 158. else: 159. optimizer = SGD(g0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) 160. 161. optimizer.add_param_group({'params': g1, 'weight_decay': hyp['weight_decay']}) # add g1 with weight_decay 162. optimizer.add_param_group({'params': g2}) # add g2 (biases) 163. LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups " 164. f"{len(g0)} weight, {len(g1)} weight (no decay), {len(g2)} bias") 165. del g0, g1, g2 166. 167. # Scheduler 168. if opt.linear_lr: 169. lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear 170. else: 171. lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf'] 172. scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs) 173. 174. # EMA 175. ema = ModelEMA(model) if RANK in [-1, 0] else None 176. 177. # Resume 178. start_epoch, best_fitness = 0, 0.0 179. if pretrained: 180. # Optimizer 181. if ckpt['optimizer'] is not None: 182. optimizer.load_state_dict(ckpt['optimizer']) 183. best_fitness = ckpt['best_fitness'] 184. 185. # EMA 186. if ema and ckpt.get('ema'): 187. ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) 188. ema.updates = ckpt['updates'] 189. 190. # Epochs 191. start_epoch = ckpt['epoch'] + 1 192. if resume: 193. assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.' 194. if epochs < start_epoch: 195. LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.") 196. epochs += ckpt['epoch'] # finetune additional epochs 197. 198. del ckpt, csd 199. 200. # DP mode 201. if cuda and RANK == -1 and torch.cuda.device_count() > 1: 202. LOGGER.warning('WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n' 203. 'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.') 204. model = torch.nn.DataParallel(model) 205. 206. # SyncBatchNorm 207. if opt.sync_bn and cuda and RANK != -1: 208. model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) 209. LOGGER.info('Using SyncBatchNorm()') 210. 211. # Trainloader 212. train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls, 213. hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, rank=LOCAL_RANK, 214. workers=workers, image_weights=opt.image_weights, quad=opt.quad, 215. prefix=colorstr('train: '), shuffle=True) 216. mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class 217. nb = len(train_loader) # number of batches 218. assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}' 219. 220. # Process 0 221. if RANK in [-1, 0]: 222. val_loader = create_dataloader(val_path, imgsz, batch_size // WORLD_SIZE * 2, gs, single_cls, 223. hyp=hyp, cache=None if noval else opt.cache, rect=True, rank=-1, 224. workers=workers, pad=0.5, 225. prefix=colorstr('val: '))[0] 226. 227. if not resume: 228. labels = np.concatenate(dataset.labels, 0) 229. # c = torch.tensor(labels[:, 0]) # classes 230. # cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency 231. # model._initialize_biases(cf.to(device)) 232. if plots: 233. plot_labels(labels, names, save_dir) 234. 235. # Anchors 236. if not opt.noautoanchor: 237. check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) 238. model.half().float() # pre-reduce anchor precision 239. 240. callbacks.run('on_pretrain_routine_end') 241. 242. # DDP mode 243. if cuda and RANK != -1: 244. model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK) 245. 246. # Model attributes 247. nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps) 248. hyp['box'] *= 3 / nl # scale to layers 249. hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers 250. hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers 251. hyp['label_smoothing'] = opt.label_smoothing 252. model.nc = nc # attach number of classes to model 253. model.hyp = hyp # attach hyperparameters to model 254. model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights 255. model.names = names 256. 257. # Start training 258. t0 = time.time() 259. nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations) 260. # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training 261. last_opt_step = -1 262. maps = np.zeros(nc) # mAP per class 263. results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) 264. scheduler.last_epoch = start_epoch - 1 # do not move 265. scaler = amp.GradScaler(enabled=cuda) 266. stopper = EarlyStopping(patience=opt.patience) 267. compute_loss = ComputeLoss(model) # init loss class 268. LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n' 269. f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n' 270. f"Logging results to {colorstr('bold', save_dir)}\n" 271. f'Starting training for {epochs} epochs...') 272. for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ 273. model.train() 274. 275. # Update image weights (optional, single-GPU only) 276. if opt.image_weights: 277. cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights 278. iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights 279. dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx 280. 281. # Update mosaic border (optional) 282. # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) 283. # dataset.mosaic_border = [b - imgsz, -b] # height, width borders 284. 285. mloss = torch.zeros(3, device=device) # mean losses 286. if RANK != -1: 287. train_loader.sampler.set_epoch(epoch) 288. pbar = enumerate(train_loader) 289. LOGGER.info(('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size')) 290. if RANK in [-1, 0]: 291. pbar = tqdm(pbar, total=nb, ncols=NCOLS, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar 292. optimizer.zero_grad() 293. for i, (imgs, targets, paths, _) in pbar: # batch ------------------------------------------------------------- 294. ni = i + nb * epoch # number integrated batches (since train start) 295. imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0 296. 297. # Warmup 298. if ni <= nw: 299. xi = [0, nw] # x interp 300. # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou) 301. accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round()) 302. for j, x in enumerate(optimizer.param_groups): 303. # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0 304. x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)]) 305. if 'momentum' in x: 306. x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']]) 307. 308. # Multi-scale 309. if opt.multi_scale: 310. sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size 311. sf = sz / max(imgs.shape[2:]) # scale factor 312. if sf != 1: 313. ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) 314. imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False) 315. 316. # Forward 317. with amp.autocast(enabled=cuda): 318. pred = model(imgs) # forward 319. loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size 320. if RANK != -1: 321. loss *= WORLD_SIZE # gradient averaged between devices in DDP mode 322. if opt.quad: 323. loss *= 4. 324. 325. # Backward 326. scaler.scale(loss).backward() 327. 328. # Optimize 329. if ni - last_opt_step >= accumulate: 330. scaler.step(optimizer) # optimizer.step 331. scaler.update() 332. optimizer.zero_grad() 333. if ema: 334. ema.update(model) 335. last_opt_step = ni 336. 337. # Log 338. if RANK in [-1, 0]: 339. mloss = (mloss * i + loss_items) / (i + 1) # update mean losses 340. mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB) 341. pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % ( 342. f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1])) 343. callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn) 344. # end batch ------------------------------------------------------------------------------------------------ 345. 346. # Scheduler 347. lr = [x['lr'] for x in optimizer.param_groups] # for loggers 348. scheduler.step() 349. 350. if RANK in [-1, 0]: 351. # mAP 352. callbacks.run('on_train_epoch_end', epoch=epoch) 353. ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights']) 354. final_epoch = (epoch + 1 == epochs) or stopper.possible_stop 355. if not noval or final_epoch: # Calculate mAP 356. results, maps, _ = val.run(data_dict, 357. batch_size=batch_size // WORLD_SIZE * 2, 358. imgsz=imgsz, 359. model=ema.ema, 360. single_cls=single_cls, 361. dataloader=val_loader, 362. save_dir=save_dir, 363. plots=False, 364. callbacks=callbacks, 365. compute_loss=compute_loss) 366. 367. # Update best mAP 368. fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] 369. if fi > best_fitness: 370. best_fitness = fi 371. log_vals = list(mloss) + list(results) + lr 372. callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi) 373. 374. # Save model 375. if (not nosave) or (final_epoch and not evolve): # if save 376. ckpt = {'epoch': epoch, 377. 'best_fitness': best_fitness, 378. 'model': deepcopy(de_parallel(model)).half(), 379. 'ema': deepcopy(ema.ema).half(), 380. 'updates': ema.updates, 381. 'optimizer': optimizer.state_dict(), 382. 'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None, 383. 'date': datetime.now().isoformat()} 384. 385. # Save last, best and delete 386. torch.save(ckpt, last) 387. if best_fitness == fi: 388. torch.save(ckpt, best) 389. if (epoch > 0) and (opt.save_period > 0) and (epoch % opt.save_period == 0): 390. torch.save(ckpt, w / f'epoch{epoch}.pt') 391. del ckpt 392. callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi) 393. 394. # Stop Single-GPU 395. if RANK == -1 and stopper(epoch=epoch, fitness=fi): 396. break 397. 398. # Stop DDP TODO: known issues shttps://github.com/ultralytics/yolov5/pull/4576 399. # stop = stopper(epoch=epoch, fitness=fi) 400. # if RANK == 0: 401. # dist.broadcast_object_list([stop], 0) # broadcast 'stop' to all ranks 402. 403. # Stop DPP 404. # with torch_distributed_zero_first(RANK): 405. # if stop: 406. # break # must break all DDP ranks 407. 408. # end epoch ---------------------------------------------------------------------------------------------------- 409. # end training ----------------------------------------------------------------------------------------------------- 410. if RANK in [-1, 0]: 411. LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.') 412. for f in last, best: 413. if f.exists(): 414. strip_optimizer(f) # strip optimizers 415. if f is best: 416. LOGGER.info(f'\nValidating {f}...') 417. results, _, _ = val.run(data_dict, 418. batch_size=batch_size // WORLD_SIZE * 2, 419. imgsz=imgsz, 420. model=attempt_load(f, device).half(), 421. iou_thres=0.65 if is_coco else 0.60, # best pycocotools results at 0.65 422. single_cls=single_cls, 423. dataloader=val_loader, 424. save_dir=save_dir, 425. save_json=is_coco, 426. verbose=True, 427. plots=True, 428. callbacks=callbacks, 429. compute_loss=compute_loss) # val best model with plots 430. if is_coco: 431. callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi) 432. 433. callbacks.run('on_train_end', last, best, plots, epoch, results) 434. LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}") 435. 436. torch.cuda.empty_cache() 437. return results 438. 439. 440. # 明天把这些模型都试试效果先,一波波给他训练完毕,找个公开的数据集测试一下。 441. def parse_opt(known=False): 442. parser = argparse.ArgumentParser() 443. parser.add_argument('--weights', type=str, default=ROOT / 'pretrained/yolov5s.pt', help='initial weights path') 444. parser.add_argument('--cfg', type=str, default=ROOT / 'models/yolov5s.yaml', help='model.yaml path') 445. parser.add_argument('--data', type=str, default=ROOT / 'data/data.yaml', help='dataset.yaml path') 446. parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch.yaml', help='hyperparameters path') 447. parser.add_argument('--epochs', type=int, default=300) 448. parser.add_argument('--batch-size', type=int, default=4, help='total batch size for all GPUs, -1 for autobatch') 449. parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)') 450. parser.add_argument('--rect', action='store_true', help='rectangular training') 451. parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') 452. parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') 453. parser.add_argument('--noval', action='store_true', help='only validate final epoch') 454. parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check') 455. parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations') 456. parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') 457. parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"') 458. parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') 459. parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 460. # parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') 461. parser.add_argument('--multi-scale', default=True, help='vary img-size +/- 50%%') 462. parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class') 463. parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') 464. parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') 465. parser.add_argument('--workers', type=int, default=0, help='max dataloader workers (per RANK in DDP mode)') 466. parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name') 467. parser.add_argument('--name', default='exp', help='save to project/name') 468. parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') 469. parser.add_argument('--quad', action='store_true', help='quad dataloader') 470. parser.add_argument('--linear-lr', action='store_true', help='linear LR') 471. parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon') 472. parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)') 473. parser.add_argument('--freeze', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24') 474. parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)') 475. parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify') 476. # Weights & Biases arguments 477. parser.add_argument('--entity', default=None, help='W&B: Entity') 478. parser.add_argument('--upload_dataset', action='store_true', help='W&B: Upload dataset as artifact table') 479. parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval') 480. parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use') 481. 482. opt = parser.parse_known_args()[0] if known else parser.parse_args() 483. return opt 484. 485. 486. def main(opt, callbacks=Callbacks()): 487. 488. # Checks 489. if RANK in [-1, 0]: 490. print_args(FILE.stem, opt) 491. check_git_status() 492. check_requirements(exclude=['thop']) 493. 494. # Resume 495. if opt.resume and not check_wandb_resume(opt) and not opt.evolve: # resume an interrupted run 496. ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path 497. assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist' 498. with open(Path(ckpt).parent.parent / 'opt.yaml', errors='ignore') as f: 499. opt = argparse.Namespace(**yaml.safe_load(f)) # replace 500. opt.cfg, opt.weights, opt.resume = '', ckpt, True # reinstate 501. LOGGER.info(f'Resuming training from {ckpt}') 502. else: 503. opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \ 504. check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks 505. assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified' 506. if opt.evolve: 507. opt.project = str(ROOT / 'runs/evolve') 508. opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume 509. opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) 510. 511. # DDP mode 512. device = select_device(opt.device, batch_size=opt.batch_size) 513. if LOCAL_RANK != -1: 514. assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command' 515. assert opt.batch_size % WORLD_SIZE == 0, '--batch-size must be multiple of CUDA device count' 516. assert not opt.image_weights, '--image-weights argument is not compatible with DDP training' 517. assert not opt.evolve, '--evolve argument is not compatible with DDP training' 518. torch.cuda.set_device(LOCAL_RANK) 519. device = torch.device('cuda', LOCAL_RANK) 520. dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo") 521. 522. # Train 523. if not opt.evolve: 524. train(opt.hyp, opt, device, callbacks) 525. if WORLD_SIZE > 1 and RANK == 0: 526. LOGGER.info('Destroying process group... ') 527. dist.destroy_process_group() 528. 529. # Evolve hyperparameters (optional) 530. else: 531. # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit) 532. meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) 533. 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) 534. 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 535. 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay 536. 'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok) 537. 'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum 538. 'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr 539. 'box': (1, 0.02, 0.2), # box loss gain 540. 'cls': (1, 0.2, 4.0), # cls loss gain 541. 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight 542. 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels) 543. 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight 544. 'iou_t': (0, 0.1, 0.7), # IoU training threshold 545. 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold 546. 'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore) 547. 'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) 548. 'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction) 549. 'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) 550. 'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) 551. 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg) 552. 'translate': (1, 0.0, 0.9), # image translation (+/- fraction) 553. 'scale': (1, 0.0, 0.9), # image scale (+/- gain) 554. 'shear': (1, 0.0, 10.0), # image shear (+/- deg) 555. 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 556. 'flipud': (1, 0.0, 1.0), # image flip up-down (probability) 557. 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability) 558. 'mosaic': (1, 0.0, 1.0), # image mixup (probability) 559. 'mixup': (1, 0.0, 1.0), # image mixup (probability) 560. 'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability) 561. 562. with open(opt.hyp, errors='ignore') as f: 563. hyp = yaml.safe_load(f) # load hyps dict 564. if 'anchors' not in hyp: # anchors commented in hyp.yaml 565. hyp['anchors'] = 3 566. opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch 567. # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices 568. evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv' 569. if opt.bucket: 570. os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {save_dir}') # download evolve.csv if exists 571. 572. for _ in range(opt.evolve): # generations to evolve 573. if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate 574. # Select parent(s) 575. parent = 'single' # parent selection method: 'single' or 'weighted' 576. x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1) 577. n = min(5, len(x)) # number of previous results to consider 578. x = x[np.argsort(-fitness(x))][:n] # top n mutations 579. w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0) 580. if parent == 'single' or len(x) == 1: 581. # x = x[random.randint(0, n - 1)] # random selection 582. x = x[random.choices(range(n), weights=w)[0]] # weighted selection 583. elif parent == 'weighted': 584. x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination 585. 586. # Mutate 587. mp, s = 0.8, 0.2 # mutation probability, sigma 588. npr = np.random 589. npr.seed(int(time.time())) 590. g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1 591. ng = len(meta) 592. v = np.ones(ng) 593. while all(v == 1): # mutate until a change occurs (prevent duplicates) 594. v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0) 595. for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300) 596. hyp[k] = float(x[i + 7] * v[i]) # mutate 597. 598. # Constrain to limits 599. for k, v in meta.items(): 600. hyp[k] = max(hyp[k], v[1]) # lower limit 601. hyp[k] = min(hyp[k], v[2]) # upper limit 602. hyp[k] = round(hyp[k], 5) # significant digits 603. 604. # Train mutation 605. results = train(hyp.copy(), opt, device, callbacks) 606. 607. # Write mutation results 608. print_mutation(results, hyp.copy(), save_dir, opt.bucket) 609. 610. # Plot results 611. plot_evolve(evolve_csv) 612. LOGGER.info(f'Hyperparameter evolution finished\n' 613. f"Results saved to {colorstr('bold', save_dir)}\n" 614. f'Use best hyperparameters example: $ python train.py --hyp {evolve_yaml}') 615. 616. 617. def run(**kwargs): 618. # Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt') 619. opt = parse_opt(True) 620. for k, v in kwargs.items(): 621. setattr(opt, k, v) 622. main(opt) 623. 624. 625. # python train.py --data mask_data.yaml --cfg mask_yolov5s.yaml --weights pretrained/yolov5s.pt --epoch 100 --batch-size 4 --device cpu 626. # python train.py --data mask_data.yaml --cfg mask_yolov5l.yaml --weights pretrained/yolov5l.pt --epoch 100 --batch-size 4 627. # python train.py --data mask_data.yaml --cfg mask_yolov5m.yaml --weights pretrained/yolov5m.pt --epoch 100 --batch-size 4 628. if __name__ == "__main__": 629. opt = parse_opt() 630. main(opt)