From e0d2fcdc5f10841737671a0e4f193bda3fb6e984 Mon Sep 17 00:00:00 2001 From: Morgan Barber Date: Wed, 27 Mar 2024 17:23:09 -0600 Subject: [PATCH 1/2] Clear cuda cache before training loop and validate loop starts --- imageai/Detection/Custom/__init__.py | 2 ++ imageai/Detection/Custom/yolo/validate.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/imageai/Detection/Custom/__init__.py b/imageai/Detection/Custom/__init__.py index e40c711d..f5bcb019 100644 --- a/imageai/Detection/Custom/__init__.py +++ b/imageai/Detection/Custom/__init__.py @@ -218,6 +218,8 @@ def trainModel(self) -> None: os.makedirs(self.__output_models_dir, exist_ok=True) os.makedirs(self.__output_json_dir, exist_ok=True) + torch.cuda.empty_cache() + mp, mr, map50, map50_95, best_fitness = 0, 0, 0, 0, 0.0 nbs = 64 # norminal batch size nb = len(self.__train_loader) # number of batches diff --git a/imageai/Detection/Custom/yolo/validate.py b/imageai/Detection/Custom/yolo/validate.py index 92416ac3..e4a09b3a 100644 --- a/imageai/Detection/Custom/yolo/validate.py +++ b/imageai/Detection/Custom/yolo/validate.py @@ -56,6 +56,8 @@ def run(model, val_dataloader, num_class, net_dim=416, nms_thresh=0.6, objectnes iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95 niou = iouv.numel() + torch.cuda.empty_cache() + p, r, f1, mp, mr, map50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 stats, ap, ap_class = [], [], [] From 283292589695812c928d4ccff58d218c3ccc0615 Mon Sep 17 00:00:00 2001 From: Morgan Barber Date: Wed, 27 Mar 2024 17:33:34 -0600 Subject: [PATCH 2/2] Improve validation loop speed --- imageai/Detection/Custom/yolo/validate.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/imageai/Detection/Custom/yolo/validate.py b/imageai/Detection/Custom/yolo/validate.py index e4a09b3a..0153c893 100644 --- a/imageai/Detection/Custom/yolo/validate.py +++ b/imageai/Detection/Custom/yolo/validate.py @@ -58,10 +58,10 @@ def run(model, val_dataloader, num_class, net_dim=416, nms_thresh=0.6, objectnes torch.cuda.empty_cache() - p, r, f1, mp, mr, map50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 - stats, ap, ap_class = [], [], [] + mp, mr, map50, map = 0.0, 0.0, 0.0, 0.0 + stats, ap = [], [] - for batch_i, (im, targets) in tqdm(enumerate(val_dataloader)): + for im, targets in tqdm(val_dataloader): im = im.to(device) targets = targets.to(device) nb = im.shape[0] # batch @@ -110,7 +110,7 @@ def run(model, val_dataloader, num_class, net_dim=416, nms_thresh=0.6, objectnes # Compute metrics stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy if len(stats) and stats[0].any(): - p, r, ap, f1, ap_class = ap_per_class(*stats) + p, r, ap, _, _ = ap_per_class(*stats) ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95 mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()