birder 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

birder/scripts/train_data2vec2.py

@@ -75,7 +75,7 @@ def train(args: argparse.Namespace) -> None:
     #
     # Initialize
     #
-    (device, device_id, disable_tqdm) = training_utils.init_training(args, logger)
+    device, device_id, disable_tqdm = training_utils.init_training(args, logger)
 
     if args.size is None:
         # Prefer mim size over encoder default size
@@ -105,7 +105,7 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("data2vec2", encoder=args.network, tag=args.tag)
 
-    backbone = registry.net_factory(args.network, sample_shape[1], 0, config=args.model_config, size=args.size)
+    backbone = registry.net_factory(args.network, 0, sample_shape[1], config=args.model_config, size=args.size)
     net = Data2Vec2(
         backbone,
         config={
@@ -121,7 +121,7 @@ def train(args: argparse.Namespace) -> None:
 
     if args.resume_epoch is not None:
         begin_epoch = args.resume_epoch + 1
-        (net, training_states) = fs_ops.load_simple_checkpoint(
+        net, training_states = fs_ops.load_simple_checkpoint(
             device, net, network_name, epoch=args.resume_epoch, strict=not args.non_strict_weights
         )
 
@@ -169,11 +169,11 @@ def train(args: argparse.Namespace) -> None:
     elif args.wds is True:
         wds_path: str | list[str]
         if args.wds_info is not None:
-            (wds_path, dataset_size) = wds_args_from_info(args.wds_info, args.wds_split)
+            wds_path, dataset_size = wds_args_from_info(args.wds_info, args.wds_split)
             if args.wds_size is not None:
                 dataset_size = args.wds_size
         else:
-            (wds_path, dataset_size) = prepare_wds_args(args.data_path[0], args.wds_size, device)
+            wds_path, dataset_size = prepare_wds_args(args.data_path[0], args.wds_size, device)
 
         training_dataset = make_wds_dataset(
             wds_path,
@@ -199,7 +199,7 @@ def train(args: argparse.Namespace) -> None:
 
     # Data loaders and samplers
     virtual_epoch_mode = args.steps_per_epoch is not None
-    (train_sampler, _) = training_utils.get_samplers(
+    train_sampler, _ = training_utils.get_samplers(
         args, training_dataset, validation_dataset=None, infinite=virtual_epoch_mode
     )
 
@@ -288,7 +288,7 @@ def train(args: argparse.Namespace) -> None:
     )
 
     # Gradient scaler and AMP related tasks
-    (scaler, amp_dtype) = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
+    scaler, amp_dtype = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
 
     # Load states
     if args.load_states is True:
@@ -400,6 +400,9 @@ def train(args: argparse.Namespace) -> None:
         tic = time.time()
         net.train()
 
+        # Clear metrics
+        running_loss.clear()
+
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)
 

birder/scripts/train_detection.py

@@ -27,7 +27,7 @@ from birder.common import training_cli
 from birder.common import training_utils
 from birder.conf import settings
 from birder.data.collators.detection import BatchRandomResizeCollator
-from birder.data.collators.detection import training_collate_fn
+from birder.data.collators.detection import DetectionCollator
 from birder.data.datasets.coco import CocoMosaicTraining
 from birder.data.datasets.coco import CocoTraining
 from birder.data.transforms.classification import get_rgb_stats
@@ -63,7 +63,7 @@ def train(args: argparse.Namespace) -> None:
     )
     model_dynamic_size = transform_dynamic_size or args.batch_multiscale is True
 
-    (device, device_id, disable_tqdm) = training_utils.init_training(
+    device, device_id, disable_tqdm = training_utils.init_training(
         args, logger, cudnn_dynamic_size=transform_dynamic_size
     )
 
@@ -92,6 +92,7 @@ def train(args: argparse.Namespace) -> None:
         args.multiscale,
         args.max_size,
         args.multiscale_min_size,
+        args.multiscale_step,
     )
     mosaic_dataset = None
     if args.mosaic_prob > 0.0:
@@ -104,6 +105,7 @@ def train(args: argparse.Namespace) -> None:
             args.multiscale,
             args.max_size,
             args.multiscale_min_size,
+            args.multiscale_step,
             post_mosaic=True,
         )
         if args.dynamic_size is True or args.multiscale is True:
@@ -177,14 +179,22 @@ def train(args: argparse.Namespace) -> None:
 
     # Data loaders and samplers
     virtual_epoch_mode = args.steps_per_epoch is not None
-    (train_sampler, validation_sampler) = training_utils.get_samplers(
+    train_sampler, validation_sampler = training_utils.get_samplers(
         args, training_dataset, validation_dataset, infinite=virtual_epoch_mode
     )
 
     if args.batch_multiscale is True:
-        train_collate_fn: Any = BatchRandomResizeCollator(0, args.size, multiscale_min_size=args.multiscale_min_size)
+        train_collate_fn: Any = BatchRandomResizeCollator(
+            0,
+            args.size,
+            size_divisible=args.multiscale_step,
+            multiscale_min_size=args.multiscale_min_size,
+            multiscale_step=args.multiscale_step,
+        )
     else:
-        train_collate_fn = training_collate_fn
+        train_collate_fn = DetectionCollator(0, size_divisible=args.multiscale_step)
+
+    validation_collate_fn = DetectionCollator(0, size_divisible=args.multiscale_step)
 
     training_loader = DataLoader(
         training_dataset,
@@ -202,7 +212,7 @@ def train(args: argparse.Namespace) -> None:
         sampler=validation_sampler,
         num_workers=args.num_workers,
         prefetch_factor=args.prefetch_factor,
-        collate_fn=training_collate_fn,
+        collate_fn=validation_collate_fn,
         pin_memory=True,
         drop_last=args.drop_last,
     )
@@ -243,7 +253,7 @@ def train(args: argparse.Namespace) -> None:
 
     if args.resume_epoch is not None:
         begin_epoch = args.resume_epoch + 1
-        (net, class_to_idx_saved, training_states) = fs_ops.load_detection_checkpoint(
+        net, class_to_idx_saved, training_states = fs_ops.load_detection_checkpoint(
             device,
             args.network,
             config=args.model_config,
@@ -262,7 +272,7 @@ def train(args: argparse.Namespace) -> None:
 
     elif args.pretrained is True:
         fs_ops.download_model_by_weights(network_name, progress_bar=training_utils.is_local_primary(args))
-        (net, class_to_idx_saved, training_states) = fs_ops.load_detection_checkpoint(
+        net, class_to_idx_saved, training_states = fs_ops.load_detection_checkpoint(
             device,
             args.network,
             config=args.model_config,
@@ -282,7 +292,7 @@ def train(args: argparse.Namespace) -> None:
     else:
         if args.backbone_epoch is not None:
             backbone: DetectorBackbone
-            (backbone, class_to_idx_saved, _) = fs_ops.load_checkpoint(
+            backbone, class_to_idx_saved, _ = fs_ops.load_checkpoint(
                 device,
                 args.backbone,
                 config=args.backbone_model_config,
@@ -297,7 +307,7 @@ def train(args: argparse.Namespace) -> None:
                 lib.get_network_name(args.backbone, tag=args.backbone_tag),
                 progress_bar=training_utils.is_local_primary(args),
             )
-            (backbone, class_to_idx_saved, _) = fs_ops.load_checkpoint(
+            backbone, class_to_idx_saved, _ = fs_ops.load_checkpoint(
                 device,
                 args.backbone,
                 config=args.backbone_model_config,
@@ -309,7 +319,7 @@ def train(args: argparse.Namespace) -> None:
 
         else:
             backbone = registry.net_factory(
-                args.backbone, sample_shape[1], num_outputs, config=args.backbone_model_config, size=args.size
+                args.backbone, num_outputs, sample_shape[1], config=args.backbone_model_config, size=args.size
             )
 
         net = registry.detection_net_factory(
@@ -386,7 +396,7 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
     # Gradient scaler and AMP related tasks
-    (scaler, amp_dtype) = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
+    scaler, amp_dtype = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
 
     # Load states
     if args.load_states is True:
@@ -546,6 +556,11 @@ def train(args: argparse.Namespace) -> None:
         tic = time.time()
         net.train()
 
+        # Clear metrics
+        running_loss.clear()
+        for tracker in loss_trackers.values():
+            tracker.clear()
+
         validation_metrics.reset()
 
         if args.distributed is True or virtual_epoch_mode is True:
@@ -586,7 +601,7 @@ def train(args: argparse.Namespace) -> None:
 
             # Forward, backward and optimize
             with torch.amp.autocast("cuda", enabled=args.amp, dtype=amp_dtype):
-                (_detections, losses) = net(inputs, targets, masks, image_sizes)
+                _detections, losses = net(inputs, targets, masks, image_sizes)
                 loss = sum(v for v in losses.values())
 
             if scaler is not None:
@@ -708,7 +723,7 @@ def train(args: argparse.Namespace) -> None:
                 masks = masks.to(device, non_blocking=True)
 
                 with torch.amp.autocast("cuda", enabled=args.amp, dtype=amp_dtype):
-                    (detections, losses) = eval_model(inputs, masks=masks, image_sizes=image_sizes)
+                    detections, losses = eval_model(inputs, masks=masks, image_sizes=image_sizes)
 
                 for target in targets:
                     # TorchMetrics can't handle "empty" images

birder/scripts/train_dino_v1.py

@@ -101,7 +101,7 @@ def train(args: argparse.Namespace) -> None:
     #
     # Initialize
     #
-    (device, device_id, disable_tqdm) = training_utils.init_training(args, logger)
+    device, device_id, disable_tqdm = training_utils.init_training(args, logger)
 
     if args.size is None:
         args.size = registry.get_default_size(args.network)
@@ -129,11 +129,11 @@ def train(args: argparse.Namespace) -> None:
     elif args.wds is True:
         wds_path: str | list[str]
         if args.wds_info is not None:
-            (wds_path, dataset_size) = wds_args_from_info(args.wds_info, args.wds_split)
+            wds_path, dataset_size = wds_args_from_info(args.wds_info, args.wds_split)
             if args.wds_size is not None:
                 dataset_size = args.wds_size
         else:
-            (wds_path, dataset_size) = prepare_wds_args(args.data_path[0], args.wds_size, device)
+            wds_path, dataset_size = prepare_wds_args(args.data_path[0], args.wds_size, device)
 
         training_dataset = make_wds_dataset(
             wds_path,
@@ -163,7 +163,7 @@ def train(args: argparse.Namespace) -> None:
 
     # Data loaders and samplers
     virtual_epoch_mode = args.steps_per_epoch is not None
-    (train_sampler, _) = training_utils.get_samplers(
+    train_sampler, _ = training_utils.get_samplers(
         args, training_dataset, validation_dataset=None, infinite=virtual_epoch_mode
     )
 
@@ -226,9 +226,9 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("dino_v1", encoder=args.network, tag=args.tag)
 
-    student_backbone = registry.net_factory(args.network, sample_shape[1], 0, config=args.model_config, size=args.size)
+    student_backbone = registry.net_factory(args.network, 0, sample_shape[1], config=args.model_config, size=args.size)
     if args.backbone_epoch is not None:
-        (student_backbone, _) = fs_ops.load_simple_checkpoint(
+        student_backbone, _ = fs_ops.load_simple_checkpoint(
             device, student_backbone, backbone_name, epoch=args.backbone_epoch, strict=not args.non_strict_weights
         )
 
@@ -239,7 +239,7 @@ def train(args: argparse.Namespace) -> None:
         teacher_model_config = {"drop_path_rate": 0.0}
 
     teacher_backbone = registry.net_factory(
-        args.network, sample_shape[1], 0, config=teacher_model_config, size=args.size
+        args.network, 0, sample_shape[1], config=teacher_model_config, size=args.size
     )
     if args.freeze_body is True:
         student_backbone.freeze(freeze_classifier=False, unfreeze_features=True)
@@ -293,7 +293,7 @@ def train(args: argparse.Namespace) -> None:
 
     if args.resume_epoch is not None:
         begin_epoch = args.resume_epoch + 1
-        (net, training_states) = fs_ops.load_simple_checkpoint(
+        net, training_states = fs_ops.load_simple_checkpoint(
             device, net, network_name, epoch=args.resume_epoch, strict=not args.non_strict_weights
         )
         student = net["student"]
@@ -368,7 +368,7 @@ def train(args: argparse.Namespace) -> None:
         wd_schedule = None
 
     # Gradient scaler and AMP related tasks
-    (scaler, amp_dtype) = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
+    scaler, amp_dtype = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
 
     # Load states
     if args.load_states is True:
@@ -488,6 +488,10 @@ def train(args: argparse.Namespace) -> None:
         tic = time.time()
         net.train()
 
+        # Clear metrics
+        running_loss.clear()
+        train_proto_agreement.clear()
+
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)
 

birder/scripts/train_dino_v2.py

@@ -178,7 +178,7 @@ def train(args: argparse.Namespace) -> None:
     #
     # Initialize
     #
-    (device, device_id, disable_tqdm) = training_utils.init_training(args, logger)
+    device, device_id, disable_tqdm = training_utils.init_training(args, logger)
 
     if args.size is None:
         args.size = registry.get_default_size(args.network)
@@ -207,7 +207,7 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("dino_v2", encoder=args.network, tag=args.tag)
 
-    student_backbone = registry.net_factory(args.network, sample_shape[1], 0, config=args.model_config, size=args.size)
+    student_backbone = registry.net_factory(args.network, 0, sample_shape[1], config=args.model_config, size=args.size)
     if args.model_config is not None:
         teacher_model_config = args.model_config.copy()
         teacher_model_config.update({"drop_path_rate": 0.0})
@@ -215,7 +215,7 @@ def train(args: argparse.Namespace) -> None:
         teacher_model_config = {"drop_path_rate": 0.0}
 
     teacher_backbone = registry.net_factory(
-        args.network, sample_shape[1], 0, config=teacher_model_config, size=args.size
+        args.network, 0, sample_shape[1], config=teacher_model_config, size=args.size
     )
     student_backbone.set_dynamic_size()
     if args.ibot_separate_head is False:
@@ -267,7 +267,7 @@ def train(args: argparse.Namespace) -> None:
 
     if args.resume_epoch is not None:
         begin_epoch = args.resume_epoch + 1
-        (net, training_states) = fs_ops.load_simple_checkpoint(
+        net, training_states = fs_ops.load_simple_checkpoint(
             device, net, network_name, epoch=args.resume_epoch, strict=not args.non_strict_weights
         )
         student = net["student"]
@@ -336,11 +336,11 @@ def train(args: argparse.Namespace) -> None:
     elif args.wds is True:
         wds_path: str | list[str]
         if args.wds_info is not None:
-            (wds_path, dataset_size) = wds_args_from_info(args.wds_info, args.wds_split)
+            wds_path, dataset_size = wds_args_from_info(args.wds_info, args.wds_split)
             if args.wds_size is not None:
                 dataset_size = args.wds_size
         else:
-            (wds_path, dataset_size) = prepare_wds_args(args.data_path[0], args.wds_size, device)
+            wds_path, dataset_size = prepare_wds_args(args.data_path[0], args.wds_size, device)
 
         training_dataset = make_wds_dataset(
             wds_path,
@@ -366,7 +366,7 @@ def train(args: argparse.Namespace) -> None:
 
     # Data loaders and samplers
     virtual_epoch_mode = args.steps_per_epoch is not None
-    (train_sampler, _) = training_utils.get_samplers(
+    train_sampler, _ = training_utils.get_samplers(
         args, training_dataset, validation_dataset=None, infinite=virtual_epoch_mode
     )
 
@@ -466,7 +466,7 @@ def train(args: argparse.Namespace) -> None:
         wd_schedule = None
 
     # Gradient scaler and AMP related tasks
-    (scaler, amp_dtype) = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
+    scaler, amp_dtype = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
 
     # Load states
     if args.load_states is True:
@@ -603,6 +603,19 @@ def train(args: argparse.Namespace) -> None:
         tic = time.time()
         net.train()
 
+        # Clear metrics
+        running_loss.clear()
+        running_loss_dino_local.clear()
+        running_loss_dino_global.clear()
+        running_loss_koleo.clear()
+        running_loss_ibot_patch.clear()
+        if track_extended_metrics is True:
+            train_proto_agreement.clear()
+            train_patch_agreement.clear()
+            running_target_entropy.clear()
+            running_dino_center_drift.clear()
+            running_ibot_center_drift.clear()
+
         if args.sinkhorn_queue_size is not None:
             queue_active = epoch > args.sinkhorn_queue_warmup_epochs
             dino_loss.set_queue_active(queue_active)
@@ -661,7 +674,7 @@ def train(args: argparse.Namespace) -> None:
             with torch.amp.autocast("cuda", enabled=args.amp, dtype=amp_dtype):
                 with torch.no_grad():
                     # Teacher
-                    (teacher_embedding_after_head, teacher_masked_patch_tokens_after_head) = teacher(
+                    teacher_embedding_after_head, teacher_masked_patch_tokens_after_head = teacher(
                         global_crops, n_global_crops, upper_bound, mask_indices_list
                     )
                     teacher_patch_tokens_raw = teacher_masked_patch_tokens_after_head
@@ -671,7 +684,7 @@ def train(args: argparse.Namespace) -> None:
                             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_dino_softmax_centered = dino_loss.softmax_center_teacher(
                             teacher_embedding_after_head, teacher_temp=teacher_temp
                         ).view(n_global_crops, -1, *teacher_embedding_after_head.shape[1:])
                         dino_loss.update_center(teacher_embedding_after_head)
@@ -684,7 +697,7 @@ def train(args: argparse.Namespace) -> None:
                         ibot_patch_loss.update_center(teacher_masked_patch_tokens_after_head[:, :n_masked_patches])
 
                     else:  # sinkhorn_knopp
-                        teacher_dino_softmax_centered_list = dino_loss.sinkhorn_knopp_teacher(
+                        teacher_dino_softmax_centered = dino_loss.sinkhorn_knopp_teacher(
                             teacher_embedding_after_head, teacher_temp=teacher_temp
                         ).view(n_global_crops, -1, *teacher_embedding_after_head.shape[1:])
 
@@ -705,7 +718,7 @@ def train(args: argparse.Namespace) -> None:
                 # Local DINO loss
                 loss_dino_local_crops = dino_loss(
                     student_local_embedding_after_head.chunk(n_local_crops),
-                    teacher_dino_softmax_centered_list,
+                    teacher_dino_softmax_centered.unbind(0),
                 ) / (n_global_crops_loss_terms + n_local_crops_loss_terms)
                 loss = args.dino_loss_weight * loss_dino_local_crops
 
@@ -715,7 +728,7 @@ def train(args: argparse.Namespace) -> None:
                     dino_loss(
                         [student_global_embedding_after_head],
                         [
-                            teacher_dino_softmax_centered_list.flatten(0, 1)
+                            teacher_dino_softmax_centered.flatten(0, 1)
                         ],  # These were chunked and stacked in reverse so A is matched to B
                     )
                     * loss_scales
@@ -809,7 +822,7 @@ def train(args: argparse.Namespace) -> None:
                 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 = teacher_dino_softmax_centered.detach()
                     p = p.reshape(-1, p.size(-1))  # (N, D)
 
                     # Mean distribution over prototypes (marginal)

birder/scripts/train_dino_v2_dist.py

@@ -179,7 +179,7 @@ def train(args: argparse.Namespace) -> None:
     #
     # Initialize
     #
-    (device, device_id, disable_tqdm) = training_utils.init_training(args, logger)
+    device, device_id, disable_tqdm = training_utils.init_training(args, logger)
 
     if args.size is None:
         args.size = registry.get_default_size(args.network)
@@ -208,17 +208,17 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("dino_v2_dist", encoder=args.network, tag=args.tag)
 
-    student_backbone = registry.net_factory(args.network, sample_shape[1], 0, config=args.model_config, size=args.size)
+    student_backbone = registry.net_factory(args.network, 0, sample_shape[1], config=args.model_config, size=args.size)
     student_backbone_ema = registry.net_factory(
-        args.network, sample_shape[1], 0, config=args.model_config, size=args.size
+        args.network, 0, sample_shape[1], 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,
-        sample_shape[1],
         0,
+        sample_shape[1],
         config=args.teacher_model_config,
         size=args.size,
     )
@@ -277,7 +277,7 @@ def train(args: argparse.Namespace) -> None:
 
     if args.resume_epoch is not None:
         begin_epoch = args.resume_epoch + 1
-        (net, training_states) = fs_ops.load_simple_checkpoint(
+        net, training_states = fs_ops.load_simple_checkpoint(
             device, net, network_name, epoch=args.resume_epoch, strict=not args.non_strict_weights
         )
         student = net["student"]
@@ -358,11 +358,11 @@ def train(args: argparse.Namespace) -> None:
     elif args.wds is True:
         wds_path: str | list[str]
         if args.wds_info is not None:
-            (wds_path, dataset_size) = wds_args_from_info(args.wds_info, args.wds_split)
+            wds_path, dataset_size = wds_args_from_info(args.wds_info, args.wds_split)
             if args.wds_size is not None:
                 dataset_size = args.wds_size
         else:
-            (wds_path, dataset_size) = prepare_wds_args(args.data_path[0], args.wds_size, device)
+            wds_path, dataset_size = prepare_wds_args(args.data_path[0], args.wds_size, device)
 
         training_dataset = make_wds_dataset(
             wds_path,
@@ -388,7 +388,7 @@ def train(args: argparse.Namespace) -> None:
 
     # Data loaders and samplers
     virtual_epoch_mode = args.steps_per_epoch is not None
-    (train_sampler, _) = training_utils.get_samplers(
+    train_sampler, _ = training_utils.get_samplers(
         args, training_dataset, validation_dataset=None, infinite=virtual_epoch_mode
     )
 
@@ -487,7 +487,7 @@ def train(args: argparse.Namespace) -> None:
         wd_schedule = None
 
     # Gradient scaler and AMP related tasks
-    (scaler, amp_dtype) = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
+    scaler, amp_dtype = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
 
     # Load states
     if args.load_states is True:
@@ -625,6 +625,19 @@ def train(args: argparse.Namespace) -> None:
         net.train()
         teacher.eval()
 
+        # Clear metrics
+        running_loss.clear()
+        running_loss_dino_local.clear()
+        running_loss_dino_global.clear()
+        running_loss_koleo.clear()
+        running_loss_ibot_patch.clear()
+        if track_extended_metrics is True:
+            train_proto_agreement.clear()
+            train_patch_agreement.clear()
+            running_target_entropy.clear()
+            running_dino_center_drift.clear()
+            running_ibot_center_drift.clear()
+
         if args.sinkhorn_queue_size is not None:
             queue_active = epoch > args.sinkhorn_queue_warmup_epochs
             dino_loss.set_queue_active(queue_active)
@@ -682,7 +695,7 @@ def train(args: argparse.Namespace) -> None:
             with torch.amp.autocast("cuda", enabled=args.amp, dtype=amp_dtype):
                 with torch.no_grad():
                     # Teacher
-                    (teacher_embedding_after_head, teacher_masked_patch_tokens_after_head) = teacher(
+                    teacher_embedding_after_head, teacher_masked_patch_tokens_after_head = teacher(
                         global_crops, n_global_crops, upper_bound, mask_indices_list
                     )
                     teacher_patch_tokens_raw = teacher_masked_patch_tokens_after_head
@@ -692,7 +705,7 @@ def train(args: argparse.Namespace) -> None:
                             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_dino_softmax_centered = dino_loss.softmax_center_teacher(
                             teacher_embedding_after_head, teacher_temp=teacher_temp
                         ).view(n_global_crops, -1, *teacher_embedding_after_head.shape[1:])
                         dino_loss.update_center(teacher_embedding_after_head)
@@ -705,7 +718,7 @@ def train(args: argparse.Namespace) -> None:
                         ibot_patch_loss.update_center(teacher_masked_patch_tokens_after_head[:, :n_masked_patches])
 
                     else:  # sinkhorn_knopp
-                        teacher_dino_softmax_centered_list = dino_loss.sinkhorn_knopp_teacher(
+                        teacher_dino_softmax_centered = dino_loss.sinkhorn_knopp_teacher(
                             teacher_embedding_after_head, teacher_temp=teacher_temp
                         ).view(n_global_crops, -1, *teacher_embedding_after_head.shape[1:])
 
@@ -726,7 +739,7 @@ def train(args: argparse.Namespace) -> None:
                 # Local DINO loss
                 loss_dino_local_crops = dino_loss(
                     student_local_embedding_after_head.chunk(n_local_crops),
-                    teacher_dino_softmax_centered_list,
+                    teacher_dino_softmax_centered.unbind(0),
                 ) / (n_global_crops_loss_terms + n_local_crops_loss_terms)
                 loss = args.dino_loss_weight * loss_dino_local_crops
 
@@ -736,7 +749,7 @@ def train(args: argparse.Namespace) -> None:
                     dino_loss(
                         [student_global_embedding_after_head],
                         [
-                            teacher_dino_softmax_centered_list.flatten(0, 1)
+                            teacher_dino_softmax_centered.flatten(0, 1)
                         ],  # These were chunked and stacked in reverse so A is matched to B
                     )
                     * loss_scales
@@ -830,7 +843,7 @@ def train(args: argparse.Namespace) -> None:
                 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 = teacher_dino_softmax_centered.detach()
                     p = p.reshape(-1, p.size(-1))  # (N, D)
 
                     # Mean distribution over prototypes (marginal)

birder/scripts/train_franca.py

@@ -205,7 +205,7 @@ def train(args: argparse.Namespace) -> None:
     #
     # Initialize
     #
-    (device, device_id, disable_tqdm) = training_utils.init_training(args, logger)
+    device, device_id, disable_tqdm = training_utils.init_training(args, logger)
 
     if args.size is None:
         args.size = registry.get_default_size(args.network)
@@ -234,7 +234,7 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("franca", encoder=args.network, tag=args.tag)
 
-    student_backbone = registry.net_factory(args.network, sample_shape[1], 0, config=args.model_config, size=args.size)
+    student_backbone = registry.net_factory(args.network, 0, sample_shape[1], config=args.model_config, size=args.size)
     if args.model_config is not None:
         teacher_model_config = args.model_config.copy()
         teacher_model_config.update({"drop_path_rate": 0.0})
@@ -242,7 +242,7 @@ def train(args: argparse.Namespace) -> None:
         teacher_model_config = {"drop_path_rate": 0.0}
 
     teacher_backbone = registry.net_factory(
-        args.network, sample_shape[1], 0, config=teacher_model_config, size=args.size
+        args.network, 0, sample_shape[1], config=teacher_model_config, size=args.size
     )
     student_backbone.set_dynamic_size()
     if args.ibot_separate_head is False:
@@ -296,7 +296,7 @@ def train(args: argparse.Namespace) -> None:
 
     if args.resume_epoch is not None:
         begin_epoch = args.resume_epoch + 1
-        (net, training_states) = fs_ops.load_simple_checkpoint(
+        net, training_states = fs_ops.load_simple_checkpoint(
             device, net, network_name, epoch=args.resume_epoch, strict=not args.non_strict_weights
         )
         student = net["student"]
@@ -363,11 +363,11 @@ def train(args: argparse.Namespace) -> None:
     elif args.wds is True:
         wds_path: str | list[str]
         if args.wds_info is not None:
-            (wds_path, dataset_size) = wds_args_from_info(args.wds_info, args.wds_split)
+            wds_path, dataset_size = wds_args_from_info(args.wds_info, args.wds_split)
             if args.wds_size is not None:
                 dataset_size = args.wds_size
         else:
-            (wds_path, dataset_size) = prepare_wds_args(args.data_path[0], args.wds_size, device)
+            wds_path, dataset_size = prepare_wds_args(args.data_path[0], args.wds_size, device)
 
         training_dataset = make_wds_dataset(
             wds_path,
@@ -393,7 +393,7 @@ def train(args: argparse.Namespace) -> None:
 
     # Data loaders and samplers
     virtual_epoch_mode = args.steps_per_epoch is not None
-    (train_sampler, _) = training_utils.get_samplers(
+    train_sampler, _ = training_utils.get_samplers(
         args, training_dataset, validation_dataset=None, infinite=virtual_epoch_mode
     )
 
@@ -493,7 +493,7 @@ def train(args: argparse.Namespace) -> None:
         wd_schedule = None
 
     # Gradient scaler and AMP related tasks
-    (scaler, amp_dtype) = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
+    scaler, amp_dtype = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
 
     # Load states
     if args.load_states is True:
@@ -623,6 +623,13 @@ def train(args: argparse.Namespace) -> None:
         tic = time.time()
         net.train()
 
+        # Clear metrics
+        running_loss.clear()
+        running_loss_dino_local.clear()
+        running_loss_dino_global.clear()
+        running_loss_koleo.clear()
+        running_loss_ibot_patch.clear()
+
         if args.sinkhorn_queue_size is not None:
             queue_active = epoch > args.sinkhorn_queue_warmup_epochs
             dino_loss.set_queue_active(queue_active)
@@ -681,7 +688,7 @@ def train(args: argparse.Namespace) -> None:
             with torch.amp.autocast("cuda", enabled=args.amp, dtype=amp_dtype):
                 with torch.no_grad():
                     # Teacher
-                    (teacher_embedding_after_head, teacher_masked_patch_tokens_after_head) = teacher(
+                    teacher_embedding_after_head, teacher_masked_patch_tokens_after_head = teacher(
                         global_crops, n_global_crops, upper_bound, mask_indices_list
                     )