Small adjustments

mindee · Jan 23, 2025 · f7fb3e1 · f7fb3e1
1 parent 87c25e1
commit f7fb3e1
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 31 deletions.
diff --git a/references/classification/train_tensorflow_character.py b/references/classification/train_tensorflow_character.py
@@ -345,22 +345,23 @@ def wandb_log_at_step(train_loss=None, val_loss=None, lr=None):
         task.upload_artifact("config", config)
 
         def clearml_log_at_step(train_loss=None, val_loss=None, lr=None):
+            logger = Logger.current_logger()
             if train_loss is not None:
-                Logger.current_logger().report_scalar(
+                logger.report_scalar(
                     title="Training Step Loss",
                     series="train_loss_step",
                     iteration=global_step,
                     value=train_loss,
                 )
             if val_loss is not None:
-                Logger.current_logger().report_scalar(
+                logger.report_scalar(
                     title="Validation Step Loss",
                     series="val_loss_step",
                     iteration=global_step,
                     value=val_loss,
                 )
             if lr is not None:
-                Logger.current_logger().report_scalar(
+                logger.report_scalar(
                     title="Step Learning Rate",
                     series="step_lr",
                     iteration=global_step,

diff --git a/references/detection/evaluate_pytorch.py b/references/detection/evaluate_pytorch.py
@@ -9,7 +9,6 @@
 
 os.environ["USE_TORCH"] = "1"
 
-import logging
 import multiprocessing as mp
 import time
 from pathlib import Path
@@ -63,7 +62,8 @@ def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
 
 
 def main(args):
-    print(args)
+    pbar = tqdm(disable=True)
+    pbar.write(str(args))
 
     if not isinstance(args.workers, int):
         args.workers = min(16, mp.cpu_count())
@@ -116,13 +116,13 @@ def main(args):
         pin_memory=torch.cuda.is_available(),
         collate_fn=ds.collate_fn,
     )
-    print(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in {len(test_loader)} batches)")
+    pbar.write(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in {len(test_loader)} batches)")
 
     batch_transforms = Normalize(mean=mean, std=std)
 
     # Resume weights
     if isinstance(args.resume, str):
-        print(f"Resuming {args.resume}")
+        pbar.write(f"Resuming {args.resume}")
         checkpoint = torch.load(args.resume, map_location="cpu")
         model.load_state_dict(checkpoint)
 
@@ -136,17 +136,17 @@ def main(args):
     elif torch.cuda.is_available():
         args.device = 0
     else:
-        logging.warning("No accessible GPU, targe device set to CPU.")
+        pbar.write("No accessible GPU, target device set to CPU.")
     if torch.cuda.is_available():
         torch.cuda.set_device(args.device)
         model = model.cuda()
 
     # Metrics
     metric = LocalizationConfusion(use_polygons=args.rotation)
 
-    print("Running evaluation")
+    pbar.write("Running evaluation")
     val_loss, recall, precision, mean_iou = evaluate(model, test_loader, batch_transforms, metric, amp=args.amp)
-    print(
+    pbar.write(
         f"Validation loss: {val_loss:.6} (Recall: {recall:.2%} | Precision: {precision:.2%} | Mean IoU: {mean_iou:.2%})"
     )
 

diff --git a/references/detection/evaluate_tensorflow.py b/references/detection/evaluate_tensorflow.py
@@ -61,7 +61,8 @@ def evaluate(model, val_loader, batch_transforms, val_metric):
 
 
 def main(args):
-    print(args)
+    pbar = tqdm(disable=True)
+    pbar.write(str(args))
 
     # AMP
     if args.amp:
@@ -78,7 +79,7 @@ def main(args):
 
     # Resume weights
     if isinstance(args.resume, str):
-        print(f"Resuming {args.resume}")
+        pbar.write(f"Resuming {args.resume}")
         model.load_weights(args.resume).expect_partial()
 
     input_shape = model.cfg["input_shape"] if input_shape is None else input_shape
@@ -116,16 +117,16 @@ def main(args):
         drop_last=False,
         shuffle=False,
     )
-    print(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in {len(test_loader)} batches)")
+    pbar.write(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in {len(test_loader)} batches)")
 
     batch_transforms = T.Normalize(mean=mean, std=std)
 
     # Metrics
     metric = LocalizationConfusion(use_polygons=args.rotation)
 
-    print("Running evaluation")
+    pbar.write("Running evaluation")
     val_loss, recall, precision, mean_iou = evaluate(model, test_loader, batch_transforms, metric)
-    print(
+    pbar.write(
         f"Validation loss: {val_loss:.6} (Recall: {recall:.2%} | Precision: {precision:.2%} | Mean IoU: {mean_iou:.2%})"
     )
 

diff --git a/references/recognition/evaluate_pytorch.py b/references/recognition/evaluate_pytorch.py
@@ -13,7 +13,11 @@
 import torch
 from torch.utils.data import DataLoader, SequentialSampler
 from torchvision.transforms import Normalize
-from tqdm import tqdm
+
+if os.getenv("TQDM_SLACK_TOKEN") and os.getenv("TQDM_SLACK_CHANNEL"):
+    from tqdm.contrib.slack import tqdm
+else:
+    from tqdm.auto import tqdm
 
 from doctr import datasets
 from doctr import transforms as T
@@ -30,7 +34,8 @@ def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
     val_metric.reset()
     # Validation loop
     val_loss, batch_cnt = 0, 0
-    for images, targets in tqdm(val_loader):
+    pbar = tqdm(val_loader)
+    for images, targets in pbar:
         try:
             if torch.cuda.is_available():
                 images = images.cuda()
@@ -50,7 +55,7 @@ def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
             val_loss += out["loss"].item()
             batch_cnt += 1
         except ValueError:
-            print(f"unexpected symbol/s in targets:\n{targets} \n--> skip batch")
+            pbar.write(f"unexpected symbol/s in targets:\n{targets} \n--> skip batch")
             continue
 
     val_loss /= batch_cnt
@@ -59,7 +64,8 @@ def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
 
 
 def main(args):
-    print(args)
+    pbar = tqdm(disable=True)
+    pbar.write(str(args))
 
     torch.backends.cudnn.benchmark = True
 
@@ -75,7 +81,7 @@ def main(args):
 
     # Resume weights
     if isinstance(args.resume, str):
-        print(f"Resuming {args.resume}")
+        pbar.write(f"Resuming {args.resume}")
         checkpoint = torch.load(args.resume, map_location="cpu")
         model.load_state_dict(checkpoint)
 
@@ -106,7 +112,7 @@ def main(args):
         pin_memory=torch.cuda.is_available(),
         collate_fn=ds.collate_fn,
     )
-    print(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in {len(test_loader)} batches)")
+    pbar.write(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in {len(test_loader)} batches)")
 
     mean, std = model.cfg["mean"], model.cfg["std"]
     batch_transforms = Normalize(mean=mean, std=std)
@@ -124,14 +130,14 @@ def main(args):
     elif torch.cuda.is_available():
         args.device = 0
     else:
-        print("No accessible GPU, targe device set to CPU.")
+        pbar.write("No accessible GPU, target device set to CPU.")
     if torch.cuda.is_available():
         torch.cuda.set_device(args.device)
         model = model.cuda()
 
-    print("Running evaluation")
+    pbar.write("Running evaluation")
     val_loss, exact_match, partial_match = evaluate(model, test_loader, batch_transforms, val_metric, amp=args.amp)
-    print(f"Validation loss: {val_loss:.6} (Exact: {exact_match:.2%} | Partial: {partial_match:.2%})")
+    pbar.write(f"Validation loss: {val_loss:.6} (Exact: {exact_match:.2%} | Partial: {partial_match:.2%})")
 
 
 def parse_args():

diff --git a/references/recognition/evaluate_tensorflow.py b/references/recognition/evaluate_tensorflow.py
@@ -16,7 +16,11 @@
 
 import tensorflow as tf
 from tensorflow.keras import mixed_precision
-from tqdm import tqdm
+
+if os.getenv("TQDM_SLACK_TOKEN") and os.getenv("TQDM_SLACK_CHANNEL"):
+    from tqdm.contrib.slack import tqdm
+else:
+    from tqdm.auto import tqdm
 
 gpu_devices = tf.config.list_physical_devices("GPU")
 if any(gpu_devices):
@@ -35,7 +39,8 @@ def evaluate(model, val_loader, batch_transforms, val_metric):
     # Validation loop
     val_loss, batch_cnt = 0, 0
     val_iter = iter(val_loader)
-    for images, targets in tqdm(val_iter):
+    pbar = tqdm(val_iter)
+    for images, targets in pbar:
         try:
             images = batch_transforms(images)
             out = model(images, target=targets, return_preds=True, training=False)
@@ -49,7 +54,7 @@ def evaluate(model, val_loader, batch_transforms, val_metric):
             val_loss += out["loss"].numpy().mean()
             batch_cnt += 1
         except ValueError:
-            print(f"unexpected symbol/s in targets:\n{targets} \n--> skip batch")
+            pbar.write(f"unexpected symbol/s in targets:\n{targets} \n--> skip batch")
             continue
 
     val_loss /= batch_cnt
@@ -58,7 +63,8 @@ def evaluate(model, val_loader, batch_transforms, val_metric):
 
 
 def main(args):
-    print(args)
+    pbar = tqdm(disable=True)
+    pbar.write(str(args))
 
     # AMP
     if args.amp:
@@ -99,17 +105,17 @@ def main(args):
         drop_last=False,
         shuffle=False,
     )
-    print(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in {len(test_loader)} batches)")
+    pbar.write(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in {len(test_loader)} batches)")
 
     mean, std = model.cfg["mean"], model.cfg["std"]
     batch_transforms = T.Normalize(mean=mean, std=std)
 
     # Metrics
     val_metric = TextMatch()
 
-    print("Running evaluation")
+    pbar.write("Running evaluation")
     val_loss, exact_match, partial_match = evaluate(model, test_loader, batch_transforms, val_metric)
-    print(f"Validation loss: {val_loss:.6} (Exact: {exact_match:.2%} | Partial: {partial_match:.2%})")
+    pbar.write(f"Validation loss: {val_loss:.6} (Exact: {exact_match:.2%} | Partial: {partial_match:.2%})")
 
 
 def parse_args():