train.py 37 KB

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