Bai360/Cotton2
0
1# YOLOv5 🚀 by Ultralytics, GPL-3.0 license2"""3Train a YOLOv5 model on a custom dataset.4Models and datasets download automatically from the latest YOLOv5 release.5 6Usage - Single-GPU training:7 $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended)8 $ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch9 10Usage - Multi-GPU DDP training:11 $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights yolov5s.pt --img 640 --device 0,1,2,312 13Models: https://github.com/ultralytics/yolov5/tree/master/models14Datasets: https://github.com/ultralytics/yolov5/tree/master/data15Tutorial: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data16"""17 18import argparse19import math20import os21import random22import sys23import time24from copy import deepcopy25from datetime import datetime26from pathlib import Path27 28import numpy as np29import torch30import torch.distributed as dist31import torch.nn as nn32import yaml33from torch.optim import lr_scheduler34from tqdm import tqdm35 36FILE = Path(__file__).resolve()37ROOT = FILE.parents[0] # YOLOv5 root directory38if str(ROOT) not in sys.path:39 sys.path.append(str(ROOT)) # add ROOT to PATH40ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative41 42import val as validate # for end-of-epoch mAP43from models.experimental import attempt_load44from models.yolo import Model45from utils.autoanchor import check_anchors46from utils.autobatch import check_train_batch_size47from utils.callbacks import Callbacks48from utils.dataloaders import create_dataloader49from utils.downloads import attempt_download, is_url50from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_info,51 check_git_status, check_img_size, check_requirements, check_suffix, check_yaml, colorstr,52 get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights,53 labels_to_image_weights, methods, one_cycle, print_args, print_mutation, strip_optimizer,54 yaml_save)55from utils.loggers import Loggers56from utils.loggers.comet.comet_utils import check_comet_resume57from utils.loss import ComputeLoss58from utils.metrics import fitness59from utils.plots import plot_evolve60from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,61 smart_resume, torch_distributed_zero_first)62 63LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html64RANK = int(os.getenv('RANK', -1))65WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))66GIT_INFO = check_git_info()67 68 69def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary70 save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \71 Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \72 opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze73 callbacks.run('on_pretrain_routine_start')74 75 # Directories76 w = save_dir / 'weights' # weights dir77 (w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir78 last, best = w / 'last.pt', w / 'best.pt'79 80 # Hyperparameters81 if isinstance(hyp, str):82 with open(hyp, errors='ignore') as f:83 hyp = yaml.safe_load(f) # load hyps dict84 LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))85 opt.hyp = hyp.copy() # for saving hyps to checkpoints86 87 # Save run settings88 if not evolve:89 yaml_save(save_dir / 'hyp.yaml', hyp)90 yaml_save(save_dir / 'opt.yaml', vars(opt))91 92 # Loggers93 data_dict = None94 if RANK in {-1, 0}:95 loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance96 97 # Register actions98 for k in methods(loggers):99 callbacks.register_action(k, callback=getattr(loggers, k))100 101 # Process custom dataset artifact link102 data_dict = loggers.remote_dataset103 if resume: # If resuming runs from remote artifact104 weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size105 106 # Config107 plots = not evolve and not opt.noplots # create plots108 cuda = device.type != 'cpu'109 init_seeds(opt.seed + 1 + RANK, deterministic=True)110 with torch_distributed_zero_first(LOCAL_RANK):111 data_dict = data_dict or check_dataset(data) # check if None112 train_path, val_path = data_dict['train'], data_dict['val']113 nc = 1 if single_cls else int(data_dict['nc']) # number of classes114 names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names115 is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset116 117 # Model118 check_suffix(weights, '.pt') # check weights119 pretrained = weights.endswith('.pt')120 if pretrained:121 with torch_distributed_zero_first(LOCAL_RANK):122 weights = attempt_download(weights) # download if not found locally123 ckpt = torch.load(weights, map_location='cpu') # load checkpoint to CPU to avoid CUDA memory leak124 model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create125 exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys126 csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32127 csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect128 model.load_state_dict(csd, strict=False) # load129 LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report130 else:131 model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create132 amp = check_amp(model) # check AMP133 134 # Freeze135 freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze136 for k, v in model.named_parameters():137 v.requires_grad = True # train all layers138 # v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results)139 if any(x in k for x in freeze):140 LOGGER.info(f'freezing {k}')141 v.requires_grad = False142 143 # Image size144 gs = max(int(model.stride.max()), 32) # grid size (max stride)145 imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple146 147 # Batch size148 if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size149 batch_size = check_train_batch_size(model, imgsz, amp)150 loggers.on_params_update({"batch_size": batch_size})151 152 # Optimizer153 nbs = 64 # nominal batch size154 accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing155 hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay156 optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])157 158 # Scheduler159 if opt.cos_lr:160 lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']161 else:162 lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear163 scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)164 165 # EMA166 ema = ModelEMA(model) if RANK in {-1, 0} else None167 168 # Resume169 best_fitness, start_epoch = 0.0, 0170 if pretrained:171 if resume:172 best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)173 del ckpt, csd174 175 # DP mode176 if cuda and RANK == -1 and torch.cuda.device_count() > 1:177 LOGGER.warning('WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'178 'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')179 model = torch.nn.DataParallel(model)180 181 # SyncBatchNorm182 if opt.sync_bn and cuda and RANK != -1:183 model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)184 LOGGER.info('Using SyncBatchNorm()')185 186 # Trainloader187 train_loader, dataset = create_dataloader(train_path,188 imgsz,189 batch_size // WORLD_SIZE,190 gs,191 single_cls,192 hyp=hyp,193 augment=True,194 cache=None if opt.cache == 'val' else opt.cache,195 rect=opt.rect,196 rank=LOCAL_RANK,197 workers=workers,198 image_weights=opt.image_weights,199 quad=opt.quad,200 prefix=colorstr('train: '),201 shuffle=True,202 seed=opt.seed)203 labels = np.concatenate(dataset.labels, 0)204 mlc = int(labels[:, 0].max()) # max label class205 assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'206 207 # Process 0208 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 222 if not resume:223 if not opt.noautoanchor:224 check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # run AutoAnchor225 model.half().float() # pre-reduce anchor precision226 227 callbacks.run('on_pretrain_routine_end', labels, names)228 229 # DDP mode230 if cuda and RANK != -1:231 model = smart_DDP(model)232 233 # Model attributes234 nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps)235 hyp['box'] *= 3 / nl # scale to layers236 hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers237 hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers238 hyp['label_smoothing'] = opt.label_smoothing239 model.nc = nc # attach number of classes to model240 model.hyp = hyp # attach hyperparameters to model241 model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights242 model.names = names243 244 # Start training245 t0 = time.time()246 nb = len(train_loader) # number of batches247 nw = max(round(hyp['warmup_epochs'] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations)248 # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training249 last_opt_step = -1250 maps = np.zeros(nc) # mAP per class251 results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)252 scheduler.last_epoch = start_epoch - 1 # do not move253 scaler = torch.cuda.amp.GradScaler(enabled=amp)254 stopper, stop = EarlyStopping(patience=opt.patience), False255 compute_loss = ComputeLoss(model) # init loss class256 callbacks.run('on_train_start')257 LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'258 f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'259 f"Logging results to {colorstr('bold', save_dir)}\n"260 f'Starting training for {epochs} epochs...')261 for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------262 callbacks.run('on_train_epoch_start')263 model.train()264 265 # Update image weights (optional, single-GPU only)266 if opt.image_weights:267 cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights268 iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights269 dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx270 271 # Update mosaic border (optional)272 # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)273 # dataset.mosaic_border = [b - imgsz, -b] # height, width borders274 275 mloss = torch.zeros(3, device=device) # mean losses276 if RANK != -1:277 train_loader.sampler.set_epoch(epoch)278 pbar = enumerate(train_loader)279 LOGGER.info(('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))280 if RANK in {-1, 0}:281 pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # progress bar282 optimizer.zero_grad()283 for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------284 callbacks.run('on_train_batch_start')285 ni = i + nb * epoch # number integrated batches (since train start)286 imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0287 288 # Warmup289 if ni <= nw:290 xi = [0, nw] # x interp291 # 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 lr0295 x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 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 299 # Multi-scale300 if opt.multi_scale:301 sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size302 sf = sz / max(imgs.shape[2:]) # scale factor303 if sf != 1:304 ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)305 imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)306 307 # Forward308 with torch.cuda.amp.autocast(amp):309 pred = model(imgs) # forward310 loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size311 if RANK != -1:312 loss *= WORLD_SIZE # gradient averaged between devices in DDP mode313 if opt.quad:314 loss *= 4.315 316 # Backward317 scaler.scale(loss).backward()318 319 # Optimize - https://pytorch.org/docs/master/notes/amp_examples.html320 if ni - last_opt_step >= accumulate:321 scaler.unscale_(optimizer) # unscale gradients322 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients323 scaler.step(optimizer) # optimizer.step324 scaler.update()325 optimizer.zero_grad()326 if ema:327 ema.update(model)328 last_opt_step = ni329 330 # Log331 if RANK in {-1, 0}:332 mloss = (mloss * i + loss_items) / (i + 1) # update mean losses333 mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)334 pbar.set_description(('%11s' * 2 + '%11.4g' * 5) %335 (f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))336 callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths, list(mloss))337 if callbacks.stop_training:338 return339 # end batch ------------------------------------------------------------------------------------------------340 341 # Scheduler342 lr = [x['lr'] for x in optimizer.param_groups] # for loggers343 scheduler.step()344 345 if RANK in {-1, 0}:346 # mAP347 callbacks.run('on_train_epoch_end', epoch=epoch)348 ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])349 final_epoch = (epoch + 1 == epochs) or stopper.possible_stop350 if not noval or final_epoch: # Calculate mAP351 results, maps, _ = validate.run(data_dict,352 batch_size=batch_size // WORLD_SIZE * 2,353 imgsz=imgsz,354 half=amp,355 model=ema.ema,356 single_cls=single_cls,357 dataloader=val_loader,358 save_dir=save_dir,359 plots=False,360 callbacks=callbacks,361 compute_loss=compute_loss)362 363 # Update best mAP364 fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]365 stop = stopper(epoch=epoch, fitness=fi) # early stop check366 if fi > best_fitness:367 best_fitness = fi368 log_vals = list(mloss) + list(results) + lr369 callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)370 371 # Save model372 if (not nosave) or (final_epoch and not evolve): # if save373 ckpt = {374 'epoch': epoch,375 'best_fitness': best_fitness,376 'model': deepcopy(de_parallel(model)).half(),377 'ema': deepcopy(ema.ema).half(),378 'updates': ema.updates,379 'optimizer': optimizer.state_dict(),380 'opt': vars(opt),381 'git': GIT_INFO, # {remote, branch, commit} if a git repo382 'date': datetime.now().isoformat()}383 384 # Save last, best and delete385 torch.save(ckpt, last)386 if best_fitness == fi:387 torch.save(ckpt, best)388 if opt.save_period > 0 and epoch % opt.save_period == 0:389 torch.save(ckpt, w / f'epoch{epoch}.pt')390 del ckpt391 callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)392 393 # EarlyStopping394 if RANK != -1: # if DDP training395 broadcast_list = [stop if RANK == 0 else None]396 dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks397 if RANK != 0:398 stop = broadcast_list[0]399 if stop:400 break # must break all DDP ranks401 402 # end epoch ----------------------------------------------------------------------------------------------------403 # end training -----------------------------------------------------------------------------------------------------404 if RANK in {-1, 0}:405 LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')406 for f in last, best:407 if f.exists():408 strip_optimizer(f) # strip optimizers409 if f is best:410 LOGGER.info(f'\nValidating {f}...')411 results, _, _ = validate.run(412 data_dict,413 batch_size=batch_size // WORLD_SIZE * 2,414 imgsz=imgsz,415 model=attempt_load(f, device).half(),416 iou_thres=0.65 if is_coco else 0.60, # best pycocotools at iou 0.65417 single_cls=single_cls,418 dataloader=val_loader,419 save_dir=save_dir,420 save_json=is_coco,421 verbose=True,422 plots=plots,423 callbacks=callbacks,424 compute_loss=compute_loss) # val best model with plots425 if is_coco:426 callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)427 428 callbacks.run('on_train_end', last, best, epoch, results)429 430 torch.cuda.empty_cache()431 return results432 433 434def parse_opt(known=False):435 parser = argparse.ArgumentParser()436 parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')437 parser.add_argument('--cfg', type=str, default='', help='model.yaml path')438 439 parser.add_argument('--data', type=str, default=ROOT / 'data/mydata.yaml', help='dataset.yaml path')440 441 parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')442 parser.add_argument('--epochs', type=int, default=200, help='total training epochs')443 parser.add_argument('--batch-size', type=int, default=40, help='total batch size for all GPUs, -1 for autobatch')444 parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')445 parser.add_argument('--rect', action='store_true', help='rectangular training')446 parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')447 parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')448 parser.add_argument('--noval', action='store_true', help='only validate final epoch')449 parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')450 parser.add_argument('--noplots', action='store_true', help='save no plot files')451 parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')452 parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')453 parser.add_argument('--cache', type=str, nargs='?', const='ram', help='image --cache ram/disk')454 parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')455 parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')456 parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')457 parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')458 parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')459 parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')460 parser.add_argument('--workers', type=int, default=1, help='max dataloader workers (per RANK in DDP mode)')461 parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')462 parser.add_argument('--name', default='exp', help='save to project/name')463 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')464 parser.add_argument('--quad', action='store_true', help='quad dataloader')465 parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')466 parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')467 parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')468 parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')469 parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')470 parser.add_argument('--seed', type=int, default=0, help='Global training seed')471 parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')472 473 # Logger arguments474 parser.add_argument('--entity', default=None, help='Entity')475 parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')476 parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')477 parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')478 479 return parser.parse_known_args()[0] if known else parser.parse_args()480 481 482def main(opt, callbacks=Callbacks()):483 # Checks484 if RANK in {-1, 0}:485 print_args(vars(opt))486 check_git_status()487 check_requirements()488 489 # Resume (from specified or most recent last.pt)490 if opt.resume and not check_comet_resume(opt) and not opt.evolve:491 last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())492 opt_yaml = last.parent.parent / 'opt.yaml' # train options yaml493 opt_data = opt.data # original dataset494 if opt_yaml.is_file():495 with open(opt_yaml, errors='ignore') as f:496 d = yaml.safe_load(f)497 else:498 d = torch.load(last, map_location='cpu')['opt']499 opt = argparse.Namespace(**d) # replace500 opt.cfg, opt.weights, opt.resume = '', str(last), True # reinstate501 if is_url(opt_data):502 opt.data = check_file(opt_data) # avoid HUB resume auth timeout503 else:504 opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \505 check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks506 assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'507 if opt.evolve:508 if opt.project == str(ROOT / 'runs/train'): # if default project name, rename to runs/evolve509 opt.project = str(ROOT / 'runs/evolve')510 opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume511 if opt.name == 'cfg':512 opt.name = Path(opt.cfg).stem # use model.yaml as name513 opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))514 515 # DDP mode516 device = select_device(opt.device, batch_size=opt.batch_size)517 if LOCAL_RANK != -1:518 msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'519 assert not opt.image_weights, f'--image-weights {msg}'520 assert not opt.evolve, f'--evolve {msg}'521 assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'522 assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'523 assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'524 torch.cuda.set_device(LOCAL_RANK)525 device = torch.device('cuda', LOCAL_RANK)526 dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")527 528 # Train529 if not opt.evolve:530 train(opt.hyp, opt, device, callbacks)531 532 # Evolve hyperparameters (optional)533 else:534 # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)535 meta = {536 'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)537 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)538 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1539 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay540 'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)541 'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum542 'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr543 'box': (1, 0.02, 0.2), # box loss gain544 'cls': (1, 0.2, 4.0), # cls loss gain545 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight546 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)547 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight548 'iou_t': (0, 0.1, 0.7), # IoU training threshold549 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold550 'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)551 'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)552 'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)553 'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)554 'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)555 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)556 'translate': (1, 0.0, 0.9), # image translation (+/- fraction)557 'scale': (1, 0.0, 0.9), # image scale (+/- gain)558 'shear': (1, 0.0, 10.0), # image shear (+/- deg)559 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001560 'flipud': (1, 0.0, 1.0), # image flip up-down (probability)561 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)562 'mosaic': (1, 0.0, 1.0), # image mixup (probability)563 'mixup': (1, 0.0, 1.0), # image mixup (probability)564 'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)565 566 with open(opt.hyp, errors='ignore') as f:567 hyp = yaml.safe_load(f) # load hyps dict568 if 'anchors' not in hyp: # anchors commented in hyp.yaml569 hyp['anchors'] = 3570 if opt.noautoanchor:571 del hyp['anchors'], meta['anchors']572 opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch573 # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices574 evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'575 if opt.bucket:576 os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {evolve_csv}') # download evolve.csv if exists577 578 for _ in range(opt.evolve): # generations to evolve579 if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate580 # Select parent(s)581 parent = 'single' # parent selection method: 'single' or 'weighted'582 x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)583 n = min(5, len(x)) # number of previous results to consider584 x = x[np.argsort(-fitness(x))][:n] # top n mutations585 w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)586 if parent == 'single' or len(x) == 1:587 # x = x[random.randint(0, n - 1)] # random selection588 x = x[random.choices(range(n), weights=w)[0]] # weighted selection589 elif parent == 'weighted':590 x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination591 592 # Mutate593 mp, s = 0.8, 0.2 # mutation probability, sigma594 npr = np.random595 npr.seed(int(time.time()))596 g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1597 ng = len(meta)598 v = np.ones(ng)599 while all(v == 1): # mutate until a change occurs (prevent duplicates)600 v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)601 for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)602 hyp[k] = float(x[i + 7] * v[i]) # mutate603 604 # Constrain to limits605 for k, v in meta.items():606 hyp[k] = max(hyp[k], v[1]) # lower limit607 hyp[k] = min(hyp[k], v[2]) # upper limit608 hyp[k] = round(hyp[k], 5) # significant digits609 610 # Train mutation611 results = train(hyp.copy(), opt, device, callbacks)612 callbacks = Callbacks()613 # Write mutation results614 keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss',615 'val/obj_loss', 'val/cls_loss')616 print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket)617 618 # Plot results619 plot_evolve(evolve_csv)620 LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'621 f"Results saved to {colorstr('bold', save_dir)}\n"622 f'Usage example: $ python train.py --hyp {evolve_yaml}')623 624 625def run(**kwargs):626 # Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')627 opt = parse_opt(True)628 for k, v in kwargs.items():629 setattr(opt, k, v)630 main(opt)631 return opt632 633 634if __name__ == "__main__":635 opt = parse_opt()636 main(opt)637