From aa637bb6f0c991ba02e526137568c25b78d2853e Mon Sep 17 00:00:00 2001 From: Jacob Brown Date: Mon, 5 Aug 2024 13:47:39 -0600 Subject: [PATCH 1/9] update torch.cuda.amp to torch.amp --- classify/val.py | 2 +- segment/train.py | 4 ++-- train.py | 4 ++-- utils/autobatch.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/classify/val.py b/classify/val.py index 8ce48f0645bf..244597a7963d 100644 --- a/classify/val.py +++ b/classify/val.py @@ -108,7 +108,7 @@ def run( action = "validating" if dataloader.dataset.root.stem == "val" else "testing" desc = f"{pbar.desc[:-36]}{action:>36}" if pbar else f"{action}" bar = tqdm(dataloader, desc, n, not training, bar_format=TQDM_BAR_FORMAT, position=0) - with torch.cuda.amp.autocast(enabled=device.type != "cpu"): + with torch.amp.autocast("cuda", enabled=device.type != "cpu"): for images, labels in bar: with dt[0]: images, labels = images.to(device, non_blocking=True), labels.to(device) diff --git a/segment/train.py b/segment/train.py index 379fed0b2f14..6654e2a9edd7 100644 --- a/segment/train.py +++ b/segment/train.py @@ -320,7 +320,7 @@ def lf(x): maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) scheduler.last_epoch = start_epoch - 1 # do not move - scaler = torch.cuda.amp.GradScaler(enabled=amp) + scaler = torch.amp.GradScaler("cuda", enabled=amp) stopper, stop = EarlyStopping(patience=opt.patience), False compute_loss = ComputeLoss(model, overlap=overlap) # init loss class # callbacks.run('on_train_start') @@ -380,7 +380,7 @@ def lf(x): imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) # Forward - with torch.cuda.amp.autocast(amp): + with torch.amp.autocast("cuda", enabled=amp): pred = model(imgs) # forward loss, loss_items = compute_loss(pred, targets.to(device), masks=masks.to(device).float()) if RANK != -1: diff --git a/train.py b/train.py index b4395d7e8d15..8ab5256b324f 100644 --- a/train.py +++ b/train.py @@ -352,7 +352,7 @@ def lf(x): maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) scheduler.last_epoch = start_epoch - 1 # do not move - scaler = torch.cuda.amp.GradScaler(enabled=amp) + scaler = torch.amp.GradScaler("cuda", enabled=amp) stopper, stop = EarlyStopping(patience=opt.patience), False compute_loss = ComputeLoss(model) # init loss class callbacks.run("on_train_start") @@ -409,7 +409,7 @@ def lf(x): imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) # Forward - with torch.cuda.amp.autocast(amp): + with torch.amp.autocast("cuda", enabled=amp): pred = model(imgs) # forward loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size if RANK != -1: diff --git a/utils/autobatch.py b/utils/autobatch.py index 08a0de841a98..23e5e6593c35 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -12,7 +12,7 @@ def check_train_batch_size(model, imgsz=640, amp=True): """Checks and computes optimal training batch size for YOLOv5 model, given image size and AMP setting.""" - with torch.cuda.amp.autocast(amp): + with torch.amp.autocast("cuda", enabled=amp): return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size From 98f4925fc700b5bac97c8d49997986cd1747426d Mon Sep 17 00:00:00 2001 From: Jacob Brown Date: Mon, 5 Aug 2024 14:25:25 -0600 Subject: [PATCH 2/9] fix to support torch versions <2.4 --- classify/val.py | 10 +++++++++- segment/train.py | 18 ++++++++++++++++-- train.py | 17 +++++++++++++++-- utils/autobatch.py | 7 +++++-- 4 files changed, 45 insertions(+), 7 deletions(-) diff --git a/classify/val.py b/classify/val.py index 244597a7963d..d28d31ae4302 100644 --- a/classify/val.py +++ b/classify/val.py @@ -42,6 +42,7 @@ Profile, check_img_size, check_requirements, + check_version, colorstr, increment_path, print_args, @@ -108,7 +109,14 @@ def run( action = "validating" if dataloader.dataset.root.stem == "val" else "testing" desc = f"{pbar.desc[:-36]}{action:>36}" if pbar else f"{action}" bar = tqdm(dataloader, desc, n, not training, bar_format=TQDM_BAR_FORMAT, position=0) - with torch.amp.autocast("cuda", enabled=device.type != "cpu"): + + amp_autocast = None + if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + amp_autocast = torch.amp.autocast("cuda", enabled=device.type != "cpu") + else: + amp_autocast = torch.cuda.amp.autocast(enabled=device.type != "cpu") + + with amp_autocast: for images, labels in bar: with dt[0]: images, labels = images.to(device, non_blocking=True), labels.to(device) diff --git a/segment/train.py b/segment/train.py index 6654e2a9edd7..194b003b471f 100644 --- a/segment/train.py +++ b/segment/train.py @@ -58,6 +58,7 @@ check_img_size, check_requirements, check_suffix, + check_version, check_yaml, colorstr, get_latest_run, @@ -320,7 +321,13 @@ def lf(x): maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) scheduler.last_epoch = start_epoch - 1 # do not move - scaler = torch.amp.GradScaler("cuda", enabled=amp) + + scaler = None + if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + scaler = torch.amp.GradScaler("cuda", enabled=amp) + else: + scaler = torch.cuda.amp.GradScaler(enabled=amp) + stopper, stop = EarlyStopping(patience=opt.patience), False compute_loss = ComputeLoss(model, overlap=overlap) # init loss class # callbacks.run('on_train_start') @@ -379,8 +386,15 @@ def lf(x): ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) + + amp_autocast = None + if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + amp_autocast = torch.amp.autocast("cuda", enabled=amp) + else: + amp_autocast = torch.cuda.amp.autocast(enabled=amp) + # Forward - with torch.amp.autocast("cuda", enabled=amp): + with amp_autocast: pred = model(imgs) # forward loss, loss_items = compute_loss(pred, targets.to(device), masks=masks.to(device).float()) if RANK != -1: diff --git a/train.py b/train.py index 8ab5256b324f..dc835d78e0b6 100644 --- a/train.py +++ b/train.py @@ -63,6 +63,7 @@ check_img_size, check_requirements, check_suffix, + check_version, check_yaml, colorstr, get_latest_run, @@ -352,7 +353,13 @@ def lf(x): maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) scheduler.last_epoch = start_epoch - 1 # do not move - scaler = torch.amp.GradScaler("cuda", enabled=amp) + + scaler = None + if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + scaler = torch.amp.GradScaler("cuda", enabled=amp) + else: + scaler = torch.cuda.amp.GradScaler(enabled=amp) + stopper, stop = EarlyStopping(patience=opt.patience), False compute_loss = ComputeLoss(model) # init loss class callbacks.run("on_train_start") @@ -408,8 +415,14 @@ def lf(x): ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) + amp_autocast = None + if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + amp_autocast = torch.amp.autocast("cuda", enabled=amp) + else: + amp_autocast = torch.cuda.amp.autocast(amp) + # Forward - with torch.amp.autocast("cuda", enabled=amp): + with amp_autocast: pred = model(imgs) # forward loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size if RANK != -1: diff --git a/utils/autobatch.py b/utils/autobatch.py index 23e5e6593c35..c8deb940fdf1 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -6,13 +6,16 @@ import numpy as np import torch -from utils.general import LOGGER, colorstr +from utils.general import LOGGER, check_version, colorstr from utils.torch_utils import profile def check_train_batch_size(model, imgsz=640, amp=True): """Checks and computes optimal training batch size for YOLOv5 model, given image size and AMP setting.""" - with torch.amp.autocast("cuda", enabled=amp): + if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + with torch.amp.autocast("cuda", enabled=amp): + return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size + with torch.cuda.amp.autocast(amp): return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size From 15f73ba2d3dc5552f73555486a0a707e852ed45e Mon Sep 17 00:00:00 2001 From: UltralyticsAssistant Date: Mon, 5 Aug 2024 20:26:18 +0000 Subject: [PATCH 3/9] Auto-format by https://ultralytics.com/actions --- classify/val.py | 2 +- segment/train.py | 1 - train.py | 4 ++-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/classify/val.py b/classify/val.py index d28d31ae4302..11e682f4891d 100644 --- a/classify/val.py +++ b/classify/val.py @@ -109,7 +109,7 @@ def run( action = "validating" if dataloader.dataset.root.stem == "val" else "testing" desc = f"{pbar.desc[:-36]}{action:>36}" if pbar else f"{action}" bar = tqdm(dataloader, desc, n, not training, bar_format=TQDM_BAR_FORMAT, position=0) - + amp_autocast = None if check_version(torch.__version__, "2.4.0", "Torch", hard=False): amp_autocast = torch.amp.autocast("cuda", enabled=device.type != "cpu") diff --git a/segment/train.py b/segment/train.py index 194b003b471f..c9924d54ec0f 100644 --- a/segment/train.py +++ b/segment/train.py @@ -386,7 +386,6 @@ def lf(x): ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) - amp_autocast = None if check_version(torch.__version__, "2.4.0", "Torch", hard=False): amp_autocast = torch.amp.autocast("cuda", enabled=amp) diff --git a/train.py b/train.py index dc835d78e0b6..44d64bded35f 100644 --- a/train.py +++ b/train.py @@ -353,7 +353,7 @@ def lf(x): maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) scheduler.last_epoch = start_epoch - 1 # do not move - + scaler = None if check_version(torch.__version__, "2.4.0", "Torch", hard=False): scaler = torch.amp.GradScaler("cuda", enabled=amp) @@ -420,7 +420,7 @@ def lf(x): amp_autocast = torch.amp.autocast("cuda", enabled=amp) else: amp_autocast = torch.cuda.amp.autocast(amp) - + # Forward with amp_autocast: pred = model(imgs) # forward From 7b1c0f40e01ede53b1c4c91647768301082dd9d0 Mon Sep 17 00:00:00 2001 From: Jacob Brown Date: Mon, 5 Aug 2024 16:08:19 -0600 Subject: [PATCH 4/9] slight edits and fixed missed torch.cuda.amp calls --- classify/train.py | 19 +- classify/val.py | 2 +- models/common.py | 8 +- segment/train.py | 4 +- train.py | 532 ++++++++++++++++++++++++++++++++++++--------- utils/autobatch.py | 26 ++- 6 files changed, 469 insertions(+), 122 deletions(-) diff --git a/classify/train.py b/classify/train.py index 9c12a66c326f..75fafc7cff01 100644 --- a/classify/train.py +++ b/classify/train.py @@ -27,7 +27,6 @@ import torch.hub as hub import torch.optim.lr_scheduler as lr_scheduler import torchvision -from torch.cuda import amp from tqdm import tqdm FILE = Path(__file__).resolve() @@ -48,6 +47,7 @@ check_git_info, check_git_status, check_requirements, + check_version, colorstr, download, increment_path, @@ -198,7 +198,13 @@ def lf(x): t0 = time.time() criterion = smartCrossEntropyLoss(label_smoothing=opt.label_smoothing) # loss function best_fitness = 0.0 - scaler = amp.GradScaler(enabled=cuda) + + scaler = None + if check_version(torch.__version__, "2.4.0"): + scaler = torch.amp.GradScaler("cuda", enabled=cuda) + else: + scaler = torch.cuda.amp.GradScaler(enabled=cuda) + val = test_dir.stem # 'val' or 'test' LOGGER.info( f'Image sizes {imgsz} train, {imgsz} test\n' @@ -218,8 +224,15 @@ def lf(x): for i, (images, labels) in pbar: # progress bar images, labels = images.to(device, non_blocking=True), labels.to(device) + + amp_autocast = None + if check_version(torch.__version__, "2.4.0"): + amp_autocast = torch.amp.autocast("cuda", enabled=device.type != "cpu") + else: + amp_autocast = torch.cuda.amp.autocast(enabled=device.type != "cpu") + # Forward - with amp.autocast(enabled=cuda): # stability issues when enabled + with amp_autocast: # stability issues when enabled loss = criterion(model(images), labels) # Backward diff --git a/classify/val.py b/classify/val.py index d28d31ae4302..4bfd42d2c593 100644 --- a/classify/val.py +++ b/classify/val.py @@ -111,7 +111,7 @@ def run( bar = tqdm(dataloader, desc, n, not training, bar_format=TQDM_BAR_FORMAT, position=0) amp_autocast = None - if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + if check_version(torch.__version__, "2.4.0"): amp_autocast = torch.amp.autocast("cuda", enabled=device.type != "cpu") else: amp_autocast = torch.cuda.amp.autocast(enabled=device.type != "cpu") diff --git a/models/common.py b/models/common.py index 1e0ffdd3abdb..d735e10fed8a 100644 --- a/models/common.py +++ b/models/common.py @@ -20,7 +20,6 @@ import torch import torch.nn as nn from PIL import Image -from torch.cuda import amp # Import 'ultralytics' package or install if missing try: @@ -839,7 +838,12 @@ def forward(self, ims, size=640, augment=False, profile=False): p = next(self.model.parameters()) if self.pt else torch.empty(1, device=self.model.device) # param autocast = self.amp and (p.device.type != "cpu") # Automatic Mixed Precision (AMP) inference if isinstance(ims, torch.Tensor): # torch - with amp.autocast(autocast): + amp_autocast = None + if check_version(torch.__version__, "2.4.0"): + amp_autocast = torch.amp.autocast("cuda", enabled=autocast) + else: + amp_autocast = torch.cuda.amp.autocast(enabled=autocast) + with amp_autocast: return self.model(ims.to(p.device).type_as(p), augment=augment) # inference # Pre-process diff --git a/segment/train.py b/segment/train.py index 194b003b471f..0bbf976435c1 100644 --- a/segment/train.py +++ b/segment/train.py @@ -323,7 +323,7 @@ def lf(x): scheduler.last_epoch = start_epoch - 1 # do not move scaler = None - if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + if check_version(torch.__version__, "2.4.0"): scaler = torch.amp.GradScaler("cuda", enabled=amp) else: scaler = torch.cuda.amp.GradScaler(enabled=amp) @@ -388,7 +388,7 @@ def lf(x): amp_autocast = None - if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + if check_version(torch.__version__, "2.4.0"): amp_autocast = torch.amp.autocast("cuda", enabled=amp) else: amp_autocast = torch.cuda.amp.autocast(enabled=amp) diff --git a/train.py b/train.py index dc835d78e0b6..3b62237495bd 100644 --- a/train.py +++ b/train.py @@ -95,7 +95,9 @@ torch_distributed_zero_first, ) -LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1)) # https://pytorch.org/docs/stable/elastic/run.html +LOCAL_RANK = int( + os.getenv("LOCAL_RANK", -1) +) # https://pytorch.org/docs/stable/elastic/run.html RANK = int(os.getenv("RANK", -1)) WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1)) GIT_INFO = check_git_info() @@ -135,7 +137,21 @@ def train(hyp, opt, device, callbacks): - Datasets: https://github.com/ultralytics/yolov5/tree/master/data - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data """ - save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = ( + ( + save_dir, + epochs, + batch_size, + weights, + single_cls, + evolve, + data, + cfg, + resume, + noval, + nosave, + workers, + freeze, + ) = ( Path(opt.save_dir), opt.epochs, opt.batch_size, @@ -161,7 +177,9 @@ def train(hyp, opt, device, callbacks): if isinstance(hyp, str): with open(hyp, errors="ignore") as f: hyp = yaml.safe_load(f) # load hyps dict - LOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items())) + LOGGER.info( + colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items()) + ) opt.hyp = hyp.copy() # for saving hyps to checkpoints # Save run settings @@ -194,7 +212,12 @@ def train(hyp, opt, device, callbacks): # Process custom dataset artifact link data_dict = loggers.remote_dataset if resume: # If resuming runs from remote artifact - weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size + weights, epochs, hyp, batch_size = ( + opt.weights, + opt.epochs, + opt.hyp, + opt.batch_size, + ) # Config plots = not evolve and not opt.noplots # create plots @@ -204,8 +227,14 @@ def train(hyp, opt, device, callbacks): data_dict = data_dict or check_dataset(data) # check if None train_path, val_path = data_dict["train"], data_dict["val"] nc = 1 if single_cls else int(data_dict["nc"]) # number of classes - names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"] # class names - is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt") # COCO dataset + names = ( + {0: "item"} + if single_cls and len(data_dict["names"]) != 1 + else data_dict["names"] + ) # class names + is_coco = isinstance(val_path, str) and val_path.endswith( + "coco/val2017.txt" + ) # COCO dataset # Model check_suffix(weights, ".pt") # check weights @@ -213,19 +242,31 @@ def train(hyp, opt, device, callbacks): if pretrained: with torch_distributed_zero_first(LOCAL_RANK): weights = attempt_download(weights) # download if not found locally - ckpt = torch.load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak - model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create - exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] # exclude keys + ckpt = torch.load( + weights, map_location="cpu" + ) # load checkpoint to CPU to avoid CUDA memory leak + model = Model( + cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors") + ).to( + device + ) # create + exclude = ( + ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] + ) # exclude keys csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32 csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(csd, strict=False) # load - LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}") # report + LOGGER.info( + f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}" + ) # report else: model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create amp = check_amp(model) # check AMP # Freeze - freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze + freeze = [ + f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0])) + ] # layers to freeze for k, v in model.named_parameters(): v.requires_grad = True # train all layers # v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results) @@ -246,7 +287,9 @@ def train(hyp, opt, device, callbacks): nbs = 64 # nominal batch size accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing hyp["weight_decay"] *= batch_size * accumulate / nbs # scale weight_decay - optimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"]) + optimizer = smart_optimizer( + model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"] + ) # Scheduler if opt.cos_lr: @@ -257,7 +300,9 @@ def lf(x): """Linear learning rate scheduler function with decay calculated by epoch proportion.""" return (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"] # linear - scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs) + scheduler = lr_scheduler.LambdaLR( + optimizer, lr_lambda=lf + ) # plot_lr_scheduler(optimizer, scheduler, epochs) # EMA ema = ModelEMA(model) if RANK in {-1, 0} else None @@ -266,7 +311,9 @@ def lf(x): best_fitness, start_epoch = 0.0, 0 if pretrained: if resume: - best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume) + best_fitness, start_epoch, epochs = smart_resume( + ckpt, optimizer, ema, weights, epochs, resume + ) del ckpt, csd # DP mode @@ -303,7 +350,9 @@ def lf(x): ) labels = np.concatenate(dataset.labels, 0) mlc = int(labels[:, 0].max()) # max label class - assert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}" + assert ( + mlc < nc + ), f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}" # Process 0 if RANK in {-1, 0}: @@ -324,7 +373,9 @@ def lf(x): if not resume: if not opt.noautoanchor: - check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz) # run AutoAnchor + check_anchors( + dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz + ) # run AutoAnchor model.half().float() # pre-reduce anchor precision callbacks.run("on_pretrain_routine_end", labels, names) @@ -341,21 +392,25 @@ def lf(x): hyp["label_smoothing"] = opt.label_smoothing model.nc = nc # attach number of classes to model model.hyp = hyp # attach hyperparameters to model - model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights + model.class_weights = ( + labels_to_class_weights(dataset.labels, nc).to(device) * nc + ) # attach class weights model.names = names # Start training t0 = time.time() nb = len(train_loader) # number of batches - nw = max(round(hyp["warmup_epochs"] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations) + nw = max( + round(hyp["warmup_epochs"] * nb), 100 + ) # number of warmup iterations, max(3 epochs, 100 iterations) # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training last_opt_step = -1 maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) scheduler.last_epoch = start_epoch - 1 # do not move - + scaler = None - if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + if check_version(torch.__version__, "2.4.0"): scaler = torch.amp.GradScaler("cuda", enabled=amp) else: scaler = torch.cuda.amp.GradScaler(enabled=amp) @@ -364,20 +419,28 @@ def lf(x): compute_loss = ComputeLoss(model) # init loss class callbacks.run("on_train_start") LOGGER.info( - f'Image sizes {imgsz} train, {imgsz} val\n' - f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n' + f"Image sizes {imgsz} train, {imgsz} val\n" + f"Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n" f"Logging results to {colorstr('bold', save_dir)}\n" - f'Starting training for {epochs} epochs...' + f"Starting training for {epochs} epochs..." ) - for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ + for epoch in range( + start_epoch, epochs + ): # epoch ------------------------------------------------------------------ callbacks.run("on_train_epoch_start") model.train() # Update image weights (optional, single-GPU only) if opt.image_weights: - cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights - iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights - dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx + cw = ( + model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc + ) # class weights + iw = labels_to_image_weights( + dataset.labels, nc=nc, class_weights=cw + ) # image weights + dataset.indices = random.choices( + range(dataset.n), weights=iw, k=dataset.n + ) # rand weighted idx # Update mosaic border (optional) # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) @@ -387,14 +450,34 @@ def lf(x): if RANK != -1: train_loader.sampler.set_epoch(epoch) pbar = enumerate(train_loader) - LOGGER.info(("\n" + "%11s" * 7) % ("Epoch", "GPU_mem", "box_loss", "obj_loss", "cls_loss", "Instances", "Size")) + LOGGER.info( + ("\n" + "%11s" * 7) + % ( + "Epoch", + "GPU_mem", + "box_loss", + "obj_loss", + "cls_loss", + "Instances", + "Size", + ) + ) if RANK in {-1, 0}: pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # progress bar optimizer.zero_grad() - for i, (imgs, targets, paths, _) in pbar: # batch ------------------------------------------------------------- + for i, ( + imgs, + targets, + paths, + _, + ) in ( + pbar + ): # batch ------------------------------------------------------------- callbacks.run("on_train_batch_start") ni = i + nb * epoch # number integrated batches (since train start) - imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0 + imgs = ( + imgs.to(device, non_blocking=True).float() / 255 + ) # uint8 to float32, 0-255 to 0.0-1.0 # Warmup if ni <= nw: @@ -403,28 +486,45 @@ def lf(x): accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round()) for j, x in enumerate(optimizer.param_groups): # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0 - x["lr"] = np.interp(ni, xi, [hyp["warmup_bias_lr"] if j == 0 else 0.0, x["initial_lr"] * lf(epoch)]) + x["lr"] = np.interp( + ni, + xi, + [ + hyp["warmup_bias_lr"] if j == 0 else 0.0, + x["initial_lr"] * lf(epoch), + ], + ) if "momentum" in x: - x["momentum"] = np.interp(ni, xi, [hyp["warmup_momentum"], hyp["momentum"]]) + x["momentum"] = np.interp( + ni, xi, [hyp["warmup_momentum"], hyp["momentum"]] + ) # Multi-scale if opt.multi_scale: - sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs # size + sz = ( + random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs + ) # size sf = sz / max(imgs.shape[2:]) # scale factor if sf != 1: - ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) - imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) + ns = [ + math.ceil(x * sf / gs) * gs for x in imgs.shape[2:] + ] # new shape (stretched to gs-multiple) + imgs = nn.functional.interpolate( + imgs, size=ns, mode="bilinear", align_corners=False + ) amp_autocast = None - if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + if check_version(torch.__version__, "2.4.0"): amp_autocast = torch.amp.autocast("cuda", enabled=amp) else: amp_autocast = torch.cuda.amp.autocast(amp) - + # Forward with amp_autocast: pred = model(imgs) # forward - loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size + loss, loss_items = compute_loss( + pred, targets.to(device) + ) # loss scaled by batch_size if RANK != -1: loss *= WORLD_SIZE # gradient averaged between devices in DDP mode if opt.quad: @@ -436,7 +536,9 @@ def lf(x): # Optimize - https://pytorch.org/docs/master/notes/amp_examples.html if ni - last_opt_step >= accumulate: scaler.unscale_(optimizer) # unscale gradients - torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients + torch.nn.utils.clip_grad_norm_( + model.parameters(), max_norm=10.0 + ) # clip gradients scaler.step(optimizer) # optimizer.step scaler.update() optimizer.zero_grad() @@ -450,9 +552,17 @@ def lf(x): mem = f"{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G" # (GB) pbar.set_description( ("%11s" * 2 + "%11.4g" * 5) - % (f"{epoch}/{epochs - 1}", mem, *mloss, targets.shape[0], imgs.shape[-1]) + % ( + f"{epoch}/{epochs - 1}", + mem, + *mloss, + targets.shape[0], + imgs.shape[-1], + ) + ) + callbacks.run( + "on_train_batch_end", model, ni, imgs, targets, paths, list(mloss) ) - callbacks.run("on_train_batch_end", model, ni, imgs, targets, paths, list(mloss)) if callbacks.stop_training: return # end batch ------------------------------------------------------------------------------------------------ @@ -464,7 +574,9 @@ def lf(x): if RANK in {-1, 0}: # mAP callbacks.run("on_train_epoch_end", epoch=epoch) - ema.update_attr(model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"]) + ema.update_attr( + model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"] + ) final_epoch = (epoch + 1 == epochs) or stopper.possible_stop if not noval or final_epoch: # Calculate mAP results, maps, _ = validate.run( @@ -482,7 +594,9 @@ def lf(x): ) # Update best mAP - fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] + fi = fitness( + np.array(results).reshape(1, -1) + ) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] stop = stopper(epoch=epoch, fitness=fi) # early stop check if fi > best_fitness: best_fitness = fi @@ -510,12 +624,16 @@ def lf(x): if opt.save_period > 0 and epoch % opt.save_period == 0: torch.save(ckpt, w / f"epoch{epoch}.pt") del ckpt - callbacks.run("on_model_save", last, epoch, final_epoch, best_fitness, fi) + callbacks.run( + "on_model_save", last, epoch, final_epoch, best_fitness, fi + ) # EarlyStopping if RANK != -1: # if DDP training broadcast_list = [stop if RANK == 0 else None] - dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks + dist.broadcast_object_list( + broadcast_list, 0 + ) # broadcast 'stop' to all ranks if RANK != 0: stop = broadcast_list[0] if stop: @@ -524,7 +642,9 @@ def lf(x): # end epoch ---------------------------------------------------------------------------------------------------- # end training ----------------------------------------------------------------------------------------------------- if RANK in {-1, 0}: - LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.") + LOGGER.info( + f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours." + ) for f in last, best: if f.exists(): strip_optimizer(f) # strip optimizers @@ -535,7 +655,9 @@ def lf(x): batch_size=batch_size // WORLD_SIZE * 2, imgsz=imgsz, model=attempt_load(f, device).half(), - iou_thres=0.65 if is_coco else 0.60, # best pycocotools at iou 0.65 + iou_thres=( + 0.65 if is_coco else 0.60 + ), # best pycocotools at iou 0.65 single_cls=single_cls, dataloader=val_loader, save_dir=save_dir, @@ -546,7 +668,13 @@ def lf(x): compute_loss=compute_loss, ) # val best model with plots if is_coco: - callbacks.run("on_fit_epoch_end", list(mloss) + list(results) + lr, epoch, best_fitness, fi) + callbacks.run( + "on_fit_epoch_end", + list(mloss) + list(results) + lr, + epoch, + best_fitness, + fi, + ) callbacks.run("on_train_end", last, best, epoch, results) @@ -577,53 +705,176 @@ def parse_opt(known=False): - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data """ parser = argparse.ArgumentParser() - parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path") + parser.add_argument( + "--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path" + ) parser.add_argument("--cfg", type=str, default="", help="model.yaml path") - parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") - parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path") + parser.add_argument( + "--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path" + ) + parser.add_argument( + "--hyp", + type=str, + default=ROOT / "data/hyps/hyp.scratch-low.yaml", + help="hyperparameters path", + ) parser.add_argument("--epochs", type=int, default=100, help="total training epochs") - parser.add_argument("--batch-size", type=int, default=16, help="total batch size for all GPUs, -1 for autobatch") - parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)") + parser.add_argument( + "--batch-size", + type=int, + default=16, + help="total batch size for all GPUs, -1 for autobatch", + ) + parser.add_argument( + "--imgsz", + "--img", + "--img-size", + type=int, + default=640, + help="train, val image size (pixels)", + ) parser.add_argument("--rect", action="store_true", help="rectangular training") - parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training") - parser.add_argument("--nosave", action="store_true", help="only save final checkpoint") - parser.add_argument("--noval", action="store_true", help="only validate final epoch") - parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor") + parser.add_argument( + "--resume", + nargs="?", + const=True, + default=False, + help="resume most recent training", + ) + parser.add_argument( + "--nosave", action="store_true", help="only save final checkpoint" + ) + parser.add_argument( + "--noval", action="store_true", help="only validate final epoch" + ) + parser.add_argument( + "--noautoanchor", action="store_true", help="disable AutoAnchor" + ) parser.add_argument("--noplots", action="store_true", help="save no plot files") - parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations") parser.add_argument( - "--evolve_population", type=str, default=ROOT / "data/hyps", help="location for loading population" + "--evolve", + type=int, + nargs="?", + const=300, + help="evolve hyperparameters for x generations", + ) + parser.add_argument( + "--evolve_population", + type=str, + default=ROOT / "data/hyps", + help="location for loading population", + ) + parser.add_argument( + "--resume_evolve", + type=str, + default=None, + help="resume evolve from last generation", ) - parser.add_argument("--resume_evolve", type=str, default=None, help="resume evolve from last generation") parser.add_argument("--bucket", type=str, default="", help="gsutil bucket") - parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk") - parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training") - parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") - parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%") - parser.add_argument("--single-cls", action="store_true", help="train multi-class data as single-class") - parser.add_argument("--optimizer", type=str, choices=["SGD", "Adam", "AdamW"], default="SGD", help="optimizer") - parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode") - parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)") - parser.add_argument("--project", default=ROOT / "runs/train", help="save to project/name") + parser.add_argument( + "--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk" + ) + parser.add_argument( + "--image-weights", + action="store_true", + help="use weighted image selection for training", + ) + parser.add_argument( + "--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu" + ) + parser.add_argument( + "--multi-scale", action="store_true", help="vary img-size +/- 50%%" + ) + parser.add_argument( + "--single-cls", + action="store_true", + help="train multi-class data as single-class", + ) + parser.add_argument( + "--optimizer", + type=str, + choices=["SGD", "Adam", "AdamW"], + default="SGD", + help="optimizer", + ) + parser.add_argument( + "--sync-bn", + action="store_true", + help="use SyncBatchNorm, only available in DDP mode", + ) + parser.add_argument( + "--workers", + type=int, + default=8, + help="max dataloader workers (per RANK in DDP mode)", + ) + parser.add_argument( + "--project", default=ROOT / "runs/train", help="save to project/name" + ) parser.add_argument("--name", default="exp", help="save to project/name") - parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment") + parser.add_argument( + "--exist-ok", + action="store_true", + help="existing project/name ok, do not increment", + ) parser.add_argument("--quad", action="store_true", help="quad dataloader") parser.add_argument("--cos-lr", action="store_true", help="cosine LR scheduler") - parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon") - parser.add_argument("--patience", type=int, default=100, help="EarlyStopping patience (epochs without improvement)") - parser.add_argument("--freeze", nargs="+", type=int, default=[0], help="Freeze layers: backbone=10, first3=0 1 2") - parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)") + parser.add_argument( + "--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon" + ) + parser.add_argument( + "--patience", + type=int, + default=100, + help="EarlyStopping patience (epochs without improvement)", + ) + parser.add_argument( + "--freeze", + nargs="+", + type=int, + default=[0], + help="Freeze layers: backbone=10, first3=0 1 2", + ) + parser.add_argument( + "--save-period", + type=int, + default=-1, + help="Save checkpoint every x epochs (disabled if < 1)", + ) parser.add_argument("--seed", type=int, default=0, help="Global training seed") - parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify") + parser.add_argument( + "--local_rank", + type=int, + default=-1, + help="Automatic DDP Multi-GPU argument, do not modify", + ) # Logger arguments parser.add_argument("--entity", default=None, help="Entity") - parser.add_argument("--upload_dataset", nargs="?", const=True, default=False, help='Upload data, "val" option') - parser.add_argument("--bbox_interval", type=int, default=-1, help="Set bounding-box image logging interval") - parser.add_argument("--artifact_alias", type=str, default="latest", help="Version of dataset artifact to use") + parser.add_argument( + "--upload_dataset", + nargs="?", + const=True, + default=False, + help='Upload data, "val" option', + ) + parser.add_argument( + "--bbox_interval", + type=int, + default=-1, + help="Set bounding-box image logging interval", + ) + parser.add_argument( + "--artifact_alias", + type=str, + default="latest", + help="Version of dataset artifact to use", + ) # NDJSON logging - parser.add_argument("--ndjson-console", action="store_true", help="Log ndjson to console") + parser.add_argument( + "--ndjson-console", action="store_true", help="Log ndjson to console" + ) parser.add_argument("--ndjson-file", action="store_true", help="Log ndjson to file") return parser.parse_known_args()[0] if known else parser.parse_args() @@ -652,7 +903,9 @@ def main(opt, callbacks=Callbacks()): # Resume (from specified or most recent last.pt) if opt.resume and not check_comet_resume(opt) and not opt.evolve: - last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run()) + last = Path( + check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run() + ) opt_yaml = last.parent.parent / "opt.yaml" # train options yaml opt_data = opt.data # original dataset if opt_yaml.is_file(): @@ -672,14 +925,23 @@ def main(opt, callbacks=Callbacks()): str(opt.weights), str(opt.project), ) # checks - assert len(opt.cfg) or len(opt.weights), "either --cfg or --weights must be specified" + assert len(opt.cfg) or len( + opt.weights + ), "either --cfg or --weights must be specified" if opt.evolve: - if opt.project == str(ROOT / "runs/train"): # if default project name, rename to runs/evolve + if opt.project == str( + ROOT / "runs/train" + ): # if default project name, rename to runs/evolve opt.project = str(ROOT / "runs/evolve") - opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume + opt.exist_ok, opt.resume = ( + opt.resume, + False, + ) # pass resume to exist_ok and disable resume if opt.name == "cfg": opt.name = Path(opt.cfg).stem # use model.yaml as name - opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) + opt.save_dir = str( + increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) + ) # DDP mode device = select_device(opt.device, batch_size=opt.batch_size) @@ -687,13 +949,20 @@ def main(opt, callbacks=Callbacks()): msg = "is not compatible with YOLOv5 Multi-GPU DDP training" assert not opt.image_weights, f"--image-weights {msg}" assert not opt.evolve, f"--evolve {msg}" - assert opt.batch_size != -1, f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size" - assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE" - assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command" + assert ( + opt.batch_size != -1 + ), f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size" + assert ( + opt.batch_size % WORLD_SIZE == 0 + ), f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE" + assert ( + torch.cuda.device_count() > LOCAL_RANK + ), "insufficient CUDA devices for DDP command" torch.cuda.set_device(LOCAL_RANK) device = torch.device("cuda", LOCAL_RANK) dist.init_process_group( - backend="nccl" if dist.is_nccl_available() else "gloo", timeout=timedelta(seconds=10800) + backend="nccl" if dist.is_nccl_available() else "gloo", + timeout=timedelta(seconds=10800), ) # Train @@ -719,7 +988,11 @@ def main(opt, callbacks=Callbacks()): "iou_t": (False, 0.1, 0.7), # IoU training threshold "anchor_t": (False, 2.0, 8.0), # anchor-multiple threshold "anchors": (False, 2.0, 10.0), # anchors per output grid (0 to ignore) - "fl_gamma": (False, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) + "fl_gamma": ( + False, + 0.0, + 2.0, + ), # focal loss gamma (efficientDet default gamma=1.5) "hsv_h": (True, 0.0, 0.1), # image HSV-Hue augmentation (fraction) "hsv_s": (True, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) "hsv_v": (True, 0.0, 0.9), # image HSV-Value augmentation (fraction) @@ -727,7 +1000,11 @@ def main(opt, callbacks=Callbacks()): "translate": (True, 0.0, 0.9), # image translation (+/- fraction) "scale": (True, 0.0, 0.9), # image scale (+/- gain) "shear": (True, 0.0, 10.0), # image shear (+/- deg) - "perspective": (True, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 + "perspective": ( + True, + 0.0, + 0.001, + ), # image perspective (+/- fraction), range 0-0.001 "flipud": (True, 0.0, 1.0), # image flip up-down (probability) "fliplr": (True, 0.0, 1.0), # image flip left-right (probability) "mosaic": (True, 0.0, 1.0), # image mixup (probability) @@ -752,7 +1029,11 @@ def main(opt, callbacks=Callbacks()): hyp["anchors"] = 3 if opt.noautoanchor: del hyp["anchors"], meta["anchors"] - opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch + opt.noval, opt.nosave, save_dir = ( + True, + True, + Path(opt.save_dir), + ) # only val/save final epoch # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices evolve_yaml, evolve_csv = save_dir / "hyp_evolve.yaml", save_dir / "evolve.csv" if opt.bucket: @@ -778,14 +1059,18 @@ def main(opt, callbacks=Callbacks()): upper_limit = np.array([meta[k][2] for k in hyp_GA.keys()]) # Create gene_ranges list to hold the range of values for each gene in the population - gene_ranges = [(lower_limit[i], upper_limit[i]) for i in range(len(upper_limit))] + gene_ranges = [ + (lower_limit[i], upper_limit[i]) for i in range(len(upper_limit)) + ] # Initialize the population with initial_values or random values initial_values = [] # If resuming evolution from a previous checkpoint if opt.resume_evolve is not None: - assert os.path.isfile(ROOT / opt.resume_evolve), "evolve population path is wrong!" + assert os.path.isfile( + ROOT / opt.resume_evolve + ), "evolve population path is wrong!" with open(ROOT / opt.resume_evolve, errors="ignore") as f: evolve_population = yaml.safe_load(f) for value in evolve_population.values(): @@ -794,7 +1079,9 @@ def main(opt, callbacks=Callbacks()): # If not resuming from a previous checkpoint, generate initial values from .yaml files in opt.evolve_population else: - yaml_files = [f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml")] + yaml_files = [ + f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml") + ] for file_name in yaml_files: with open(os.path.join(opt.evolve_population, file_name)) as yaml_file: value = yaml.safe_load(yaml_file) @@ -803,9 +1090,14 @@ def main(opt, callbacks=Callbacks()): # Generate random values within the search space for the rest of the population if initial_values is None: - population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size)] + population = [ + generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size) + ] elif pop_size > 1: - population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size - len(initial_values))] + population = [ + generate_individual(gene_ranges, len(hyp_GA)) + for _ in range(pop_size - len(initial_values)) + ] for initial_value in initial_values: population = [initial_value] + population @@ -815,14 +1107,19 @@ def main(opt, callbacks=Callbacks()): if generation >= 1: save_dict = {} for i in range(len(population)): - little_dict = {list_keys[j]: float(population[i][j]) for j in range(len(population[i]))} + little_dict = { + list_keys[j]: float(population[i][j]) + for j in range(len(population[i])) + } save_dict[f"gen{str(generation)}number{str(i)}"] = little_dict with open(save_dir / "evolve_population.yaml", "w") as outfile: yaml.dump(save_dict, outfile, default_flow_style=False) # Adaptive elite size - elite_size = min_elite_size + int((max_elite_size - min_elite_size) * (generation / opt.evolve)) + elite_size = min_elite_size + int( + (max_elite_size - min_elite_size) * (generation / opt.evolve) + ) # Evaluate the fitness of each individual in the population fitness_scores = [] for individual in population: @@ -850,16 +1147,25 @@ def main(opt, callbacks=Callbacks()): # Adaptive tournament size tournament_size = max( max(2, tournament_size_min), - int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10))), + int( + min(tournament_size_max, pop_size) + - (generation / (opt.evolve / 10)) + ), ) # Perform tournament selection to choose the best individual tournament_indices = random.sample(range(pop_size), tournament_size) tournament_fitness = [fitness_scores[j] for j in tournament_indices] - winner_index = tournament_indices[tournament_fitness.index(max(tournament_fitness))] + winner_index = tournament_indices[ + tournament_fitness.index(max(tournament_fitness)) + ] selected_indices.append(winner_index) # Add the elite individuals to the selected indices - elite_indices = [i for i in range(pop_size) if fitness_scores[i] in sorted(fitness_scores)[-elite_size:]] + elite_indices = [ + i + for i in range(pop_size) + if fitness_scores[i] in sorted(fitness_scores)[-elite_size:] + ] selected_indices.extend(elite_indices) # Create the next generation through crossover and mutation next_generation = [] @@ -868,21 +1174,33 @@ def main(opt, callbacks=Callbacks()): parent2_index = selected_indices[random.randint(0, pop_size - 1)] # Adaptive crossover rate crossover_rate = max( - crossover_rate_min, min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve)) + crossover_rate_min, + min( + crossover_rate_max, + crossover_rate_max - (generation / opt.evolve), + ), ) if random.uniform(0, 1) < crossover_rate: crossover_point = random.randint(1, len(hyp_GA) - 1) - child = population[parent1_index][:crossover_point] + population[parent2_index][crossover_point:] + child = ( + population[parent1_index][:crossover_point] + + population[parent2_index][crossover_point:] + ) else: child = population[parent1_index] # Adaptive mutation rate mutation_rate = max( - mutation_rate_min, min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve)) + mutation_rate_min, + min( + mutation_rate_max, mutation_rate_max - (generation / opt.evolve) + ), ) for j in range(len(hyp_GA)): if random.uniform(0, 1) < mutation_rate: child[j] += random.uniform(-0.1, 0.1) - child[j] = min(max(child[j], gene_ranges[j][0]), gene_ranges[j][1]) + child[j] = min( + max(child[j], gene_ranges[j][0]), gene_ranges[j][1] + ) next_generation.append(child) # Replace the old population with the new generation population = next_generation @@ -893,9 +1211,9 @@ def main(opt, callbacks=Callbacks()): # Plot results plot_evolve(evolve_csv) LOGGER.info( - f'Hyperparameter evolution finished {opt.evolve} generations\n' + f"Hyperparameter evolution finished {opt.evolve} generations\n" f"Results saved to {colorstr('bold', save_dir)}\n" - f'Usage example: $ python train.py --hyp {evolve_yaml}' + f"Usage example: $ python train.py --hyp {evolve_yaml}" ) diff --git a/utils/autobatch.py b/utils/autobatch.py index c8deb940fdf1..850e9d07b419 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -12,9 +12,11 @@ def check_train_batch_size(model, imgsz=640, amp=True): """Checks and computes optimal training batch size for YOLOv5 model, given image size and AMP setting.""" - if check_version(torch.__version__, "2.4.0", "Torch", hard=False): + if check_version(torch.__version__, "2.4.0"): with torch.amp.autocast("cuda", enabled=amp): - return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size + return autobatch( + deepcopy(model).train(), imgsz + ) # compute optimal batch size with torch.cuda.amp.autocast(amp): return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size @@ -32,10 +34,14 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): LOGGER.info(f"{prefix}Computing optimal batch size for --imgsz {imgsz}") device = next(model.parameters()).device # get model device if device.type == "cpu": - LOGGER.info(f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}") + LOGGER.info( + f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}" + ) return batch_size if torch.backends.cudnn.benchmark: - LOGGER.info(f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}") + LOGGER.info( + f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}" + ) return batch_size # Inspect CUDA memory @@ -46,7 +52,9 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): r = torch.cuda.memory_reserved(device) / gb # GiB reserved a = torch.cuda.memory_allocated(device) / gb # GiB allocated f = t - (r + a) # GiB free - LOGGER.info(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free") + LOGGER.info( + f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free" + ) # Profile batch sizes batch_sizes = [1, 2, 4, 8, 16] @@ -66,8 +74,12 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): b = batch_sizes[max(i - 1, 0)] # select prior safe point if b < 1 or b > 1024: # b outside of safe range b = batch_size - LOGGER.warning(f"{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.") + LOGGER.warning( + f"{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command." + ) fraction = (np.polyval(p, b) + r + a) / t # actual fraction predicted - LOGGER.info(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅") + LOGGER.info( + f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅" + ) return b From 4736f445d7910867c1f8aa8f697656ce2a4d22ea Mon Sep 17 00:00:00 2001 From: UltralyticsAssistant Date: Mon, 5 Aug 2024 22:09:00 +0000 Subject: [PATCH 5/9] Auto-format by https://ultralytics.com/actions --- classify/train.py | 1 - train.py | 274 +++++++++++---------------------------------- utils/autobatch.py | 24 +--- 3 files changed, 72 insertions(+), 227 deletions(-) diff --git a/classify/train.py b/classify/train.py index 75fafc7cff01..596f85a8d442 100644 --- a/classify/train.py +++ b/classify/train.py @@ -224,7 +224,6 @@ def lf(x): for i, (images, labels) in pbar: # progress bar images, labels = images.to(device, non_blocking=True), labels.to(device) - amp_autocast = None if check_version(torch.__version__, "2.4.0"): amp_autocast = torch.amp.autocast("cuda", enabled=device.type != "cpu") diff --git a/train.py b/train.py index 3b62237495bd..6b83c8ba6126 100644 --- a/train.py +++ b/train.py @@ -95,9 +95,7 @@ torch_distributed_zero_first, ) -LOCAL_RANK = int( - os.getenv("LOCAL_RANK", -1) -) # https://pytorch.org/docs/stable/elastic/run.html +LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1)) # https://pytorch.org/docs/stable/elastic/run.html RANK = int(os.getenv("RANK", -1)) WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1)) GIT_INFO = check_git_info() @@ -177,9 +175,7 @@ def train(hyp, opt, device, callbacks): if isinstance(hyp, str): with open(hyp, errors="ignore") as f: hyp = yaml.safe_load(f) # load hyps dict - LOGGER.info( - colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items()) - ) + LOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items())) opt.hyp = hyp.copy() # for saving hyps to checkpoints # Save run settings @@ -227,14 +223,8 @@ def train(hyp, opt, device, callbacks): data_dict = data_dict or check_dataset(data) # check if None train_path, val_path = data_dict["train"], data_dict["val"] nc = 1 if single_cls else int(data_dict["nc"]) # number of classes - names = ( - {0: "item"} - if single_cls and len(data_dict["names"]) != 1 - else data_dict["names"] - ) # class names - is_coco = isinstance(val_path, str) and val_path.endswith( - "coco/val2017.txt" - ) # COCO dataset + names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"] # class names + is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt") # COCO dataset # Model check_suffix(weights, ".pt") # check weights @@ -242,31 +232,19 @@ def train(hyp, opt, device, callbacks): if pretrained: with torch_distributed_zero_first(LOCAL_RANK): weights = attempt_download(weights) # download if not found locally - ckpt = torch.load( - weights, map_location="cpu" - ) # load checkpoint to CPU to avoid CUDA memory leak - model = Model( - cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors") - ).to( - device - ) # create - exclude = ( - ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] - ) # exclude keys + ckpt = torch.load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak + model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create + exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] # exclude keys csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32 csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(csd, strict=False) # load - LOGGER.info( - f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}" - ) # report + LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}") # report else: model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create amp = check_amp(model) # check AMP # Freeze - freeze = [ - f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0])) - ] # layers to freeze + freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze for k, v in model.named_parameters(): v.requires_grad = True # train all layers # v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results) @@ -287,9 +265,7 @@ def train(hyp, opt, device, callbacks): nbs = 64 # nominal batch size accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing hyp["weight_decay"] *= batch_size * accumulate / nbs # scale weight_decay - optimizer = smart_optimizer( - model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"] - ) + optimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"]) # Scheduler if opt.cos_lr: @@ -300,9 +276,7 @@ def lf(x): """Linear learning rate scheduler function with decay calculated by epoch proportion.""" return (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"] # linear - scheduler = lr_scheduler.LambdaLR( - optimizer, lr_lambda=lf - ) # plot_lr_scheduler(optimizer, scheduler, epochs) + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs) # EMA ema = ModelEMA(model) if RANK in {-1, 0} else None @@ -311,9 +285,7 @@ def lf(x): best_fitness, start_epoch = 0.0, 0 if pretrained: if resume: - best_fitness, start_epoch, epochs = smart_resume( - ckpt, optimizer, ema, weights, epochs, resume - ) + best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume) del ckpt, csd # DP mode @@ -350,9 +322,7 @@ def lf(x): ) labels = np.concatenate(dataset.labels, 0) mlc = int(labels[:, 0].max()) # max label class - assert ( - mlc < nc - ), f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}" + assert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}" # Process 0 if RANK in {-1, 0}: @@ -373,9 +343,7 @@ def lf(x): if not resume: if not opt.noautoanchor: - check_anchors( - dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz - ) # run AutoAnchor + check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz) # run AutoAnchor model.half().float() # pre-reduce anchor precision callbacks.run("on_pretrain_routine_end", labels, names) @@ -392,17 +360,13 @@ def lf(x): hyp["label_smoothing"] = opt.label_smoothing model.nc = nc # attach number of classes to model model.hyp = hyp # attach hyperparameters to model - model.class_weights = ( - labels_to_class_weights(dataset.labels, nc).to(device) * nc - ) # attach class weights + model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights model.names = names # Start training t0 = time.time() nb = len(train_loader) # number of batches - nw = max( - round(hyp["warmup_epochs"] * nb), 100 - ) # number of warmup iterations, max(3 epochs, 100 iterations) + nw = max(round(hyp["warmup_epochs"] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations) # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training last_opt_step = -1 maps = np.zeros(nc) # mAP per class @@ -424,23 +388,15 @@ def lf(x): f"Logging results to {colorstr('bold', save_dir)}\n" f"Starting training for {epochs} epochs..." ) - for epoch in range( - start_epoch, epochs - ): # epoch ------------------------------------------------------------------ + for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ callbacks.run("on_train_epoch_start") model.train() # Update image weights (optional, single-GPU only) if opt.image_weights: - cw = ( - model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc - ) # class weights - iw = labels_to_image_weights( - dataset.labels, nc=nc, class_weights=cw - ) # image weights - dataset.indices = random.choices( - range(dataset.n), weights=iw, k=dataset.n - ) # rand weighted idx + cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights + iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights + dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx # Update mosaic border (optional) # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) @@ -470,14 +426,10 @@ def lf(x): targets, paths, _, - ) in ( - pbar - ): # batch ------------------------------------------------------------- + ) in pbar: # batch ------------------------------------------------------------- callbacks.run("on_train_batch_start") ni = i + nb * epoch # number integrated batches (since train start) - imgs = ( - imgs.to(device, non_blocking=True).float() / 255 - ) # uint8 to float32, 0-255 to 0.0-1.0 + imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0 # Warmup if ni <= nw: @@ -495,23 +447,15 @@ def lf(x): ], ) if "momentum" in x: - x["momentum"] = np.interp( - ni, xi, [hyp["warmup_momentum"], hyp["momentum"]] - ) + x["momentum"] = np.interp(ni, xi, [hyp["warmup_momentum"], hyp["momentum"]]) # Multi-scale if opt.multi_scale: - sz = ( - random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs - ) # size + sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs # size sf = sz / max(imgs.shape[2:]) # scale factor if sf != 1: - ns = [ - math.ceil(x * sf / gs) * gs for x in imgs.shape[2:] - ] # new shape (stretched to gs-multiple) - imgs = nn.functional.interpolate( - imgs, size=ns, mode="bilinear", align_corners=False - ) + ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) + imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) amp_autocast = None if check_version(torch.__version__, "2.4.0"): @@ -522,9 +466,7 @@ def lf(x): # Forward with amp_autocast: pred = model(imgs) # forward - loss, loss_items = compute_loss( - pred, targets.to(device) - ) # loss scaled by batch_size + loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size if RANK != -1: loss *= WORLD_SIZE # gradient averaged between devices in DDP mode if opt.quad: @@ -536,9 +478,7 @@ def lf(x): # Optimize - https://pytorch.org/docs/master/notes/amp_examples.html if ni - last_opt_step >= accumulate: scaler.unscale_(optimizer) # unscale gradients - torch.nn.utils.clip_grad_norm_( - model.parameters(), max_norm=10.0 - ) # clip gradients + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients scaler.step(optimizer) # optimizer.step scaler.update() optimizer.zero_grad() @@ -560,9 +500,7 @@ def lf(x): imgs.shape[-1], ) ) - callbacks.run( - "on_train_batch_end", model, ni, imgs, targets, paths, list(mloss) - ) + callbacks.run("on_train_batch_end", model, ni, imgs, targets, paths, list(mloss)) if callbacks.stop_training: return # end batch ------------------------------------------------------------------------------------------------ @@ -574,9 +512,7 @@ def lf(x): if RANK in {-1, 0}: # mAP callbacks.run("on_train_epoch_end", epoch=epoch) - ema.update_attr( - model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"] - ) + ema.update_attr(model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"]) final_epoch = (epoch + 1 == epochs) or stopper.possible_stop if not noval or final_epoch: # Calculate mAP results, maps, _ = validate.run( @@ -594,9 +530,7 @@ def lf(x): ) # Update best mAP - fi = fitness( - np.array(results).reshape(1, -1) - ) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] + fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] stop = stopper(epoch=epoch, fitness=fi) # early stop check if fi > best_fitness: best_fitness = fi @@ -624,16 +558,12 @@ def lf(x): if opt.save_period > 0 and epoch % opt.save_period == 0: torch.save(ckpt, w / f"epoch{epoch}.pt") del ckpt - callbacks.run( - "on_model_save", last, epoch, final_epoch, best_fitness, fi - ) + callbacks.run("on_model_save", last, epoch, final_epoch, best_fitness, fi) # EarlyStopping if RANK != -1: # if DDP training broadcast_list = [stop if RANK == 0 else None] - dist.broadcast_object_list( - broadcast_list, 0 - ) # broadcast 'stop' to all ranks + dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks if RANK != 0: stop = broadcast_list[0] if stop: @@ -642,9 +572,7 @@ def lf(x): # end epoch ---------------------------------------------------------------------------------------------------- # end training ----------------------------------------------------------------------------------------------------- if RANK in {-1, 0}: - LOGGER.info( - f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours." - ) + LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.") for f in last, best: if f.exists(): strip_optimizer(f) # strip optimizers @@ -655,9 +583,7 @@ def lf(x): batch_size=batch_size // WORLD_SIZE * 2, imgsz=imgsz, model=attempt_load(f, device).half(), - iou_thres=( - 0.65 if is_coco else 0.60 - ), # best pycocotools at iou 0.65 + iou_thres=(0.65 if is_coco else 0.60), # best pycocotools at iou 0.65 single_cls=single_cls, dataloader=val_loader, save_dir=save_dir, @@ -705,13 +631,9 @@ def parse_opt(known=False): - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data """ parser = argparse.ArgumentParser() - parser.add_argument( - "--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path" - ) + parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path") parser.add_argument("--cfg", type=str, default="", help="model.yaml path") - parser.add_argument( - "--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path" - ) + parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") parser.add_argument( "--hyp", type=str, @@ -741,15 +663,9 @@ def parse_opt(known=False): default=False, help="resume most recent training", ) - parser.add_argument( - "--nosave", action="store_true", help="only save final checkpoint" - ) - parser.add_argument( - "--noval", action="store_true", help="only validate final epoch" - ) - parser.add_argument( - "--noautoanchor", action="store_true", help="disable AutoAnchor" - ) + parser.add_argument("--nosave", action="store_true", help="only save final checkpoint") + parser.add_argument("--noval", action="store_true", help="only validate final epoch") + parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor") parser.add_argument("--noplots", action="store_true", help="save no plot files") parser.add_argument( "--evolve", @@ -771,20 +687,14 @@ def parse_opt(known=False): help="resume evolve from last generation", ) parser.add_argument("--bucket", type=str, default="", help="gsutil bucket") - parser.add_argument( - "--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk" - ) + parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk") parser.add_argument( "--image-weights", action="store_true", help="use weighted image selection for training", ) - parser.add_argument( - "--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu" - ) - parser.add_argument( - "--multi-scale", action="store_true", help="vary img-size +/- 50%%" - ) + parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") + parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%") parser.add_argument( "--single-cls", action="store_true", @@ -808,9 +718,7 @@ def parse_opt(known=False): default=8, help="max dataloader workers (per RANK in DDP mode)", ) - parser.add_argument( - "--project", default=ROOT / "runs/train", help="save to project/name" - ) + parser.add_argument("--project", default=ROOT / "runs/train", help="save to project/name") parser.add_argument("--name", default="exp", help="save to project/name") parser.add_argument( "--exist-ok", @@ -819,9 +727,7 @@ def parse_opt(known=False): ) parser.add_argument("--quad", action="store_true", help="quad dataloader") parser.add_argument("--cos-lr", action="store_true", help="cosine LR scheduler") - parser.add_argument( - "--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon" - ) + parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon") parser.add_argument( "--patience", type=int, @@ -872,9 +778,7 @@ def parse_opt(known=False): ) # NDJSON logging - parser.add_argument( - "--ndjson-console", action="store_true", help="Log ndjson to console" - ) + parser.add_argument("--ndjson-console", action="store_true", help="Log ndjson to console") parser.add_argument("--ndjson-file", action="store_true", help="Log ndjson to file") return parser.parse_known_args()[0] if known else parser.parse_args() @@ -903,9 +807,7 @@ def main(opt, callbacks=Callbacks()): # Resume (from specified or most recent last.pt) if opt.resume and not check_comet_resume(opt) and not opt.evolve: - last = Path( - check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run() - ) + last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run()) opt_yaml = last.parent.parent / "opt.yaml" # train options yaml opt_data = opt.data # original dataset if opt_yaml.is_file(): @@ -925,13 +827,9 @@ def main(opt, callbacks=Callbacks()): str(opt.weights), str(opt.project), ) # checks - assert len(opt.cfg) or len( - opt.weights - ), "either --cfg or --weights must be specified" + assert len(opt.cfg) or len(opt.weights), "either --cfg or --weights must be specified" if opt.evolve: - if opt.project == str( - ROOT / "runs/train" - ): # if default project name, rename to runs/evolve + if opt.project == str(ROOT / "runs/train"): # if default project name, rename to runs/evolve opt.project = str(ROOT / "runs/evolve") opt.exist_ok, opt.resume = ( opt.resume, @@ -939,9 +837,7 @@ def main(opt, callbacks=Callbacks()): ) # pass resume to exist_ok and disable resume if opt.name == "cfg": opt.name = Path(opt.cfg).stem # use model.yaml as name - opt.save_dir = str( - increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) - ) + opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # DDP mode device = select_device(opt.device, batch_size=opt.batch_size) @@ -949,15 +845,9 @@ def main(opt, callbacks=Callbacks()): msg = "is not compatible with YOLOv5 Multi-GPU DDP training" assert not opt.image_weights, f"--image-weights {msg}" assert not opt.evolve, f"--evolve {msg}" - assert ( - opt.batch_size != -1 - ), f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size" - assert ( - opt.batch_size % WORLD_SIZE == 0 - ), f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE" - assert ( - torch.cuda.device_count() > LOCAL_RANK - ), "insufficient CUDA devices for DDP command" + assert opt.batch_size != -1, f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size" + assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE" + assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command" torch.cuda.set_device(LOCAL_RANK) device = torch.device("cuda", LOCAL_RANK) dist.init_process_group( @@ -1059,18 +949,14 @@ def main(opt, callbacks=Callbacks()): upper_limit = np.array([meta[k][2] for k in hyp_GA.keys()]) # Create gene_ranges list to hold the range of values for each gene in the population - gene_ranges = [ - (lower_limit[i], upper_limit[i]) for i in range(len(upper_limit)) - ] + gene_ranges = [(lower_limit[i], upper_limit[i]) for i in range(len(upper_limit))] # Initialize the population with initial_values or random values initial_values = [] # If resuming evolution from a previous checkpoint if opt.resume_evolve is not None: - assert os.path.isfile( - ROOT / opt.resume_evolve - ), "evolve population path is wrong!" + assert os.path.isfile(ROOT / opt.resume_evolve), "evolve population path is wrong!" with open(ROOT / opt.resume_evolve, errors="ignore") as f: evolve_population = yaml.safe_load(f) for value in evolve_population.values(): @@ -1079,9 +965,7 @@ def main(opt, callbacks=Callbacks()): # If not resuming from a previous checkpoint, generate initial values from .yaml files in opt.evolve_population else: - yaml_files = [ - f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml") - ] + yaml_files = [f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml")] for file_name in yaml_files: with open(os.path.join(opt.evolve_population, file_name)) as yaml_file: value = yaml.safe_load(yaml_file) @@ -1090,14 +974,9 @@ def main(opt, callbacks=Callbacks()): # Generate random values within the search space for the rest of the population if initial_values is None: - population = [ - generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size) - ] + population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size)] elif pop_size > 1: - population = [ - generate_individual(gene_ranges, len(hyp_GA)) - for _ in range(pop_size - len(initial_values)) - ] + population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size - len(initial_values))] for initial_value in initial_values: population = [initial_value] + population @@ -1107,19 +986,14 @@ def main(opt, callbacks=Callbacks()): if generation >= 1: save_dict = {} for i in range(len(population)): - little_dict = { - list_keys[j]: float(population[i][j]) - for j in range(len(population[i])) - } + little_dict = {list_keys[j]: float(population[i][j]) for j in range(len(population[i]))} save_dict[f"gen{str(generation)}number{str(i)}"] = little_dict with open(save_dir / "evolve_population.yaml", "w") as outfile: yaml.dump(save_dict, outfile, default_flow_style=False) # Adaptive elite size - elite_size = min_elite_size + int( - (max_elite_size - min_elite_size) * (generation / opt.evolve) - ) + elite_size = min_elite_size + int((max_elite_size - min_elite_size) * (generation / opt.evolve)) # Evaluate the fitness of each individual in the population fitness_scores = [] for individual in population: @@ -1147,25 +1021,16 @@ def main(opt, callbacks=Callbacks()): # Adaptive tournament size tournament_size = max( max(2, tournament_size_min), - int( - min(tournament_size_max, pop_size) - - (generation / (opt.evolve / 10)) - ), + int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10))), ) # Perform tournament selection to choose the best individual tournament_indices = random.sample(range(pop_size), tournament_size) tournament_fitness = [fitness_scores[j] for j in tournament_indices] - winner_index = tournament_indices[ - tournament_fitness.index(max(tournament_fitness)) - ] + winner_index = tournament_indices[tournament_fitness.index(max(tournament_fitness))] selected_indices.append(winner_index) # Add the elite individuals to the selected indices - elite_indices = [ - i - for i in range(pop_size) - if fitness_scores[i] in sorted(fitness_scores)[-elite_size:] - ] + elite_indices = [i for i in range(pop_size) if fitness_scores[i] in sorted(fitness_scores)[-elite_size:]] selected_indices.extend(elite_indices) # Create the next generation through crossover and mutation next_generation = [] @@ -1182,25 +1047,18 @@ def main(opt, callbacks=Callbacks()): ) if random.uniform(0, 1) < crossover_rate: crossover_point = random.randint(1, len(hyp_GA) - 1) - child = ( - population[parent1_index][:crossover_point] - + population[parent2_index][crossover_point:] - ) + child = population[parent1_index][:crossover_point] + population[parent2_index][crossover_point:] else: child = population[parent1_index] # Adaptive mutation rate mutation_rate = max( mutation_rate_min, - min( - mutation_rate_max, mutation_rate_max - (generation / opt.evolve) - ), + min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve)), ) for j in range(len(hyp_GA)): if random.uniform(0, 1) < mutation_rate: child[j] += random.uniform(-0.1, 0.1) - child[j] = min( - max(child[j], gene_ranges[j][0]), gene_ranges[j][1] - ) + child[j] = min(max(child[j], gene_ranges[j][0]), gene_ranges[j][1]) next_generation.append(child) # Replace the old population with the new generation population = next_generation diff --git a/utils/autobatch.py b/utils/autobatch.py index 850e9d07b419..201756e860c6 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -14,9 +14,7 @@ def check_train_batch_size(model, imgsz=640, amp=True): """Checks and computes optimal training batch size for YOLOv5 model, given image size and AMP setting.""" if check_version(torch.__version__, "2.4.0"): with torch.amp.autocast("cuda", enabled=amp): - return autobatch( - deepcopy(model).train(), imgsz - ) # compute optimal batch size + return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size with torch.cuda.amp.autocast(amp): return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size @@ -34,14 +32,10 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): LOGGER.info(f"{prefix}Computing optimal batch size for --imgsz {imgsz}") device = next(model.parameters()).device # get model device if device.type == "cpu": - LOGGER.info( - f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}" - ) + LOGGER.info(f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}") return batch_size if torch.backends.cudnn.benchmark: - LOGGER.info( - f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}" - ) + LOGGER.info(f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}") return batch_size # Inspect CUDA memory @@ -52,9 +46,7 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): r = torch.cuda.memory_reserved(device) / gb # GiB reserved a = torch.cuda.memory_allocated(device) / gb # GiB allocated f = t - (r + a) # GiB free - LOGGER.info( - f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free" - ) + LOGGER.info(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free") # Profile batch sizes batch_sizes = [1, 2, 4, 8, 16] @@ -74,12 +66,8 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): b = batch_sizes[max(i - 1, 0)] # select prior safe point if b < 1 or b > 1024: # b outside of safe range b = batch_size - LOGGER.warning( - f"{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command." - ) + LOGGER.warning(f"{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.") fraction = (np.polyval(p, b) + r + a) / t # actual fraction predicted - LOGGER.info( - f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅" - ) + LOGGER.info(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅") return b From 24ab9de10c97e3427965cde7e1bdfba917307910 Mon Sep 17 00:00:00 2001 From: Jacob Brown Date: Mon, 5 Aug 2024 16:12:49 -0600 Subject: [PATCH 6/9] removed formatting changes --- train.py | 258 +++++++++---------------------------------------------- 1 file changed, 41 insertions(+), 217 deletions(-) diff --git a/train.py b/train.py index 6b83c8ba6126..bc736c67c65c 100644 --- a/train.py +++ b/train.py @@ -135,21 +135,7 @@ def train(hyp, opt, device, callbacks): - Datasets: https://github.com/ultralytics/yolov5/tree/master/data - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data """ - ( - save_dir, - epochs, - batch_size, - weights, - single_cls, - evolve, - data, - cfg, - resume, - noval, - nosave, - workers, - freeze, - ) = ( + save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = ( Path(opt.save_dir), opt.epochs, opt.batch_size, @@ -208,12 +194,7 @@ def train(hyp, opt, device, callbacks): # Process custom dataset artifact link data_dict = loggers.remote_dataset if resume: # If resuming runs from remote artifact - weights, epochs, hyp, batch_size = ( - opt.weights, - opt.epochs, - opt.hyp, - opt.batch_size, - ) + weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size # Config plots = not evolve and not opt.noplots # create plots @@ -383,10 +364,10 @@ def lf(x): compute_loss = ComputeLoss(model) # init loss class callbacks.run("on_train_start") LOGGER.info( - f"Image sizes {imgsz} train, {imgsz} val\n" - f"Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n" + f'Image sizes {imgsz} train, {imgsz} val\n' + f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n' f"Logging results to {colorstr('bold', save_dir)}\n" - f"Starting training for {epochs} epochs..." + f'Starting training for {epochs} epochs...' ) for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ callbacks.run("on_train_epoch_start") @@ -406,27 +387,11 @@ def lf(x): if RANK != -1: train_loader.sampler.set_epoch(epoch) pbar = enumerate(train_loader) - LOGGER.info( - ("\n" + "%11s" * 7) - % ( - "Epoch", - "GPU_mem", - "box_loss", - "obj_loss", - "cls_loss", - "Instances", - "Size", - ) - ) + LOGGER.info(("\n" + "%11s" * 7) % ("Epoch", "GPU_mem", "box_loss", "obj_loss", "cls_loss", "Instances", "Size")) if RANK in {-1, 0}: pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # progress bar optimizer.zero_grad() - for i, ( - imgs, - targets, - paths, - _, - ) in pbar: # batch ------------------------------------------------------------- + for i, (imgs, targets, paths, _) in pbar: # batch ------------------------------------------------------------- callbacks.run("on_train_batch_start") ni = i + nb * epoch # number integrated batches (since train start) imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0 @@ -438,14 +403,7 @@ def lf(x): accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round()) for j, x in enumerate(optimizer.param_groups): # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0 - x["lr"] = np.interp( - ni, - xi, - [ - hyp["warmup_bias_lr"] if j == 0 else 0.0, - x["initial_lr"] * lf(epoch), - ], - ) + x["lr"] = np.interp(ni, xi, [hyp["warmup_bias_lr"] if j == 0 else 0.0, x["initial_lr"] * lf(epoch)]) if "momentum" in x: x["momentum"] = np.interp(ni, xi, [hyp["warmup_momentum"], hyp["momentum"]]) @@ -492,13 +450,7 @@ def lf(x): mem = f"{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G" # (GB) pbar.set_description( ("%11s" * 2 + "%11.4g" * 5) - % ( - f"{epoch}/{epochs - 1}", - mem, - *mloss, - targets.shape[0], - imgs.shape[-1], - ) + % (f"{epoch}/{epochs - 1}", mem, *mloss, targets.shape[0], imgs.shape[-1]) ) callbacks.run("on_train_batch_end", model, ni, imgs, targets, paths, list(mloss)) if callbacks.stop_training: @@ -583,7 +535,7 @@ def lf(x): batch_size=batch_size // WORLD_SIZE * 2, imgsz=imgsz, model=attempt_load(f, device).half(), - iou_thres=(0.65 if is_coco else 0.60), # best pycocotools at iou 0.65 + iou_thres=0.65 if is_coco else 0.60, # best pycocotools at iou 0.65 single_cls=single_cls, dataloader=val_loader, save_dir=save_dir, @@ -594,13 +546,7 @@ def lf(x): compute_loss=compute_loss, ) # val best model with plots if is_coco: - callbacks.run( - "on_fit_epoch_end", - list(mloss) + list(results) + lr, - epoch, - best_fitness, - fi, - ) + callbacks.run("on_fit_epoch_end", list(mloss) + list(results) + lr, epoch, best_fitness, fi) callbacks.run("on_train_end", last, best, epoch, results) @@ -634,148 +580,47 @@ def parse_opt(known=False): parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path") parser.add_argument("--cfg", type=str, default="", help="model.yaml path") parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") - parser.add_argument( - "--hyp", - type=str, - default=ROOT / "data/hyps/hyp.scratch-low.yaml", - help="hyperparameters path", - ) + parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path") parser.add_argument("--epochs", type=int, default=100, help="total training epochs") - parser.add_argument( - "--batch-size", - type=int, - default=16, - help="total batch size for all GPUs, -1 for autobatch", - ) - parser.add_argument( - "--imgsz", - "--img", - "--img-size", - type=int, - default=640, - help="train, val image size (pixels)", - ) + parser.add_argument("--batch-size", type=int, default=16, help="total batch size for all GPUs, -1 for autobatch") + parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)") parser.add_argument("--rect", action="store_true", help="rectangular training") - parser.add_argument( - "--resume", - nargs="?", - const=True, - default=False, - help="resume most recent training", - ) + parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training") parser.add_argument("--nosave", action="store_true", help="only save final checkpoint") parser.add_argument("--noval", action="store_true", help="only validate final epoch") parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor") parser.add_argument("--noplots", action="store_true", help="save no plot files") + parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations") parser.add_argument( - "--evolve", - type=int, - nargs="?", - const=300, - help="evolve hyperparameters for x generations", - ) - parser.add_argument( - "--evolve_population", - type=str, - default=ROOT / "data/hyps", - help="location for loading population", - ) - parser.add_argument( - "--resume_evolve", - type=str, - default=None, - help="resume evolve from last generation", + "--evolve_population", type=str, default=ROOT / "data/hyps", help="location for loading population" ) + parser.add_argument("--resume_evolve", type=str, default=None, help="resume evolve from last generation") parser.add_argument("--bucket", type=str, default="", help="gsutil bucket") parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/disk") - parser.add_argument( - "--image-weights", - action="store_true", - help="use weighted image selection for training", - ) + parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training") parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%") - parser.add_argument( - "--single-cls", - action="store_true", - help="train multi-class data as single-class", - ) - parser.add_argument( - "--optimizer", - type=str, - choices=["SGD", "Adam", "AdamW"], - default="SGD", - help="optimizer", - ) - parser.add_argument( - "--sync-bn", - action="store_true", - help="use SyncBatchNorm, only available in DDP mode", - ) - parser.add_argument( - "--workers", - type=int, - default=8, - help="max dataloader workers (per RANK in DDP mode)", - ) + parser.add_argument("--single-cls", action="store_true", help="train multi-class data as single-class") + parser.add_argument("--optimizer", type=str, choices=["SGD", "Adam", "AdamW"], default="SGD", help="optimizer") + parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode") + parser.add_argument("--workers", type=int, default=8, help="max dataloader workers (per RANK in DDP mode)") parser.add_argument("--project", default=ROOT / "runs/train", help="save to project/name") parser.add_argument("--name", default="exp", help="save to project/name") - parser.add_argument( - "--exist-ok", - action="store_true", - help="existing project/name ok, do not increment", - ) + parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment") parser.add_argument("--quad", action="store_true", help="quad dataloader") parser.add_argument("--cos-lr", action="store_true", help="cosine LR scheduler") parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon") - parser.add_argument( - "--patience", - type=int, - default=100, - help="EarlyStopping patience (epochs without improvement)", - ) - parser.add_argument( - "--freeze", - nargs="+", - type=int, - default=[0], - help="Freeze layers: backbone=10, first3=0 1 2", - ) - parser.add_argument( - "--save-period", - type=int, - default=-1, - help="Save checkpoint every x epochs (disabled if < 1)", - ) + parser.add_argument("--patience", type=int, default=100, help="EarlyStopping patience (epochs without improvement)") + parser.add_argument("--freeze", nargs="+", type=int, default=[0], help="Freeze layers: backbone=10, first3=0 1 2") + parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)") parser.add_argument("--seed", type=int, default=0, help="Global training seed") - parser.add_argument( - "--local_rank", - type=int, - default=-1, - help="Automatic DDP Multi-GPU argument, do not modify", - ) + parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify") # Logger arguments parser.add_argument("--entity", default=None, help="Entity") - parser.add_argument( - "--upload_dataset", - nargs="?", - const=True, - default=False, - help='Upload data, "val" option', - ) - parser.add_argument( - "--bbox_interval", - type=int, - default=-1, - help="Set bounding-box image logging interval", - ) - parser.add_argument( - "--artifact_alias", - type=str, - default="latest", - help="Version of dataset artifact to use", - ) + parser.add_argument("--upload_dataset", nargs="?", const=True, default=False, help='Upload data, "val" option') + parser.add_argument("--bbox_interval", type=int, default=-1, help="Set bounding-box image logging interval") + parser.add_argument("--artifact_alias", type=str, default="latest", help="Version of dataset artifact to use") # NDJSON logging parser.add_argument("--ndjson-console", action="store_true", help="Log ndjson to console") @@ -831,10 +676,7 @@ def main(opt, callbacks=Callbacks()): if opt.evolve: if opt.project == str(ROOT / "runs/train"): # if default project name, rename to runs/evolve opt.project = str(ROOT / "runs/evolve") - opt.exist_ok, opt.resume = ( - opt.resume, - False, - ) # pass resume to exist_ok and disable resume + opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume if opt.name == "cfg": opt.name = Path(opt.cfg).stem # use model.yaml as name opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) @@ -851,8 +693,7 @@ def main(opt, callbacks=Callbacks()): torch.cuda.set_device(LOCAL_RANK) device = torch.device("cuda", LOCAL_RANK) dist.init_process_group( - backend="nccl" if dist.is_nccl_available() else "gloo", - timeout=timedelta(seconds=10800), + backend="nccl" if dist.is_nccl_available() else "gloo", timeout=timedelta(seconds=10800) ) # Train @@ -878,11 +719,7 @@ def main(opt, callbacks=Callbacks()): "iou_t": (False, 0.1, 0.7), # IoU training threshold "anchor_t": (False, 2.0, 8.0), # anchor-multiple threshold "anchors": (False, 2.0, 10.0), # anchors per output grid (0 to ignore) - "fl_gamma": ( - False, - 0.0, - 2.0, - ), # focal loss gamma (efficientDet default gamma=1.5) + "fl_gamma": (False, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) "hsv_h": (True, 0.0, 0.1), # image HSV-Hue augmentation (fraction) "hsv_s": (True, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) "hsv_v": (True, 0.0, 0.9), # image HSV-Value augmentation (fraction) @@ -890,11 +727,7 @@ def main(opt, callbacks=Callbacks()): "translate": (True, 0.0, 0.9), # image translation (+/- fraction) "scale": (True, 0.0, 0.9), # image scale (+/- gain) "shear": (True, 0.0, 10.0), # image shear (+/- deg) - "perspective": ( - True, - 0.0, - 0.001, - ), # image perspective (+/- fraction), range 0-0.001 + "perspective": (True, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 "flipud": (True, 0.0, 1.0), # image flip up-down (probability) "fliplr": (True, 0.0, 1.0), # image flip left-right (probability) "mosaic": (True, 0.0, 1.0), # image mixup (probability) @@ -919,11 +752,7 @@ def main(opt, callbacks=Callbacks()): hyp["anchors"] = 3 if opt.noautoanchor: del hyp["anchors"], meta["anchors"] - opt.noval, opt.nosave, save_dir = ( - True, - True, - Path(opt.save_dir), - ) # only val/save final epoch + opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices evolve_yaml, evolve_csv = save_dir / "hyp_evolve.yaml", save_dir / "evolve.csv" if opt.bucket: @@ -1039,11 +868,7 @@ def main(opt, callbacks=Callbacks()): parent2_index = selected_indices[random.randint(0, pop_size - 1)] # Adaptive crossover rate crossover_rate = max( - crossover_rate_min, - min( - crossover_rate_max, - crossover_rate_max - (generation / opt.evolve), - ), + crossover_rate_min, min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve)) ) if random.uniform(0, 1) < crossover_rate: crossover_point = random.randint(1, len(hyp_GA) - 1) @@ -1052,8 +877,7 @@ def main(opt, callbacks=Callbacks()): child = population[parent1_index] # Adaptive mutation rate mutation_rate = max( - mutation_rate_min, - min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve)), + mutation_rate_min, min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve)) ) for j in range(len(hyp_GA)): if random.uniform(0, 1) < mutation_rate: @@ -1069,9 +893,9 @@ def main(opt, callbacks=Callbacks()): # Plot results plot_evolve(evolve_csv) LOGGER.info( - f"Hyperparameter evolution finished {opt.evolve} generations\n" + f'Hyperparameter evolution finished {opt.evolve} generations\n' f"Results saved to {colorstr('bold', save_dir)}\n" - f"Usage example: $ python train.py --hyp {evolve_yaml}" + f'Usage example: $ python train.py --hyp {evolve_yaml}' ) @@ -1172,4 +996,4 @@ def run(**kwargs): if __name__ == "__main__": opt = parse_opt() - main(opt) + main(opt) \ No newline at end of file From 0c6aaa0665267c03c8f526564e661bb6d841d15e Mon Sep 17 00:00:00 2001 From: UltralyticsAssistant Date: Mon, 5 Aug 2024 22:13:12 +0000 Subject: [PATCH 7/9] Auto-format by https://ultralytics.com/actions --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index bc736c67c65c..f5c6a52a95e9 100644 --- a/train.py +++ b/train.py @@ -996,4 +996,4 @@ def run(**kwargs): if __name__ == "__main__": opt = parse_opt() - main(opt) \ No newline at end of file + main(opt) From 3e9b43e8d2c74816b3737e55c0d993f8d172e450 Mon Sep 17 00:00:00 2001 From: Jacob Brown Date: Mon, 5 Aug 2024 16:19:18 -0600 Subject: [PATCH 8/9] fixed a missed amp.autocast --- models/common.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/models/common.py b/models/common.py index d735e10fed8a..4841c09fcf86 100644 --- a/models/common.py +++ b/models/common.py @@ -870,7 +870,13 @@ def forward(self, ims, size=640, augment=False, profile=False): x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) # stack and BHWC to BCHW x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32 - with amp.autocast(autocast): + amp_autocast = None + if check_version(torch.__version__, "2.4.0"): + amp_autocast = torch.amp.autocast("cuda", enabled=autocast) + else: + amp_autocast = torch.cuda.amp.autocast(enabled=autocast) + + with amp_autocast: # Inference with dt[1]: y = self.model(x, augment=augment) # forward From 4d818f6b7af5a33a9f1fac6876a6c9ba13020623 Mon Sep 17 00:00:00 2001 From: UltralyticsAssistant Date: Sat, 24 Aug 2024 21:43:09 +0000 Subject: [PATCH 9/9] Auto-format by https://ultralytics.com/actions --- export.py | 1 + utils/augmentations.py | 1 - utils/callbacks.py | 1 - utils/dataloaders.py | 7 ++++--- utils/general.py | 2 -- utils/loggers/__init__.py | 3 ++- utils/loggers/clearml/clearml_utils.py | 14 +++++++------- utils/loggers/wandb/wandb_utils.py | 12 ++++++------ utils/metrics.py | 8 +++----- utils/segment/augmentations.py | 1 - utils/segment/general.py | 3 --- utils/triton.py | 3 +-- 12 files changed, 24 insertions(+), 32 deletions(-) diff --git a/export.py b/export.py index dfb1c06fb5e2..f3216a564290 100644 --- a/export.py +++ b/export.py @@ -449,6 +449,7 @@ def transform_fn(data_item): Quantization transform function. Extracts and preprocess input data from dataloader item for quantization. + Parameters: data_item: Tuple with data item produced by DataLoader during iteration Returns: diff --git a/utils/augmentations.py b/utils/augmentations.py index 4a6e441d7c45..bdbe07712716 100644 --- a/utils/augmentations.py +++ b/utils/augmentations.py @@ -156,7 +156,6 @@ def random_perspective( ): # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10)) # targets = [cls, xyxy] - """Applies random perspective transformation to an image, modifying the image and corresponding labels.""" height = im.shape[0] + border[0] * 2 # shape(h,w,c) width = im.shape[1] + border[1] * 2 diff --git a/utils/callbacks.py b/utils/callbacks.py index 0a0bcbdb2b96..21c587bd74c6 100644 --- a/utils/callbacks.py +++ b/utils/callbacks.py @@ -64,7 +64,6 @@ def run(self, hook, *args, thread=False, **kwargs): thread: (boolean) Run callbacks in daemon thread kwargs: Keyword Arguments to receive from YOLOv5 """ - assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" for logger in self._callbacks[hook]: if thread: diff --git a/utils/dataloaders.py b/utils/dataloaders.py index 21308f0cedbd..bdeffec465e7 100644 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -1104,7 +1104,8 @@ def extract_boxes(path=DATASETS_DIR / "coco128"): def autosplit(path=DATASETS_DIR / "coco128/images", weights=(0.9, 0.1, 0.0), annotated_only=False): """Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files Usage: from utils.dataloaders import *; autosplit() - Arguments + + Arguments: path: Path to images directory weights: Train, val, test weights (list, tuple) annotated_only: Only use images with an annotated txt file @@ -1183,7 +1184,7 @@ class HUBDatasetStats: """ Class for generating HUB dataset JSON and `-hub` dataset directory. - Arguments + Arguments: path: Path to data.yaml or data.zip (with data.yaml inside data.zip) autodownload: Attempt to download dataset if not found locally @@ -1314,7 +1315,7 @@ class ClassificationDataset(torchvision.datasets.ImageFolder): """ YOLOv5 Classification Dataset. - Arguments + Arguments: root: Dataset path transform: torchvision transforms, used by default album_transform: Albumentations transforms, used if installed diff --git a/utils/general.py b/utils/general.py index e311504b3031..57db68a7ac76 100644 --- a/utils/general.py +++ b/utils/general.py @@ -518,7 +518,6 @@ def check_font(font=FONT, progress=False): def check_dataset(data, autodownload=True): """Validates and/or auto-downloads a dataset, returning its configuration as a dictionary.""" - # Download (optional) extract_dir = "" if isinstance(data, (str, Path)) and (is_zipfile(data) or is_tarfile(data)): @@ -1023,7 +1022,6 @@ def non_max_suppression( Returns: list of detections, on (n,6) tensor per image [xyxy, conf, cls] """ - # Checks assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0" assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0" diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py index 2bd8583d2ade..7051e8da0a29 100644 --- a/utils/loggers/__init__.py +++ b/utils/loggers/__init__.py @@ -350,7 +350,8 @@ class GenericLogger: """ YOLOv5 General purpose logger for non-task specific logging Usage: from utils.loggers import GenericLogger; logger = GenericLogger(...) - Arguments + + Arguments: opt: Run arguments console_logger: Console logger include: loggers to include diff --git a/utils/loggers/clearml/clearml_utils.py b/utils/loggers/clearml/clearml_utils.py index 2b5351ef8533..de4129e08a16 100644 --- a/utils/loggers/clearml/clearml_utils.py +++ b/utils/loggers/clearml/clearml_utils.py @@ -80,7 +80,7 @@ def __init__(self, opt, hyp): - Initialize ClearML Task, this object will capture the experiment - Upload dataset version to ClearML Data if opt.upload_dataset is True - arguments: + Arguments: opt (namespace) -- Commandline arguments for this run hyp (dict) -- Hyperparameters for this run @@ -133,7 +133,7 @@ def log_scalars(self, metrics, epoch): """ Log scalars/metrics to ClearML. - arguments: + Arguments: metrics (dict) Metrics in dict format: {"metrics/mAP": 0.8, ...} epoch (int) iteration number for the current set of metrics """ @@ -145,7 +145,7 @@ def log_model(self, model_path, model_name, epoch=0): """ Log model weights to ClearML. - arguments: + Arguments: model_path (PosixPath or str) Path to the model weights model_name (str) Name of the model visible in ClearML epoch (int) Iteration / epoch of the model weights @@ -158,7 +158,7 @@ def log_summary(self, metrics): """ Log final metrics to a summary table. - arguments: + Arguments: metrics (dict) Metrics in dict format: {"metrics/mAP": 0.8, ...} """ for k, v in metrics.items(): @@ -168,7 +168,7 @@ def log_plot(self, title, plot_path): """ Log image as plot in the plot section of ClearML. - arguments: + Arguments: title (str) Title of the plot plot_path (PosixPath or str) Path to the saved image file """ @@ -183,7 +183,7 @@ def log_debug_samples(self, files, title="Debug Samples"): """ Log files (images) as debug samples in the ClearML task. - arguments: + Arguments: files (List(PosixPath)) a list of file paths in PosixPath format title (str) A title that groups together images with the same values """ @@ -199,7 +199,7 @@ def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_thres """ Draw the bounding boxes on a single image and report the result as a ClearML debug sample. - arguments: + Arguments: image_path (PosixPath) the path the original image file boxes (list): list of scaled predictions in the format - [xmin, ymin, xmax, ymax, confidence, class] class_names (dict): dict containing mapping of class int to class name diff --git a/utils/loggers/wandb/wandb_utils.py b/utils/loggers/wandb/wandb_utils.py index 930f2c7543af..6a32c8cc7b03 100644 --- a/utils/loggers/wandb/wandb_utils.py +++ b/utils/loggers/wandb/wandb_utils.py @@ -49,7 +49,7 @@ def __init__(self, opt, run_id=None, job_type="Training"): - Upload dataset if opt.upload_dataset is True - Setup training processes if job_type is 'Training' - arguments: + Arguments: opt (namespace) -- Commandline arguments for this run run_id (str) -- Run ID of W&B run to be resumed job_type (str) -- To set the job_type for this run @@ -90,7 +90,7 @@ def setup_training(self, opt): - Update data_dict, to contain info of previous run if resumed and the paths of dataset artifact if downloaded - Setup log_dict, initialize bbox_interval - arguments: + Arguments: opt (namespace) -- commandline arguments for this run """ @@ -120,7 +120,7 @@ def log_model(self, path, opt, epoch, fitness_score, best_model=False): """ Log the model checkpoint as W&B artifact. - arguments: + Arguments: path (Path) -- Path of directory containing the checkpoints opt (namespace) -- Command line arguments for this run epoch (int) -- Current epoch number @@ -159,7 +159,7 @@ def log(self, log_dict): """ Save the metrics to the logging dictionary. - arguments: + Arguments: log_dict (Dict) -- metrics/media to be logged in current step """ if self.wandb_run: @@ -170,7 +170,7 @@ def end_epoch(self): """ Commit the log_dict, model artifacts and Tables to W&B and flush the log_dict. - arguments: + Arguments: best_result (boolean): Boolean representing if the result of this evaluation is best or not """ if self.wandb_run: @@ -197,7 +197,7 @@ def finish_run(self): @contextmanager def all_logging_disabled(highest_level=logging.CRITICAL): - """source - https://gist.github.com/simon-weber/7853144 + """Source - https://gist.github.com/simon-weber/7853144 A context manager that will prevent any logging messages triggered during the body from being processed. :param highest_level: the maximum logging level in use. This would only need to be changed if a custom level greater than CRITICAL is defined. diff --git a/utils/metrics.py b/utils/metrics.py index 385fdc471748..9acc38591f96 100644 --- a/utils/metrics.py +++ b/utils/metrics.py @@ -41,7 +41,6 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir=".", names # Returns The average precision as computed in py-faster-rcnn. """ - # Sort by objectness i = np.argsort(-conf) tp, conf, pred_cls = tp[i], conf[i], pred_cls[i] @@ -103,7 +102,6 @@ def compute_ap(recall, precision): # Returns Average precision, precision curve, recall curve """ - # Append sentinel values to beginning and end mrec = np.concatenate(([0.0], recall, [1.0])) mpre = np.concatenate(([1.0], precision, [0.0])) @@ -137,6 +135,7 @@ def process_batch(self, detections, labels): Return intersection-over-union (Jaccard index) of boxes. Both sets of boxes are expected to be in (x1, y1, x2, y2) format. + Arguments: detections (Array[N, 6]), x1, y1, x2, y2, conf, class labels (Array[M, 5]), class, x1, y1, x2, y2 @@ -233,7 +232,6 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7 Input shapes are box1(1,4) to box2(n,4). """ - # Get the coordinates of bounding boxes if xywh: # transform from xywh to xyxy (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1) @@ -279,14 +277,15 @@ def box_iou(box1, box2, eps=1e-7): Return intersection-over-union (Jaccard index) of boxes. Both sets of boxes are expected to be in (x1, y1, x2, y2) format. + Arguments: box1 (Tensor[N, 4]) box2 (Tensor[M, 4]) + Returns: iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2 """ - # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) (a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2) inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2) @@ -304,7 +303,6 @@ def bbox_ioa(box1, box2, eps=1e-7): box2: np.array of shape(nx4) returns: np.array of shape(n) """ - # Get the coordinates of bounding boxes b1_x1, b1_y1, b1_x2, b1_y2 = box1 b2_x1, b2_y1, b2_x2, b2_y2 = box2.T diff --git a/utils/segment/augmentations.py b/utils/segment/augmentations.py index d7dd8aec6691..2e1dca1198b0 100644 --- a/utils/segment/augmentations.py +++ b/utils/segment/augmentations.py @@ -29,7 +29,6 @@ def random_perspective( ): # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10)) # targets = [cls, xyxy] - """Applies random perspective, rotation, scale, shear, and translation augmentations to an image and targets.""" height = im.shape[0] + border[0] * 2 # shape(h,w,c) width = im.shape[1] + border[1] * 2 diff --git a/utils/segment/general.py b/utils/segment/general.py index 2f65d60238dd..0793470a95e4 100644 --- a/utils/segment/general.py +++ b/utils/segment/general.py @@ -14,7 +14,6 @@ def crop_mask(masks, boxes): - masks should be a size [n, h, w] tensor of masks - boxes should be a size [n, 4] tensor of bbox coords in relative point form """ - n, h, w = masks.shape x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(1,1,n) r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :] # rows shape(1,w,1) @@ -33,7 +32,6 @@ def process_mask_upsample(protos, masks_in, bboxes, shape): return: h, w, n """ - c, mh, mw = protos.shape # CHW masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0] # CHW @@ -51,7 +49,6 @@ def process_mask(protos, masks_in, bboxes, shape, upsample=False): return: h, w, n """ - c, mh, mw = protos.shape # CHW ih, iw = shape masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) # CHW diff --git a/utils/triton.py b/utils/triton.py index 3d529ec88a07..2fee42815517 100644 --- a/utils/triton.py +++ b/utils/triton.py @@ -17,10 +17,9 @@ class TritonRemoteModel: def __init__(self, url: str): """ - Keyword arguments: + Keyword Arguments: url: Fully qualified address of the Triton server - for e.g. grpc://localhost:8000 """ - parsed_url = urlparse(url) if parsed_url.scheme == "grpc": from tritonclient.grpc import InferenceServerClient, InferInput