fix logging call

amazon-science · Apr 22, 2024 · 0964452 · 0964452
1 parent 9b0e60a
commit 0964452
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 13 deletions.
diff --git a/training/tuning_grid.py b/training/tuning_grid.py
@@ -94,7 +94,7 @@ def main(args):
         single_training = False
     for lr in args.lr_grid:
         for weight_decay in args.wd_grid:
-            tune_log(f"Train for lr={lr} and wd={weight_decay}")
+            tune_log.info(f"Train for lr={lr} and wd={weight_decay}")
             args.weight_decay = weight_decay
             args.lr = lr
             train_acc1, train_acc5, test_acc1, test_acc5 = train_one_model(
@@ -140,8 +140,8 @@ def main(args):
     total_tuning_time_str = str(
         datetime.timedelta(seconds=int(total_tuning_time))
     )
-    tune_log(f"Tuning time: {total_tuning_time_str}")
-    tune_log("Done!!!")
+    tune_log.info(f"Tuning time: {total_tuning_time_str}")
+    tune_log.info("Done!!!")
 
 
 def train_one_epoch(
@@ -312,7 +312,7 @@ def evaluate(
         )
 
     metric_logger.synchronize_between_processes()
-    tune_log(
+    tune_log.info(
         (
             f"{header} Acc@1 {metric_logger.acc1.global_avg:.3f}"
             f"Acc@5 {metric_logger.acc5.global_avg:.3f}"
@@ -472,7 +472,7 @@ def collate_fn(batch):
         pin_memory=True,
     )
 
-    tune_log("Creating model")
+    tune_log.info("Creating model")
     model = registry.MLP_MODEL[args.model_name]["wrapper"]()
     model.to(device)
 
@@ -640,7 +640,7 @@ def collate_fn(batch):
             )
         return
 
-    tune_log("Start training")
+    tune_log.info("Start training")
     start_time = time.time()
     for epoch in range(args.start_epoch, args.epochs):
         if args.distributed:
@@ -699,7 +699,7 @@ def collate_fn(batch):
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-    tune_log(f"Training time {total_time_str}")
+    tune_log.info(f"Training time {total_time_str}")
     if single_training and args.wandb:
         wandb.finish()
     return train_acc1, train_acc5, test_acc1, test_acc5

diff --git a/training/utils.py b/training/utils.py
@@ -185,7 +185,7 @@ def log_every(self, iterable, print_freq, header=None):
                 eta_seconds = iter_time.global_avg * (len(iterable) - i)
                 eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                 if torch.cuda.is_available():
-                    tune_log(
+                    tune_log.info(
                         log_msg.format(
                             i,
                             len(iterable),
@@ -197,7 +197,7 @@ def log_every(self, iterable, print_freq, header=None):
                         )
                     )
                 else:
-                    tune_log(
+                    tune_log.info(
                         log_msg.format(
                             i,
                             len(iterable),
@@ -211,7 +211,7 @@ def log_every(self, iterable, print_freq, header=None):
             end = time.time()
         total_time = time.time() - start_time
         total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-        tune_log(f"{header} Total time: {total_time_str}")
+        tune_log.info(f"{header} Total time: {total_time_str}")
 
 
 class ExponentialMovingAverage(torch.optim.swa_utils.AveragedModel):
@@ -321,9 +321,7 @@ def init_distributed_mode(args):
 
     torch.cuda.set_device(args.gpu)
     args.dist_backend = "nccl"
-    tune_log.info(
-        f"| distributed init (rank {args.rank}): {args.dist_url}", flush=True
-    )
+    tune_log.info(f"| distributed init (rank {args.rank}): {args.dist_url}")
     torch.distributed.init_process_group(
         backend=args.dist_backend,
         init_method=args.dist_url,