您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

671 行
34KB

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