mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

refractored the code

pull/13483/head
Bala-Vignesh-Reddy 2025-02-23 20:19:31 +05:30
parent 8da9408121 a63bfd38c1
commit b65447a098
5 changed files with 43 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
classify/train.py
View File

@ -29,6 +29,14 @@ import torch.optim.lr_scheduler as lr_scheduler
import torchvision
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:
@ -197,13 +205,7 @@ def train(opt, device):
t0 = time.time()
criterion = smartCrossEntropyLoss(label_smoothing=opt.label_smoothing) # loss function
best_fitness = 0.0
# adding a check to torch version
scaler = None
if torch.__version__.startswith("1.8"):
scaler = torch.cuda.amp.GradScaler(enabled=cuda)
else:
scaler = torch.amp.GradScaler("cuda", enabled=cuda)
scaler = GradScaler(enabled=cuda)
val = test_dir.stem # 'val' or 'test'
LOGGER.info(
f"Image sizes {imgsz} train, {imgsz} test\n"
@ -224,12 +226,7 @@ def train(opt, device):
images, labels = images.to(device, non_blocking=True), labels.to(device)
# Forward
amp_autocast = None
if torch.__version__.startswith("1.8"):
amp_autocast = torch.cuda.amp.autocast(enabled=device.type != "cpu")
else:
amp_autocast = torch.amp.autocast("cuda", enabled=device.type != "cpu")
with amp_autocast:
with Autocast(enabled=device.type != "cpu"): # stability issues when enabled
loss = criterion(model(images), labels)
# Backward

15
classify/val.py
View File

@ -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,15 +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)
# checking the version
amp_autocast = None
if torch.__version__.startswith("1.8"):
amp_autocast = torch.cuda.amp.autocast(enabled=device.type != "cpu")
else:
amp_autocast = torch.amp.autocast("cuda", enabled=device.type != "cpu")
with amp_autocast:
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)

19
models/common.py
View File

@ -55,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):
"""
@ -863,12 +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
amp_autocast = None
if torch.__version__.startswith("1.8"):
amp_autocast = torch.cuda.amp.autocast(enabled=autocast)
else:
amp_autocast = torch.amp.autocast("cuda", enabled=autocast)
with amp_autocast:
with Autocast(enabled=autocast):
return self.model(ims.to(p.device).type_as(p), augment=augment) # inference
# Pre-process
@ -895,12 +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
amp_autocast = None
if torch.__version__.startswith("1.8"):
amp_autocast = torch.cuda.amp.autocast(enabled=autocast)
else:
amp_autocast = torch.amp.autocast("cuda", enabled=autocast)
with amp_autocast:
with Autocast(enabled=autocast):
# Inference
with dt[1]:
y = self.model(x, augment=augment) # forward

18
segment/train.py
View File

@ -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,13 +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 = None
if torch.__version__.startswith("1.8"):
scaler = torch.cuda.amp.GradScaler(enabled=amp)
else:
scaler = torch.amp.GradScaler("cuda", 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')
@ -391,7 +393,7 @@ def train(hyp, opt, device, callbacks):
else:
amp_autocast = torch.amp.autocast("cuda", enabled=amp)
# Forward
with amp_autocast:
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:

21
train.py
View File

@ -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,11 +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 = None
if torch.__version__.startswith("1.8"):
scaler = torch.cuda.amp.GradScaler(enabled=amp)
else:
scaler = torch.amp.GradScaler("cuda", 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")
@ -413,12 +417,7 @@ def train(hyp, opt, device, callbacks):
imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
# Forward
amp_autocast = None
if torch.__version__.startswith("1.8"):
amp_autocast = torch.cuda.amp.autocast(enabled=amp)
else:
amp_autocast = torch.amp.autocast("cuda", enabled=amp)
with amp_autocast:
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: