birder 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

birder/scripts/train_data2vec.py

@@ -384,11 +384,12 @@ def train(args: argparse.Namespace) -> None:
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+
     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()
 
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)
@@ -463,7 +464,7 @@ def train(args: argparse.Namespace) -> None:
             running_loss.update(loss.detach())
 
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -603,6 +604,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -617,7 +619,6 @@ def get_args_parser() -> argparse.ArgumentParser:
         default=0.999,
         help="base EMA parameter for teacher update, set a higher value with small batches",
     )
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     training_cli.add_optimization_args(parser)
     training_cli.add_lr_wd_args(parser)
     training_cli.add_lr_scheduler_args(parser)

birder/scripts/train_data2vec2.py

@@ -393,11 +393,12 @@ def train(args: argparse.Namespace) -> None:
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+
     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()
 
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)
@@ -473,7 +474,7 @@ def train(args: argparse.Namespace) -> None:
             running_loss.update(loss.detach())
 
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -615,6 +616,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -635,7 +637,6 @@ def get_args_parser() -> argparse.ArgumentParser:
         default=0.9998,
         help="base EMA parameter for teacher update, set a higher value with small batches",
     )
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     training_cli.add_optimization_args(parser)
     training_cli.add_lr_wd_args(parser)
     training_cli.add_lr_scheduler_args(parser)

birder/scripts/train_detection.py

@@ -538,12 +538,14 @@ def train(args: argparse.Namespace) -> None:
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+    loss_trackers: dict[str, training_utils.SmoothedValue] = {}
+
     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 or virtual_epoch_mode is True:
@@ -634,7 +636,7 @@ def train(args: argparse.Namespace) -> None:
                 loss_trackers[key].update(value.detach())
 
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -889,6 +891,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -897,8 +900,8 @@ def get_args_parser() -> argparse.ArgumentParser:
             "('drop_path_rate=0.2' or '{\"units\": [3, 24, 36, 3], \"dropout\": 0.2}'"
         ),
     )
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument("--backbone", type=str, help="the neural network to used as backbone")
+    parser.add_argument("--backbone-tag", type=str, help="backbone training log tag (loading only)")
     parser.add_argument(
         "--backbone-model-config",
         action=cli.FlexibleDictAction,
@@ -907,7 +910,6 @@ def get_args_parser() -> argparse.ArgumentParser:
             "('drop_path_rate=0.2' or '{\"units\": [3, 24, 36, 3], \"dropout\": 0.2}'"
         ),
     )
-    parser.add_argument("--backbone-tag", type=str, help="backbone training log tag (loading only)")
     parser.add_argument("--backbone-epoch", type=int, help="load backbone weights from selected epoch")
     parser.add_argument(
         "--backbone-pretrained",

birder/scripts/train_dino_v1.py

@@ -480,12 +480,13 @@ def train(args: argparse.Namespace) -> None:
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+    train_proto_agreement = training_utils.SmoothedValue()
+
     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()
-        train_proto_agreement = training_utils.SmoothedValue()
 
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)
@@ -581,7 +582,7 @@ def train(args: argparse.Namespace) -> None:
             train_proto_agreement.update(training_utils.accuracy(pred_teacher, pred_student))
 
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -733,6 +734,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -788,7 +790,6 @@ def get_args_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--local-crop-size", type=int, nargs="+", default=[96, 96], metavar=("H", "W"), help="local view size"
     )
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--backbone-epoch",
         type=int,

birder/scripts/train_dino_v2.py

@@ -582,22 +582,26 @@ def train(args: argparse.Namespace) -> None:
     #
     # Training loop
     #
-    track_agreement = not args.no_agreement_metrics
+    track_extended_metrics = not args.no_extended_metrics
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+    running_loss_dino_local = training_utils.SmoothedValue()
+    running_loss_dino_global = training_utils.SmoothedValue()
+    running_loss_koleo = training_utils.SmoothedValue()
+    running_loss_ibot_patch = training_utils.SmoothedValue()
+    if track_extended_metrics is True:
+        train_proto_agreement = training_utils.SmoothedValue()
+        train_patch_agreement = training_utils.SmoothedValue()
+        running_target_entropy = training_utils.SmoothedValue()
+        running_dino_center_drift = training_utils.SmoothedValue()
+        running_ibot_center_drift = training_utils.SmoothedValue()
+
     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()
-        running_loss_dino_local = training_utils.SmoothedValue()
-        running_loss_dino_global = training_utils.SmoothedValue()
-        running_loss_koleo = training_utils.SmoothedValue()
-        running_loss_ibot_patch = training_utils.SmoothedValue()
-        if track_agreement is True:
-            train_proto_agreement = training_utils.SmoothedValue()
-            train_patch_agreement = training_utils.SmoothedValue()
 
         if args.sinkhorn_queue_size is not None:
             queue_active = epoch > args.sinkhorn_queue_warmup_epochs
@@ -662,6 +666,11 @@ def train(args: argparse.Namespace) -> None:
                     )
                     teacher_patch_tokens_raw = teacher_masked_patch_tokens_after_head
                     if args.centering == "centering":
+                        # Track centers before update for drift computation
+                        if track_extended_metrics is True:
+                            prev_dino_center = dino_loss.center.clone()
+                            prev_ibot_center = ibot_patch_loss.center.clone()
+
                         teacher_dino_softmax_centered_list = dino_loss.softmax_center_teacher(
                             teacher_embedding_after_head, teacher_temp=teacher_temp
                         ).view(n_global_crops, -1, *teacher_embedding_after_head.shape[1:])
@@ -788,7 +797,7 @@ def train(args: argparse.Namespace) -> None:
             running_loss_koleo.update(loss_koleo.detach())
             running_loss_ibot_patch.update(loss_ibot_patch.detach())
 
-            if track_agreement is True:
+            if track_extended_metrics is True:
                 probs_teacher = teacher_embedding_after_head.chunk(n_global_crops)
                 probs_student = student_global_embedding_after_head.chunk(n_global_crops)
                 pred_teacher = probs_teacher[0].argmax(dim=1)
@@ -799,8 +808,27 @@ def train(args: argparse.Namespace) -> None:
                 pred_patch_student = student_global_masked_patch_tokens_after_head.argmax(dim=1)
                 train_patch_agreement.update(training_utils.accuracy(pred_patch_teacher, pred_patch_student))
 
+                with torch.no_grad():
+                    p = teacher_dino_softmax_centered_list.detach()
+                    p = p.reshape(-1, p.size(-1))  # (N, D)
+
+                    # Mean distribution over prototypes (marginal)
+                    m = p.mean(dim=0).clamp_min(1e-12)
+
+                    # Entropy of the marginal
+                    entropy = -(m * m.log()).sum()
+
+                running_target_entropy.update(entropy.detach())
+
+                # Compute center drift
+                if args.centering == "centering":
+                    dino_center_drift = torch.norm(dino_loss.center - prev_dino_center, p=2).detach()
+                    ibot_center_drift = torch.norm(ibot_patch_loss.center - prev_ibot_center, p=2).detach()
+                    running_dino_center_drift.update(dino_center_drift)
+                    running_ibot_center_drift.update(ibot_center_drift)
+
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -819,9 +847,13 @@ def train(args: argparse.Namespace) -> None:
                 running_loss_dino_global.synchronize_between_processes(device)
                 running_loss_koleo.synchronize_between_processes(device)
                 running_loss_ibot_patch.synchronize_between_processes(device)
-                if track_agreement is True:
+                if track_extended_metrics is True:
                     train_proto_agreement.synchronize_between_processes(device)
                     train_patch_agreement.synchronize_between_processes(device)
+                    running_target_entropy.synchronize_between_processes(device)
+                    if args.centering == "centering":
+                        running_dino_center_drift.synchronize_between_processes(device)
+                        running_ibot_center_drift.synchronize_between_processes(device)
 
                 with training_utils.single_handler_logging(logger, file_handler, enabled=not disable_tqdm) as log:
                     log.info(
@@ -846,13 +878,19 @@ def train(args: argparse.Namespace) -> None:
                         },
                         ((epoch - 1) * epoch_samples) + ((i + 1) * batch_size * args.world_size),
                     )
-                    if track_agreement is True:
+                    if track_extended_metrics is True:
+                        metrics = {
+                            "prototype_agreement": train_proto_agreement.avg,
+                            "patch_agreement": train_patch_agreement.avg,
+                            "target_entropy": running_target_entropy.avg,
+                        }
+                        if args.centering == "centering":
+                            metrics["dino_center_drift"] = running_dino_center_drift.avg
+                            metrics["ibot_center_drift"] = running_ibot_center_drift.avg
+
                         summary_writer.add_scalars(
                             "performance",
-                            {
-                                "prototype_agreement": train_proto_agreement.avg,
-                                "patch_agreement": train_patch_agreement.avg,
-                            },
+                            metrics,
                             ((epoch - 1) * epoch_samples) + ((i + 1) * batch_size * args.world_size),
                         )
 
@@ -867,9 +905,17 @@ def train(args: argparse.Namespace) -> None:
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} dino_global_loss: {running_loss_dino_global.global_avg:.4f}")
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} koleo_loss: {running_loss_koleo.global_avg:.4f}")
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} ibot_patch_loss: {running_loss_ibot_patch.global_avg:.4f}")
-        if track_agreement is True:
+        if track_extended_metrics is True:
             logger.info(f"[Trn] Epoch {epoch}/{epochs-1} prototype_agreement: {train_proto_agreement.global_avg:.4f}")
             logger.info(f"[Trn] Epoch {epoch}/{epochs-1} patch_agreement: {train_patch_agreement.global_avg:.4f}")
+            logger.info(f"[Trn] Epoch {epoch}/{epochs-1} target_entropy: {running_target_entropy.global_avg:.4f}")
+            if args.centering == "centering":
+                logger.info(
+                    f"[Trn] Epoch {epoch}/{epochs-1} dino_center_drift: {running_dino_center_drift.global_avg:.4f}"
+                )
+                logger.info(
+                    f"[Trn] Epoch {epoch}/{epochs-1} ibot_center_drift: {running_ibot_center_drift.global_avg:.4f}"
+                )
 
         # Learning rate scheduler update
         if step_update is False:
@@ -976,6 +1022,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -1042,9 +1089,11 @@ def get_args_parser() -> argparse.ArgumentParser:
         help="number of initial epochs to disable Sinkhorn queueing",
     )
     parser.add_argument(
-        "--no-agreement-metrics", default=False, action="store_true", help="disable prototype/patch agreement tracking"
+        "--no-extended-metrics",
+        default=False,
+        action="store_true",
+        help="disable extended metrics (prototype/patch agreement, target entropy, center drift)",
     )
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     training_cli.add_optimization_args(parser)
     training_cli.add_lr_wd_args(parser, wd_end=True)
     training_cli.add_lr_scheduler_args(parser)

birder/scripts/train_dino_v2_dist.py

@@ -603,23 +603,27 @@ def train(args: argparse.Namespace) -> None:
     #
     # Training loop
     #
-    track_agreement = not args.no_agreement_metrics
+    track_extended_metrics = not args.no_extended_metrics
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+    running_loss_dino_local = training_utils.SmoothedValue()
+    running_loss_dino_global = training_utils.SmoothedValue()
+    running_loss_koleo = training_utils.SmoothedValue()
+    running_loss_ibot_patch = training_utils.SmoothedValue()
+    if track_extended_metrics is True:
+        train_proto_agreement = training_utils.SmoothedValue()
+        train_patch_agreement = training_utils.SmoothedValue()
+        running_target_entropy = training_utils.SmoothedValue()
+        running_dino_center_drift = training_utils.SmoothedValue()
+        running_ibot_center_drift = training_utils.SmoothedValue()
+
     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()
         teacher.eval()
-        running_loss = training_utils.SmoothedValue()
-        running_loss_dino_local = training_utils.SmoothedValue()
-        running_loss_dino_global = training_utils.SmoothedValue()
-        running_loss_koleo = training_utils.SmoothedValue()
-        running_loss_ibot_patch = training_utils.SmoothedValue()
-        if track_agreement is True:
-            train_proto_agreement = training_utils.SmoothedValue()
-            train_patch_agreement = training_utils.SmoothedValue()
 
         if args.sinkhorn_queue_size is not None:
             queue_active = epoch > args.sinkhorn_queue_warmup_epochs
@@ -683,6 +687,11 @@ def train(args: argparse.Namespace) -> None:
                     )
                     teacher_patch_tokens_raw = teacher_masked_patch_tokens_after_head
                     if args.centering == "centering":
+                        # Track centers before update for drift computation
+                        if track_extended_metrics is True:
+                            prev_dino_center = dino_loss.center.clone()
+                            prev_ibot_center = ibot_patch_loss.center.clone()
+
                         teacher_dino_softmax_centered_list = dino_loss.softmax_center_teacher(
                             teacher_embedding_after_head, teacher_temp=teacher_temp
                         ).view(n_global_crops, -1, *teacher_embedding_after_head.shape[1:])
@@ -809,7 +818,7 @@ def train(args: argparse.Namespace) -> None:
             running_loss_koleo.update(loss_koleo.detach())
             running_loss_ibot_patch.update(loss_ibot_patch.detach())
 
-            if track_agreement is True:
+            if track_extended_metrics is True:
                 probs_teacher = teacher_embedding_after_head.chunk(n_global_crops)
                 probs_student = student_global_embedding_after_head.chunk(n_global_crops)
                 pred_teacher = probs_teacher[0].argmax(dim=1)
@@ -820,8 +829,27 @@ def train(args: argparse.Namespace) -> None:
                 pred_patch_student = student_global_masked_patch_tokens_after_head.argmax(dim=1)
                 train_patch_agreement.update(training_utils.accuracy(pred_patch_teacher, pred_patch_student))
 
+                with torch.no_grad():
+                    p = teacher_dino_softmax_centered_list.detach()
+                    p = p.reshape(-1, p.size(-1))  # (N, D)
+
+                    # Mean distribution over prototypes (marginal)
+                    m = p.mean(dim=0).clamp_min(1e-12)
+
+                    # Entropy of the marginal
+                    entropy = -(m * m.log()).sum()
+
+                running_target_entropy.update(entropy.detach())
+
+                # Compute center drift
+                if args.centering == "centering":
+                    dino_center_drift = torch.norm(dino_loss.center - prev_dino_center, p=2).detach()
+                    ibot_center_drift = torch.norm(ibot_patch_loss.center - prev_ibot_center, p=2).detach()
+                    running_dino_center_drift.update(dino_center_drift)
+                    running_ibot_center_drift.update(ibot_center_drift)
+
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -840,9 +868,13 @@ def train(args: argparse.Namespace) -> None:
                 running_loss_dino_global.synchronize_between_processes(device)
                 running_loss_koleo.synchronize_between_processes(device)
                 running_loss_ibot_patch.synchronize_between_processes(device)
-                if track_agreement is True:
+                if track_extended_metrics is True:
                     train_proto_agreement.synchronize_between_processes(device)
                     train_patch_agreement.synchronize_between_processes(device)
+                    running_target_entropy.synchronize_between_processes(device)
+                    if args.centering == "centering":
+                        running_dino_center_drift.synchronize_between_processes(device)
+                        running_ibot_center_drift.synchronize_between_processes(device)
 
                 with training_utils.single_handler_logging(logger, file_handler, enabled=not disable_tqdm) as log:
                     log.info(
@@ -867,13 +899,19 @@ def train(args: argparse.Namespace) -> None:
                         },
                         ((epoch - 1) * epoch_samples) + ((i + 1) * batch_size * args.world_size),
                     )
-                    if track_agreement is True:
+                    if track_extended_metrics is True:
+                        metrics = {
+                            "prototype_agreement": train_proto_agreement.avg,
+                            "patch_agreement": train_patch_agreement.avg,
+                            "target_entropy": running_target_entropy.avg,
+                        }
+                        if args.centering == "centering":
+                            metrics["dino_center_drift"] = running_dino_center_drift.avg
+                            metrics["ibot_center_drift"] = running_ibot_center_drift.avg
+
                         summary_writer.add_scalars(
                             "performance",
-                            {
-                                "prototype_agreement": train_proto_agreement.avg,
-                                "patch_agreement": train_patch_agreement.avg,
-                            },
+                            metrics,
                             ((epoch - 1) * epoch_samples) + ((i + 1) * batch_size * args.world_size),
                         )
 
@@ -888,9 +926,17 @@ def train(args: argparse.Namespace) -> None:
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} dino_global_loss: {running_loss_dino_global.global_avg:.4f}")
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} koleo_loss: {running_loss_koleo.global_avg:.4f}")
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} ibot_patch_loss: {running_loss_ibot_patch.global_avg:.4f}")
-        if track_agreement is True:
+        if track_extended_metrics is True:
             logger.info(f"[Trn] Epoch {epoch}/{epochs-1} prototype_agreement: {train_proto_agreement.global_avg:.4f}")
             logger.info(f"[Trn] Epoch {epoch}/{epochs-1} patch_agreement: {train_patch_agreement.global_avg:.4f}")
+            logger.info(f"[Trn] Epoch {epoch}/{epochs-1} target_entropy: {running_target_entropy.global_avg:.4f}")
+            if args.centering == "centering":
+                logger.info(
+                    f"[Trn] Epoch {epoch}/{epochs-1} dino_center_drift: {running_dino_center_drift.global_avg:.4f}"
+                )
+                logger.info(
+                    f"[Trn] Epoch {epoch}/{epochs-1} ibot_center_drift: {running_ibot_center_drift.global_avg:.4f}"
+                )
 
         # Learning rate scheduler update
         if step_update is False:
@@ -998,6 +1044,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -1006,8 +1053,8 @@ def get_args_parser() -> argparse.ArgumentParser:
             "('drop_path_rate=0.2' or '{\"units\": [3, 24, 36, 3], \"dropout\": 0.2}'"
         ),
     )
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument("--teacher", type=str, help="the neural network to use as teacher")
+    parser.add_argument("--teacher-tag", type=str, help="teacher training logs tag")
     parser.add_argument(
         "--teacher-model-config",
         action=cli.FlexibleDictAction,
@@ -1016,7 +1063,6 @@ def get_args_parser() -> argparse.ArgumentParser:
             "('drop_path_rate=0.2' or '{\"units\": [3, 24, 36, 3], \"dropout\": 0.2}'"
         ),
     )
-    parser.add_argument("--teacher-tag", type=str, help="teacher training logs tag")
     parser.add_argument("--teacher-epoch", type=int, metavar="N", help="load teacher weights from selected epoch")
     parser.add_argument("--dino-loss-weight", type=float, default=1.0, help="weight for the DINO loss component")
     parser.add_argument("--dino-out-dim", type=int, default=65536, help="dimensionality of the DINO head output")
@@ -1070,7 +1116,10 @@ def get_args_parser() -> argparse.ArgumentParser:
         help="number of initial epochs to disable Sinkhorn queueing",
     )
     parser.add_argument(
-        "--no-agreement-metrics", default=False, action="store_true", help="disable prototype/patch agreement tracking"
+        "--no-extended-metrics",
+        default=False,
+        action="store_true",
+        help="disable extended metrics (prototype/patch agreement, target entropy, center drift)",
     )
     training_cli.add_optimization_args(parser)
     training_cli.add_lr_wd_args(parser, wd_end=True)

birder/scripts/train_franca.py

@@ -612,15 +612,16 @@ def train(args: argparse.Namespace) -> None:
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+    running_loss_dino_local = training_utils.SmoothedValue()
+    running_loss_dino_global = training_utils.SmoothedValue()
+    running_loss_koleo = training_utils.SmoothedValue()
+    running_loss_ibot_patch = training_utils.SmoothedValue()
+
     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()
-        running_loss_dino_local = training_utils.SmoothedValue()
-        running_loss_dino_global = training_utils.SmoothedValue()
-        running_loss_koleo = training_utils.SmoothedValue()
-        running_loss_ibot_patch = training_utils.SmoothedValue()
 
         if args.sinkhorn_queue_size is not None:
             queue_active = epoch > args.sinkhorn_queue_warmup_epochs
@@ -804,7 +805,7 @@ def train(args: argparse.Namespace) -> None:
             running_loss_ibot_patch.update(loss_ibot_patch.detach())
 
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -963,6 +964,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -1024,7 +1026,6 @@ def get_args_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--local-crop-size", type=int, nargs="+", default=[96, 96], metavar=("H", "W"), help="local view size"
     )
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     training_cli.add_optimization_args(parser)
     training_cli.add_lr_wd_args(parser, wd_end=True)
     training_cli.add_lr_scheduler_args(parser)

birder/scripts/train_i_jepa.py

@@ -433,11 +433,12 @@ def train(args: argparse.Namespace) -> None:
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+
     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()
 
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)
@@ -534,7 +535,7 @@ def train(args: argparse.Namespace) -> None:
             running_loss.update(loss.detach())
 
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -677,6 +678,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -688,7 +690,6 @@ def get_args_parser() -> argparse.ArgumentParser:
     parser.add_argument("--predictor-embed-dim", type=int, default=384, help="predictor embedding dimension")
     parser.add_argument("--predictor-num-heads", type=int, default=12, help="predictor number of heads")
     parser.add_argument("--predictor-depth", type=int, default=12, help="predictor number of layers")
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     training_cli.add_optimization_args(parser)
     training_cli.add_lr_wd_args(parser, wd_end=True)
     training_cli.add_lr_scheduler_args(parser)

birder/scripts/train_ibot.py

@@ -499,12 +499,13 @@ def train(args: argparse.Namespace) -> None:
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    running_loss = training_utils.SmoothedValue()
+    train_proto_agreement = training_utils.SmoothedValue()
+
     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()
-        train_proto_agreement = training_utils.SmoothedValue()
 
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)
@@ -617,7 +618,7 @@ def train(args: argparse.Namespace) -> None:
             train_proto_agreement.update(training_utils.accuracy(pred_teacher, pred_student))
 
             # Write statistics
-            if i % args.log_interval == 0 or i == last_batch_idx:
+            if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
                 time_now = time.time()
                 time_cost = time_now - start_time
                 iters_processed_in_interval = i - last_idx
@@ -774,6 +775,7 @@ def get_args_parser() -> argparse.ArgumentParser:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("-n", "--network", type=str, help="the neural network to use")
+    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument(
         "--model-config",
         action=cli.FlexibleDictAction,
@@ -832,7 +834,6 @@ def get_args_parser() -> argparse.ArgumentParser:
             "try increasing this value if the loss does not decrease"
         ),
     )
-    parser.add_argument("-t", "--tag", type=str, help="add model tag")
     training_cli.add_optimization_args(parser)
     training_cli.add_lr_wd_args(parser, wd_end=True)
     training_cli.add_lr_scheduler_args(parser)