From a63bfd38c18962b29a74a477f7e7d68364f1a983 Mon Sep 17 00:00:00 2001 From: Bala-Vignesh-Reddy Date: Sun, 23 Feb 2025 20:07:59 +0530 Subject: [PATCH] refractored the code --- classify/train.py | 13 ++++++++++--- classify/val.py | 7 ++++++- models/common.py | 10 +++++++--- segment/train.py | 12 ++++++++++-- train.py | 12 ++++++++++-- utils/autobatch.py | 2 +- 6 files changed, 44 insertions(+), 12 deletions(-) diff --git a/classify/train.py b/classify/train.py index d454c7187..93bc700ca 100644 --- a/classify/train.py +++ b/classify/train.py @@ -27,9 +27,16 @@ import torch.distributed as dist import torch.hub as hub import torch.optim.lr_scheduler as lr_scheduler import torchvision -from torch.cuda import amp from tqdm import tqdm +# version check +if torch.__version__.startswith("1.8"): + Autocast = torch.cuda.amp.autocast + GradScaler = torch.cuda.amp.GradScaler +else: + Autocast = torch.amp.autocast + GradScaler = torch.amp.GradScaler + FILE = Path(__file__).resolve() ROOT = FILE.parents[1] # YOLOv5 root directory if str(ROOT) not in sys.path: @@ -198,7 +205,7 @@ def train(opt, device): t0 = time.time() criterion = smartCrossEntropyLoss(label_smoothing=opt.label_smoothing) # loss function best_fitness = 0.0 - scaler = amp.GradScaler(enabled=cuda) + scaler = GradScaler(enabled=cuda) val = test_dir.stem # 'val' or 'test' LOGGER.info( f"Image sizes {imgsz} train, {imgsz} test\n" @@ -219,7 +226,7 @@ def train(opt, device): images, labels = images.to(device, non_blocking=True), labels.to(device) # Forward - with amp.autocast(enabled=cuda): # stability issues when enabled + with Autocast(enabled=device.type != "cpu"): # stability issues when enabled loss = criterion(model(images), labels) # Backward diff --git a/classify/val.py b/classify/val.py index 72bd0e14e..48104e39e 100644 --- a/classify/val.py +++ b/classify/val.py @@ -48,6 +48,11 @@ from utils.general import ( ) from utils.torch_utils import select_device, smart_inference_mode +#version check +if torch.__version__.startswith("1.8"): + Autocast = torch.cuda.amp.autocast +else: + Autocast = torch.amp.autocast @smart_inference_mode() def run( @@ -108,7 +113,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 Autocast(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/models/common.py b/models/common.py index ea893db4b..50dd5b937 100644 --- a/models/common.py +++ b/models/common.py @@ -20,7 +20,6 @@ import requests import torch import torch.nn as nn from PIL import Image -from torch.cuda import amp # Import 'ultralytics' package or install if missing try: @@ -56,6 +55,11 @@ from utils.general import ( ) from utils.torch_utils import copy_attr, smart_inference_mode +# version check +if torch.__version__.startswith("1.8"): + Autocast = torch.cuda.amp.autocast +else: + Autocast = torch.amp.autocast def autopad(k, p=None, d=1): """ @@ -864,7 +868,7 @@ class AutoShape(nn.Module): 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): + with Autocast(enabled=autocast): return self.model(ims.to(p.device).type_as(p), augment=augment) # inference # Pre-process @@ -891,7 +895,7 @@ class AutoShape(nn.Module): 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): + with Autocast(enabled=autocast): # Inference with dt[1]: y = self.model(x, augment=augment) # forward diff --git a/segment/train.py b/segment/train.py index 815c97ce1..e37907ea7 100644 --- a/segment/train.py +++ b/segment/train.py @@ -89,6 +89,14 @@ from utils.torch_utils import ( torch_distributed_zero_first, ) +# version check +if torch.__version__.startswith("1.8"): + Autocast = torch.cuda.amp.autocast + GradScaler = torch.cuda.amp.GradScaler +else: + Autocast = torch.amp.autocast + GradScaler = torch.amp.GradScaler + 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)) @@ -320,7 +328,7 @@ def train(hyp, opt, device, callbacks): 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 = GradScaler(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 +388,7 @@ def train(hyp, opt, device, callbacks): imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) # Forward - with torch.cuda.amp.autocast(amp): + with Autocast(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 1401ccb96..dfd3f90d8 100644 --- a/train.py +++ b/train.py @@ -94,6 +94,14 @@ from utils.torch_utils import ( torch_distributed_zero_first, ) +# version check +if torch.__version__.startswith("1.8"): + Autocast = torch.cuda.amp.autocast + GradScaler = torch.cuda.amp.GradScaler +else: + Autocast = torch.amp.autocast + GradScaler = torch.amp.GradScaler + 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)) @@ -352,7 +360,7 @@ def train(hyp, opt, device, callbacks): 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 = GradScaler(enabled=amp) stopper, stop = EarlyStopping(patience=opt.patience), False compute_loss = ComputeLoss(model) # init loss class callbacks.run("on_train_start") @@ -409,7 +417,7 @@ def train(hyp, opt, device, callbacks): imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) # Forward - with torch.cuda.amp.autocast(amp): + with Autocast(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 9d5ea0a94..eaaf3218b 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -12,7 +12,7 @@ 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.cuda.amp.autocast(amp): + with torch.amp.autocast("cuda", enabled=amp): return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size