You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

573 lines
29KB

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