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.

626 lines
32KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Train a YOLOv5 model on a custom dataset
  4. Usage:
  5. $ python path/to/train.py --data coco128.yaml --weights yolov5s.pt --img 640
  6. """
  7. import argparse
  8. import logging
  9. import math
  10. import os
  11. import random
  12. import sys
  13. import time
  14. from copy import deepcopy
  15. from pathlib import Path
  16. import numpy as np
  17. import torch
  18. import torch.distributed as dist
  19. import torch.nn as nn
  20. import yaml
  21. from torch.cuda import amp
  22. from torch.nn.parallel import DistributedDataParallel as DDP
  23. from torch.optim import Adam, SGD, lr_scheduler
  24. from tqdm import tqdm
  25. FILE = Path(__file__).resolve()
  26. ROOT = FILE.parents[0] # YOLOv5 root directory
  27. if str(ROOT) not in sys.path:
  28. sys.path.append(str(ROOT)) # add ROOT to PATH
  29. ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
  30. import val # for end-of-epoch mAP
  31. from models.experimental import attempt_load
  32. from models.yolo import Model
  33. from utils.autoanchor import check_anchors
  34. from utils.autobatch import check_train_batch_size
  35. from utils.datasets import create_dataloader
  36. from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \
  37. strip_optimizer, get_latest_run, check_dataset, check_git_status, check_img_size, check_requirements, \
  38. check_file, check_yaml, check_suffix, print_args, print_mutation, one_cycle, colorstr, methods, LOGGER
  39. from utils.downloads import attempt_download
  40. from utils.loss import ComputeLoss
  41. from utils.plots import plot_labels, plot_evolve
  42. from utils.torch_utils import EarlyStopping, ModelEMA, de_parallel, intersect_dicts, select_device, \
  43. torch_distributed_zero_first
  44. from utils.loggers.wandb.wandb_utils import check_wandb_resume
  45. from utils.metrics import fitness
  46. from utils.loggers import Loggers
  47. from utils.callbacks import Callbacks
  48. LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
  49. RANK = int(os.getenv('RANK', -1))
  50. WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
  51. def train(hyp, # path/to/hyp.yaml or hyp dictionary
  52. opt,
  53. device,
  54. callbacks
  55. ):
  56. save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, = \
  57. Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
  58. opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
  59. # Directories
  60. w = save_dir / 'weights' # weights dir
  61. (w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
  62. last, best = w / 'last.pt', w / 'best.pt'
  63. # Hyperparameters
  64. if isinstance(hyp, str):
  65. with open(hyp, errors='ignore') as f:
  66. hyp = yaml.safe_load(f) # load hyps dict
  67. LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
  68. # Save run settings
  69. with open(save_dir / 'hyp.yaml', 'w') as f:
  70. yaml.safe_dump(hyp, f, sort_keys=False)
  71. with open(save_dir / 'opt.yaml', 'w') as f:
  72. yaml.safe_dump(vars(opt), f, sort_keys=False)
  73. data_dict = None
  74. # Loggers
  75. if RANK in [-1, 0]:
  76. loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
  77. if loggers.wandb:
  78. data_dict = loggers.wandb.data_dict
  79. if resume:
  80. weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp
  81. # Register actions
  82. for k in methods(loggers):
  83. callbacks.register_action(k, callback=getattr(loggers, k))
  84. # Config
  85. plots = not evolve # create plots
  86. cuda = device.type != 'cpu'
  87. init_seeds(1 + RANK)
  88. with torch_distributed_zero_first(LOCAL_RANK):
  89. data_dict = data_dict or check_dataset(data) # check if None
  90. train_path, val_path = data_dict['train'], data_dict['val']
  91. nc = 1 if single_cls else int(data_dict['nc']) # number of classes
  92. names = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
  93. assert len(names) == nc, f'{len(names)} names found for nc={nc} dataset in {data}' # check
  94. is_coco = data.endswith('coco.yaml') and nc == 80 # COCO dataset
  95. # Model
  96. check_suffix(weights, '.pt') # check weights
  97. pretrained = weights.endswith('.pt')
  98. if pretrained:
  99. with torch_distributed_zero_first(LOCAL_RANK):
  100. weights = attempt_download(weights) # download if not found locally
  101. ckpt = torch.load(weights, map_location=device) # load checkpoint
  102. model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
  103. exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
  104. csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
  105. csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect
  106. model.load_state_dict(csd, strict=False) # load
  107. LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report
  108. else:
  109. model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
  110. # Freeze
  111. freeze = [f'model.{x}.' for x in range(freeze)] # layers to freeze
  112. for k, v in model.named_parameters():
  113. v.requires_grad = True # train all layers
  114. if any(x in k for x in freeze):
  115. LOGGER.info(f'freezing {k}')
  116. v.requires_grad = False
  117. # Image size
  118. gs = max(int(model.stride.max()), 32) # grid size (max stride)
  119. imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
  120. # Batch size
  121. if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size
  122. batch_size = check_train_batch_size(model, imgsz)
  123. # Optimizer
  124. nbs = 64 # nominal batch size
  125. accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
  126. hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
  127. LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
  128. g0, g1, g2 = [], [], [] # optimizer parameter groups
  129. for v in model.modules():
  130. if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias
  131. g2.append(v.bias)
  132. if isinstance(v, nn.BatchNorm2d): # weight (no decay)
  133. g0.append(v.weight)
  134. elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): # weight (with decay)
  135. g1.append(v.weight)
  136. if opt.adam:
  137. optimizer = Adam(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
  138. else:
  139. optimizer = SGD(g0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
  140. optimizer.add_param_group({'params': g1, 'weight_decay': hyp['weight_decay']}) # add g1 with weight_decay
  141. optimizer.add_param_group({'params': g2}) # add g2 (biases)
  142. LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups "
  143. f"{len(g0)} weight, {len(g1)} weight (no decay), {len(g2)} bias")
  144. del g0, g1, g2
  145. # Scheduler
  146. if opt.linear_lr:
  147. lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
  148. else:
  149. lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
  150. scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
  151. # EMA
  152. ema = ModelEMA(model) if RANK in [-1, 0] else None
  153. # Resume
  154. start_epoch, best_fitness = 0, 0.0
  155. if pretrained:
  156. # Optimizer
  157. if ckpt['optimizer'] is not None:
  158. optimizer.load_state_dict(ckpt['optimizer'])
  159. best_fitness = ckpt['best_fitness']
  160. # EMA
  161. if ema and ckpt.get('ema'):
  162. ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
  163. ema.updates = ckpt['updates']
  164. # Epochs
  165. start_epoch = ckpt['epoch'] + 1
  166. if resume:
  167. assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.'
  168. if epochs < start_epoch:
  169. LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.")
  170. epochs += ckpt['epoch'] # finetune additional epochs
  171. del ckpt, csd
  172. # DP mode
  173. if cuda and RANK == -1 and torch.cuda.device_count() > 1:
  174. logging.warning('DP not recommended, instead use torch.distributed.run for best DDP Multi-GPU results.\n'
  175. 'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')
  176. model = torch.nn.DataParallel(model)
  177. # SyncBatchNorm
  178. if opt.sync_bn and cuda and RANK != -1:
  179. model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
  180. LOGGER.info('Using SyncBatchNorm()')
  181. # Trainloader
  182. train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls,
  183. hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, rank=LOCAL_RANK,
  184. workers=workers, image_weights=opt.image_weights, quad=opt.quad,
  185. prefix=colorstr('train: '))
  186. mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class
  187. nb = len(train_loader) # number of batches
  188. assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
  189. # Process 0
  190. if RANK in [-1, 0]:
  191. val_loader = create_dataloader(val_path, imgsz, batch_size // WORLD_SIZE * 2, gs, single_cls,
  192. hyp=hyp, cache=None if noval else opt.cache, rect=True, rank=-1,
  193. workers=workers, pad=0.5,
  194. prefix=colorstr('val: '))[0]
  195. if not resume:
  196. labels = np.concatenate(dataset.labels, 0)
  197. # c = torch.tensor(labels[:, 0]) # classes
  198. # cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
  199. # model._initialize_biases(cf.to(device))
  200. if plots:
  201. plot_labels(labels, names, save_dir)
  202. # Anchors
  203. if not opt.noautoanchor:
  204. check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
  205. model.half().float() # pre-reduce anchor precision
  206. callbacks.run('on_pretrain_routine_end')
  207. # DDP mode
  208. if cuda and RANK != -1:
  209. model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
  210. # Model parameters
  211. nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps)
  212. hyp['box'] *= 3. / nl # scale to layers
  213. hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
  214. hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl # scale to image size and layers
  215. hyp['label_smoothing'] = opt.label_smoothing
  216. model.nc = nc # attach number of classes to model
  217. model.hyp = hyp # attach hyperparameters to model
  218. model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
  219. model.names = names
  220. # Start training
  221. t0 = time.time()
  222. nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations)
  223. # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
  224. last_opt_step = -1
  225. maps = np.zeros(nc) # mAP per class
  226. results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
  227. scheduler.last_epoch = start_epoch - 1 # do not move
  228. scaler = amp.GradScaler(enabled=cuda)
  229. stopper = EarlyStopping(patience=opt.patience)
  230. compute_loss = ComputeLoss(model) # init loss class
  231. LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
  232. f'Using {train_loader.num_workers} dataloader workers\n'
  233. f"Logging results to {colorstr('bold', save_dir)}\n"
  234. f'Starting training for {epochs} epochs...')
  235. for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
  236. model.train()
  237. # Update image weights (optional, single-GPU only)
  238. if opt.image_weights:
  239. cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
  240. iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
  241. dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
  242. # Update mosaic border (optional)
  243. # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
  244. # dataset.mosaic_border = [b - imgsz, -b] # height, width borders
  245. mloss = torch.zeros(3, device=device) # mean losses
  246. if RANK != -1:
  247. train_loader.sampler.set_epoch(epoch)
  248. pbar = enumerate(train_loader)
  249. LOGGER.info(('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size'))
  250. if RANK in [-1, 0]:
  251. pbar = tqdm(pbar, total=nb) # progress bar
  252. optimizer.zero_grad()
  253. for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
  254. ni = i + nb * epoch # number integrated batches (since train start)
  255. imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
  256. # Warmup
  257. if ni <= nw:
  258. xi = [0, nw] # x interp
  259. # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
  260. accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
  261. for j, x in enumerate(optimizer.param_groups):
  262. # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
  263. x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
  264. if 'momentum' in x:
  265. x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
  266. # Multi-scale
  267. if opt.multi_scale:
  268. sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
  269. sf = sz / max(imgs.shape[2:]) # scale factor
  270. if sf != 1:
  271. ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
  272. imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
  273. # Forward
  274. with amp.autocast(enabled=cuda):
  275. pred = model(imgs) # forward
  276. loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
  277. if RANK != -1:
  278. loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
  279. if opt.quad:
  280. loss *= 4.
  281. # Backward
  282. scaler.scale(loss).backward()
  283. # Optimize
  284. if ni - last_opt_step >= accumulate:
  285. scaler.step(optimizer) # optimizer.step
  286. scaler.update()
  287. optimizer.zero_grad()
  288. if ema:
  289. ema.update(model)
  290. last_opt_step = ni
  291. # Log
  292. if RANK in [-1, 0]:
  293. mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
  294. mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
  295. pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % (
  296. f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
  297. callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn)
  298. # end batch ------------------------------------------------------------------------------------------------
  299. # Scheduler
  300. lr = [x['lr'] for x in optimizer.param_groups] # for loggers
  301. scheduler.step()
  302. if RANK in [-1, 0]:
  303. # mAP
  304. callbacks.run('on_train_epoch_end', epoch=epoch)
  305. ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
  306. final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
  307. if not noval or final_epoch: # Calculate mAP
  308. results, maps, _ = val.run(data_dict,
  309. batch_size=batch_size // WORLD_SIZE * 2,
  310. imgsz=imgsz,
  311. model=ema.ema,
  312. single_cls=single_cls,
  313. dataloader=val_loader,
  314. save_dir=save_dir,
  315. plots=False,
  316. callbacks=callbacks,
  317. compute_loss=compute_loss)
  318. # Update best mAP
  319. fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
  320. if fi > best_fitness:
  321. best_fitness = fi
  322. log_vals = list(mloss) + list(results) + lr
  323. callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
  324. # Save model
  325. if (not nosave) or (final_epoch and not evolve): # if save
  326. ckpt = {'epoch': epoch,
  327. 'best_fitness': best_fitness,
  328. 'model': deepcopy(de_parallel(model)).half(),
  329. 'ema': deepcopy(ema.ema).half(),
  330. 'updates': ema.updates,
  331. 'optimizer': optimizer.state_dict(),
  332. 'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None}
  333. # Save last, best and delete
  334. torch.save(ckpt, last)
  335. if best_fitness == fi:
  336. torch.save(ckpt, best)
  337. if (epoch > 0) and (opt.save_period > 0) and (epoch % opt.save_period == 0):
  338. torch.save(ckpt, w / f'epoch{epoch}.pt')
  339. del ckpt
  340. callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
  341. # Stop Single-GPU
  342. if RANK == -1 and stopper(epoch=epoch, fitness=fi):
  343. break
  344. # Stop DDP TODO: known issues shttps://github.com/ultralytics/yolov5/pull/4576
  345. # stop = stopper(epoch=epoch, fitness=fi)
  346. # if RANK == 0:
  347. # dist.broadcast_object_list([stop], 0) # broadcast 'stop' to all ranks
  348. # Stop DPP
  349. # with torch_distributed_zero_first(RANK):
  350. # if stop:
  351. # break # must break all DDP ranks
  352. # end epoch ----------------------------------------------------------------------------------------------------
  353. # end training -----------------------------------------------------------------------------------------------------
  354. if RANK in [-1, 0]:
  355. LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
  356. for f in last, best:
  357. if f.exists():
  358. strip_optimizer(f) # strip optimizers
  359. if f is best:
  360. LOGGER.info(f'\nValidating {f}...')
  361. results, _, _ = val.run(data_dict,
  362. batch_size=batch_size // WORLD_SIZE * 2,
  363. imgsz=imgsz,
  364. model=attempt_load(f, device).half(),
  365. iou_thres=0.65 if is_coco else 0.60, # best pycocotools results at 0.65
  366. single_cls=single_cls,
  367. dataloader=val_loader,
  368. save_dir=save_dir,
  369. save_json=is_coco,
  370. verbose=True,
  371. plots=True,
  372. callbacks=callbacks,
  373. compute_loss=compute_loss) # val best model with plots
  374. if is_coco:
  375. callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
  376. callbacks.run('on_train_end', last, best, plots, epoch, results)
  377. LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
  378. torch.cuda.empty_cache()
  379. return results
  380. def parse_opt(known=False):
  381. parser = argparse.ArgumentParser()
  382. parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
  383. parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
  384. parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
  385. parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch.yaml', help='hyperparameters path')
  386. parser.add_argument('--epochs', type=int, default=300)
  387. parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
  388. parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
  389. parser.add_argument('--rect', action='store_true', help='rectangular training')
  390. parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
  391. parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
  392. parser.add_argument('--noval', action='store_true', help='only validate final epoch')
  393. parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
  394. parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
  395. parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
  396. parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
  397. parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
  398. parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  399. parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
  400. parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
  401. parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
  402. parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
  403. parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
  404. parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
  405. parser.add_argument('--name', default='exp', help='save to project/name')
  406. parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  407. parser.add_argument('--quad', action='store_true', help='quad dataloader')
  408. parser.add_argument('--linear-lr', action='store_true', help='linear LR')
  409. parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
  410. parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
  411. parser.add_argument('--freeze', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24')
  412. parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
  413. parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
  414. # Weights & Biases arguments
  415. parser.add_argument('--entity', default=None, help='W&B: Entity')
  416. parser.add_argument('--upload_dataset', action='store_true', help='W&B: Upload dataset as artifact table')
  417. parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
  418. parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')
  419. opt = parser.parse_known_args()[0] if known else parser.parse_args()
  420. return opt
  421. def main(opt, callbacks=Callbacks()):
  422. # Checks
  423. if RANK in [-1, 0]:
  424. print_args(FILE.stem, opt)
  425. check_git_status()
  426. check_requirements(exclude=['thop'])
  427. # Resume
  428. if opt.resume and not check_wandb_resume(opt) and not opt.evolve: # resume an interrupted run
  429. ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path
  430. assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
  431. with open(Path(ckpt).parent.parent / 'opt.yaml', errors='ignore') as f:
  432. opt = argparse.Namespace(**yaml.safe_load(f)) # replace
  433. opt.cfg, opt.weights, opt.resume = '', ckpt, True # reinstate
  434. LOGGER.info(f'Resuming training from {ckpt}')
  435. else:
  436. opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
  437. check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks
  438. assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
  439. if opt.evolve:
  440. opt.project = str(ROOT / 'runs/evolve')
  441. opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
  442. opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
  443. # DDP mode
  444. device = select_device(opt.device, batch_size=opt.batch_size)
  445. if LOCAL_RANK != -1:
  446. assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
  447. assert opt.batch_size % WORLD_SIZE == 0, '--batch-size must be multiple of CUDA device count'
  448. assert not opt.image_weights, '--image-weights argument is not compatible with DDP training'
  449. assert not opt.evolve, '--evolve argument is not compatible with DDP training'
  450. torch.cuda.set_device(LOCAL_RANK)
  451. device = torch.device('cuda', LOCAL_RANK)
  452. dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
  453. # Train
  454. if not opt.evolve:
  455. train(opt.hyp, opt, device, callbacks)
  456. if WORLD_SIZE > 1 and RANK == 0:
  457. LOGGER.info('Destroying process group... ')
  458. dist.destroy_process_group()
  459. # Evolve hyperparameters (optional)
  460. else:
  461. # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
  462. meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
  463. 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
  464. 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
  465. 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
  466. 'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
  467. 'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
  468. 'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
  469. 'box': (1, 0.02, 0.2), # box loss gain
  470. 'cls': (1, 0.2, 4.0), # cls loss gain
  471. 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
  472. 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
  473. 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
  474. 'iou_t': (0, 0.1, 0.7), # IoU training threshold
  475. 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
  476. 'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
  477. 'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
  478. 'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
  479. 'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
  480. 'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
  481. 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
  482. 'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
  483. 'scale': (1, 0.0, 0.9), # image scale (+/- gain)
  484. 'shear': (1, 0.0, 10.0), # image shear (+/- deg)
  485. 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
  486. 'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
  487. 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
  488. 'mosaic': (1, 0.0, 1.0), # image mixup (probability)
  489. 'mixup': (1, 0.0, 1.0), # image mixup (probability)
  490. 'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)
  491. with open(opt.hyp, errors='ignore') as f:
  492. hyp = yaml.safe_load(f) # load hyps dict
  493. if 'anchors' not in hyp: # anchors commented in hyp.yaml
  494. hyp['anchors'] = 3
  495. opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch
  496. # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
  497. evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
  498. if opt.bucket:
  499. os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {save_dir}') # download evolve.csv if exists
  500. for _ in range(opt.evolve): # generations to evolve
  501. if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate
  502. # Select parent(s)
  503. parent = 'single' # parent selection method: 'single' or 'weighted'
  504. x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
  505. n = min(5, len(x)) # number of previous results to consider
  506. x = x[np.argsort(-fitness(x))][:n] # top n mutations
  507. w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)
  508. if parent == 'single' or len(x) == 1:
  509. # x = x[random.randint(0, n - 1)] # random selection
  510. x = x[random.choices(range(n), weights=w)[0]] # weighted selection
  511. elif parent == 'weighted':
  512. x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination
  513. # Mutate
  514. mp, s = 0.8, 0.2 # mutation probability, sigma
  515. npr = np.random
  516. npr.seed(int(time.time()))
  517. g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1
  518. ng = len(meta)
  519. v = np.ones(ng)
  520. while all(v == 1): # mutate until a change occurs (prevent duplicates)
  521. v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
  522. for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
  523. hyp[k] = float(x[i + 7] * v[i]) # mutate
  524. # Constrain to limits
  525. for k, v in meta.items():
  526. hyp[k] = max(hyp[k], v[1]) # lower limit
  527. hyp[k] = min(hyp[k], v[2]) # upper limit
  528. hyp[k] = round(hyp[k], 5) # significant digits
  529. # Train mutation
  530. results = train(hyp.copy(), opt, device, callbacks)
  531. # Write mutation results
  532. print_mutation(results, hyp.copy(), save_dir, opt.bucket)
  533. # Plot results
  534. plot_evolve(evolve_csv)
  535. LOGGER.info(f'Hyperparameter evolution finished\n'
  536. f"Results saved to {colorstr('bold', save_dir)}\n"
  537. f'Use best hyperparameters example: $ python train.py --hyp {evolve_yaml}')
  538. def run(**kwargs):
  539. # Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
  540. opt = parse_opt(True)
  541. for k, v in kwargs.items():
  542. setattr(opt, k, v)
  543. main(opt)
  544. if __name__ == "__main__":
  545. opt = parse_opt()
  546. main(opt)