birder 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

birder/scripts/train.py

@@ -160,8 +160,9 @@ def train(args: argparse.Namespace) -> None:
 
     num_outputs = len(class_to_idx)
     batch_size: int = args.batch_size
-    model_ema_steps: int = args.model_ema_steps * args.grad_accum_steps
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    model_ema_steps: int = args.model_ema_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     # Set data iterators
     if args.mixup_alpha is not None or args.cutmix is True:
@@ -220,8 +221,8 @@ def train(args: argparse.Namespace) -> None:
             pin_memory=True,
         )
 
-    optimizer_steps_per_epoch = math.ceil(len(training_loader) / args.grad_accum_steps)
-    assert args.model_ema is False or args.model_ema_steps <= optimizer_steps_per_epoch
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+    assert args.model_ema is False or model_ema_steps <= optimizer_steps_per_epoch
 
     last_batch_idx = len(training_loader) - 1
     begin_epoch = 1
@@ -317,20 +318,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -356,11 +356,14 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False),
+            lrs,
+        )
         plt.show()
         raise SystemExit(0)
 
@@ -368,15 +371,15 @@ def train(args: argparse.Namespace) -> None:
     # Distributed (DDP) and Model EMA
     #
     if args.model_ema_warmup is not None:
-        ema_warmup_epochs = args.model_ema_warmup
+        ema_warmup_steps = args.model_ema_warmup * optimizer_steps_per_epoch
     elif args.warmup_epochs is not None:
-        ema_warmup_epochs = args.warmup_epochs
+        ema_warmup_steps = args.warmup_epochs * optimizer_steps_per_epoch
     elif args.warmup_steps is not None:
-        ema_warmup_epochs = args.warmup_steps // steps_per_epoch
+        ema_warmup_steps = args.warmup_steps
     else:
-        ema_warmup_epochs = 0
+        ema_warmup_steps = 0
 
-    logger.debug(f"EMA warmup epochs = {ema_warmup_epochs}")
+    logger.debug(f"EMA warmup steps = {ema_warmup_steps}")
     net_without_ddp = net
     if args.distributed is True:
         net = torch.nn.parallel.DistributedDataParallel(
@@ -474,6 +477,7 @@ def train(args: argparse.Namespace) -> None:
     #
     # Training loop
     #
+    optimizer_step = (begin_epoch - 1) * optimizer_steps_per_epoch
     logger.info(f"Starting training with learning rate of {last_lr}")
     for epoch in range(begin_epoch, args.stop_epoch):
         tic = time.time()
@@ -542,10 +546,13 @@ def train(args: argparse.Namespace) -> None:
                     if step_update is True:
                         scheduler.step()
 
+            if optimizer_update is True:
+                optimizer_step += 1
+
             # Exponential moving average
-            if args.model_ema is True and i % model_ema_steps == 0:
+            if args.model_ema is True and optimizer_update is True and optimizer_step % model_ema_steps == 0:
                 model_ema.update_parameters(net)
-                if epoch <= ema_warmup_epochs:
+                if ema_warmup_steps > 0 and optimizer_step <= ema_warmup_steps:
                     # Reset ema buffer to keep copying weights during warmup period
                     model_ema.n_averaged.fill_(0)  # pylint: disable=no-member
 

birder/scripts/train_barlow_twins.py

@@ -149,7 +149,8 @@ def train(args: argparse.Namespace) -> None:
     logger.info(f"Training on {len(training_dataset):,} samples")
 
     batch_size: int = args.batch_size
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     # Data loaders and samplers
     if args.distributed is True:
@@ -181,6 +182,7 @@ def train(args: argparse.Namespace) -> None:
             drop_last=args.drop_last,
         )
 
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
     last_batch_idx = len(training_loader) - 1
     begin_epoch = 1
     epochs = args.epochs + 1
@@ -244,20 +246,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -283,11 +284,13 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs
+        )
         plt.show()
         raise SystemExit(0)
 

birder/scripts/train_byol.py

@@ -151,7 +151,8 @@ def train(args: argparse.Namespace) -> None:
     logger.info(f"Training on {len(training_dataset):,} samples")
 
     batch_size: int = args.batch_size
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     # Data loaders and samplers
     if args.distributed is True:
@@ -183,6 +184,7 @@ def train(args: argparse.Namespace) -> None:
             drop_last=args.drop_last,
         )
 
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
     last_batch_idx = len(training_loader) - 1
     begin_epoch = 1
     epochs = args.epochs + 1
@@ -256,20 +258,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -295,11 +296,13 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs
+        )
         plt.show()
         raise SystemExit(0)
 

birder/scripts/train_capi.py

@@ -115,7 +115,8 @@ def train(args: argparse.Namespace) -> None:
     torch.autograd.set_detect_anomaly(args.grad_anomaly_detection)
 
     batch_size: int = args.batch_size
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     begin_epoch = 1
     epochs = args.epochs + 1
@@ -278,6 +279,7 @@ def train(args: argparse.Namespace) -> None:
             drop_last=args.drop_last,
         )
 
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
     last_batch_idx = len(training_loader) - 1
 
     #
@@ -300,22 +302,21 @@ def train(args: argparse.Namespace) -> None:
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
     clustering_lr = lr / 2
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
     clustering_optimizer = torch.optim.AdamW(teacher.head.parameters(), lr=clustering_lr, betas=[0.9, 0.95])
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
-    clustering_scheduler = training_utils.get_scheduler(clustering_optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
+    clustering_scheduler = training_utils.get_scheduler(clustering_optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
         clustering_optimizer.step = torch.compile(clustering_optimizer.step, fullgraph=False)
@@ -324,7 +325,7 @@ def train(args: argparse.Namespace) -> None:
     if args.warmup_epochs is not None:
         warmup_epochs = args.warmup_epochs
     elif args.warmup_steps is not None:
-        warmup_epochs = args.warmup_steps / steps_per_epoch
+        warmup_epochs = args.warmup_steps / scheduler_steps_per_epoch
     else:
         warmup_epochs = 0.0
 
@@ -353,11 +354,13 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs
+        )
         plt.show()
         raise SystemExit(0)
 

birder/scripts/train_data2vec.py

@@ -106,7 +106,8 @@ def train(args: argparse.Namespace) -> None:
     torch.autograd.set_detect_anomaly(args.grad_anomaly_detection)
 
     batch_size: int = args.batch_size
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     begin_epoch = 1
     epochs = args.epochs + 1
@@ -245,6 +246,7 @@ def train(args: argparse.Namespace) -> None:
             drop_last=args.drop_last,
         )
 
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
     last_batch_idx = len(training_loader) - 1
 
     #
@@ -266,20 +268,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -310,11 +311,13 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs
+        )
         plt.show()
         raise SystemExit(0)
 

birder/scripts/train_data2vec2.py

@@ -112,7 +112,8 @@ def train(args: argparse.Namespace) -> None:
     torch.autograd.set_detect_anomaly(args.grad_anomaly_detection)
 
     batch_size: int = args.batch_size
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     begin_epoch = 1
     epochs = args.epochs + 1
@@ -254,6 +255,7 @@ def train(args: argparse.Namespace) -> None:
             drop_last=args.drop_last,
         )
 
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
     last_batch_idx = len(training_loader) - 1
 
     #
@@ -275,20 +277,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -319,11 +320,13 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs
+        )
         plt.show()
         raise SystemExit(0)
 

birder/scripts/train_detection.py

@@ -192,8 +192,9 @@ def train(args: argparse.Namespace) -> None:
 
     num_outputs = len(class_to_idx)  # Does not include background class
     batch_size: int = args.batch_size
-    model_ema_steps: int = args.model_ema_steps * args.grad_accum_steps
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    model_ema_steps: int = args.model_ema_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     # Data loaders and samplers
     (train_sampler, validation_sampler) = training_utils.get_samplers(args, training_dataset, validation_dataset)
@@ -224,8 +225,8 @@ def train(args: argparse.Namespace) -> None:
         drop_last=args.drop_last,
     )
 
-    optimizer_steps_per_epoch = math.ceil(len(training_loader) / args.grad_accum_steps)
-    assert args.model_ema is False or args.model_ema_steps <= optimizer_steps_per_epoch
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+    assert args.model_ema is False or model_ema_steps <= optimizer_steps_per_epoch
 
     last_batch_idx = len(training_loader) - 1
     begin_epoch = 1
@@ -369,20 +370,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -408,11 +408,14 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False),
+            lrs,
+        )
         plt.show()
         raise SystemExit(0)
 
@@ -420,15 +423,15 @@ def train(args: argparse.Namespace) -> None:
     # Distributed (DDP) and Model EMA
     #
     if args.model_ema_warmup is not None:
-        ema_warmup_epochs = args.model_ema_warmup
+        ema_warmup_steps = args.model_ema_warmup * optimizer_steps_per_epoch
     elif args.warmup_epochs is not None:
-        ema_warmup_epochs = args.warmup_epochs
+        ema_warmup_steps = args.warmup_epochs * optimizer_steps_per_epoch
     elif args.warmup_steps is not None:
-        ema_warmup_epochs = args.warmup_steps // steps_per_epoch
+        ema_warmup_steps = args.warmup_steps
     else:
-        ema_warmup_epochs = 0
+        ema_warmup_steps = 0
 
-    logger.debug(f"EMA warmup epochs = {ema_warmup_epochs}")
+    logger.debug(f"EMA warmup steps = {ema_warmup_steps}")
     net_without_ddp = net
     if args.distributed is True:
         net = torch.nn.parallel.DistributedDataParallel(
@@ -532,11 +535,13 @@ def train(args: argparse.Namespace) -> None:
     #
     # Training loop
     #
+    optimizer_step = (begin_epoch - 1) * optimizer_steps_per_epoch
     logger.info(f"Starting training with learning rate of {last_lr}")
     for epoch in range(begin_epoch, args.stop_epoch):
         tic = time.time()
         net.train()
         running_loss = training_utils.SmoothedValue()
+        loss_trackers: dict[str, training_utils.SmoothedValue] = {}
         validation_metrics.reset()
 
         if args.distributed is True:
@@ -598,16 +603,28 @@ def train(args: argparse.Namespace) -> None:
                     if step_update is True:
                         scheduler.step()
 
+            if optimizer_update is True:
+                optimizer_step += 1
+
             # Exponential moving average
-            if args.model_ema is True and i % model_ema_steps == 0:
+            if args.model_ema is True and optimizer_update is True and optimizer_step % model_ema_steps == 0:
                 model_ema.update_parameters(net)
-                if epoch <= ema_warmup_epochs:
+                if ema_warmup_steps > 0 and optimizer_step <= ema_warmup_steps:
                     # Reset ema buffer to keep copying weights during warmup period
                     model_ema.n_averaged.fill_(0)  # pylint: disable=no-member
 
             # Statistics
             running_loss.update(loss.detach())
 
+            # Dynamically create trackers on first batch
+            if len(loss_trackers) == 0:
+                for key in losses.keys():
+                    loss_trackers[key] = training_utils.SmoothedValue()
+
+            # Update individual loss trackers
+            for key, value in losses.items():
+                loss_trackers[key].update(value.detach())
+
             # Write statistics
             if (i == last_batch_idx) or (i + 1) % args.log_interval == 0:
                 time_now = time.time()
@@ -624,6 +641,9 @@ def train(args: argparse.Namespace) -> None:
                 cur_lr = float(max(scheduler.get_last_lr()))
 
                 running_loss.synchronize_between_processes(device)
+                for tracker in loss_trackers.values():
+                    tracker.synchronize_between_processes(device)
+
                 with training_utils.single_handler_logging(logger, file_handler, enabled=not disable_tqdm) as log:
                     log.info(
                         f"[Trn] Epoch {epoch}/{epochs-1}, iter {i+1}/{last_batch_idx+1}  "
@@ -636,9 +656,11 @@ def train(args: argparse.Namespace) -> None:
                     )
 
                 if training_utils.is_local_primary(args) is True:
+                    loss_dict = {"training": running_loss.avg}
+                    loss_dict.update({k: v.avg for k, v in loss_trackers.items()})
                     summary_writer.add_scalars(
                         "loss",
-                        {"training": running_loss.avg},
+                        loss_dict,
                         ((epoch - 1) * len(training_dataset)) + (i * batch_size * args.world_size),
                     )
 
@@ -649,6 +671,8 @@ def train(args: argparse.Namespace) -> None:
 
         # Epoch training metrics
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} training_loss: {running_loss.global_avg:.4f}")
+        for key, tracker in loss_trackers.items():
+            logger.info(f"[Trn] Epoch {epoch}/{epochs-1} {key}: {tracker.global_avg:.4f}")
 
         # Validation
         eval_model.eval()

birder/scripts/train_dino_v1.py

@@ -186,7 +186,8 @@ def train(args: argparse.Namespace) -> None:
     logger.info(f"Training on {len(training_dataset):,} samples")
 
     batch_size: int = args.batch_size
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     # Data loaders and samplers
     if args.distributed is True:
@@ -218,6 +219,7 @@ def train(args: argparse.Namespace) -> None:
             drop_last=args.drop_last,
         )
 
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
     last_batch_idx = len(training_loader) - 1
     begin_epoch = 1
     epochs = args.epochs + 1
@@ -352,20 +354,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -398,11 +399,13 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs
+        )
         plt.show()
         raise SystemExit(0)