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

birder/scripts/train_dino_v2.py

@@ -214,7 +214,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
@@ -417,6 +418,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
 
     #
@@ -438,20 +440,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)
 
@@ -492,11 +493,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_dino_v2_dist.py

@@ -215,7 +215,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
@@ -240,6 +241,7 @@ def train(args: argparse.Namespace) -> None:
         args.network, sample_shape[1], 0, config=args.model_config, size=args.size
     )
     student_backbone_ema.load_state_dict(student_backbone.state_dict())
+    student_backbone_ema.requires_grad_(False)
 
     teacher_backbone = registry.net_factory(
         args.teacher,
@@ -248,6 +250,11 @@ def train(args: argparse.Namespace) -> None:
         config=args.teacher_model_config,
         size=args.size,
     )
+    assert student_backbone.max_stride == teacher_backbone.max_stride, (
+        "Student and teacher max_stride must match for distillation "
+        f"(student={student_backbone.max_stride}, teacher={teacher_backbone.max_stride})"
+    )
+
     student_backbone.set_dynamic_size()
     if args.ibot_separate_head is False:
         args.ibot_out_dim = args.dino_out_dim
@@ -433,6 +440,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
 
     #
@@ -454,20 +462,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)
 
@@ -507,11 +514,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)
 
@@ -604,6 +613,7 @@ def train(args: argparse.Namespace) -> None:
     for epoch in range(begin_epoch, args.stop_epoch):
         tic = time.time()
         net.train()
+        teacher.eval()
         running_loss = training_utils.SmoothedValue()
         running_loss_dino_local = training_utils.SmoothedValue()
         running_loss_dino_global = training_utils.SmoothedValue()

birder/scripts/train_franca.py

@@ -241,7 +241,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
@@ -444,6 +445,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
 
     #
@@ -465,20 +467,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)
 
@@ -519,11 +520,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_i_jepa.py

@@ -120,7 +120,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
@@ -284,6 +285,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
 
     #
@@ -305,20 +307,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)
 
@@ -351,11 +352,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)
 
@@ -521,12 +524,13 @@ def train(args: argparse.Namespace) -> None:
                     if step_update is True:
                         scheduler.step()
 
-            # EMA update for the target encoder
-            with torch.no_grad():
-                m = momentum_schedule[global_iter]
-                torch._foreach_lerp_(  # pylint: disable=protected-access
-                    list(target_encoder.parameters()), list(encoder.parameters()), weight=1 - m
-                )
+            if optimizer_update is True:
+                # EMA update for the target encoder
+                with torch.no_grad():
+                    m = momentum_schedule[global_iter]
+                    torch._foreach_lerp_(  # pylint: disable=protected-access
+                        list(target_encoder.parameters()), list(encoder.parameters()), weight=1 - m
+                    )
 
             # Statistics
             running_loss.update(loss.detach())

birder/scripts/train_ibot.py

@@ -143,7 +143,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
@@ -351,6 +352,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
 
     #
@@ -372,20 +374,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)
 
@@ -418,11 +419,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_kd.py

@@ -186,8 +186,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:
@@ -246,8 +247,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
@@ -336,20 +337,18 @@ 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)
 
@@ -375,11 +374,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)
 
@@ -387,15 +389,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 = student
     if args.distributed is True:
         student = torch.nn.parallel.DistributedDataParallel(
@@ -493,6 +495,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()
@@ -571,10 +574,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(student)
-                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_mim.py

@@ -130,7 +130,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:
@@ -162,6 +163,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
@@ -254,20 +256,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)
 
@@ -293,11 +294,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)
 
@@ -599,9 +602,7 @@ def get_args_parser() -> argparse.ArgumentParser:
             "('drop_path_rate=0.2' or '{\"units\": [3, 24, 36, 3], \"dropout\": 0.2}'"
         ),
     )
-    parser.add_argument(
-        "--mask-ratio", type=float, default=None, help="mask ratio for MIM training (default: model-specific)"
-    )
+    parser.add_argument("--mask-ratio", type=float, help="mask ratio for MIM training (default: model-specific)")
     parser.add_argument("--min-mask-size", type=int, default=1, help="minimum mask unit size in patches")
     training_cli.add_optimization_args(parser)
     training_cli.add_lr_wd_args(parser)

birder/scripts/train_mmcr.py

@@ -154,7 +154,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:
@@ -186,6 +187,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_rotnet.py

@@ -168,7 +168,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:
@@ -200,6 +201,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
@@ -268,20 +270,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)
 
@@ -307,11 +308,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)