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

birder/scripts/train_i_jepa.py

@@ -74,7 +74,7 @@ class TrainCollator:
     def __call__(self, batch: Any) -> tuple[torch.Tensor, list[torch.Tensor], list[torch.Tensor]]:
         B = len(batch)
         collated_batch = torch.utils.data.default_collate(batch)
-        (enc_masks, pred_masks) = self.mask_generator(B)
+        enc_masks, pred_masks = self.mask_generator(B)
 
         return (collated_batch, enc_masks, pred_masks)
 
@@ -84,7 +84,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)
@@ -119,9 +119,9 @@ def train(args: argparse.Namespace) -> None:
     else:
         model_config = {"drop_path_rate": 0.0}
 
-    backbone = registry.net_factory(args.network, sample_shape[1], 0, config=model_config, size=args.size)
+    backbone = registry.net_factory(args.network, 0, sample_shape[1], config=model_config, size=args.size)
     num_special_tokens = backbone.num_special_tokens
-    target_backbone = registry.net_factory(args.network, sample_shape[1], 0, config=model_config, size=args.size)
+    target_backbone = registry.net_factory(args.network, 0, sample_shape[1], config=model_config, size=args.size)
     encoder = I_JEPA(backbone)
     target_encoder = I_JEPA(target_backbone)
     target_encoder.load_state_dict(encoder.state_dict())
@@ -148,7 +148,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
         )
         encoder = net["encoder"]
@@ -198,11 +198,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,
@@ -228,7 +228,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
     )
 
@@ -320,7 +320,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:
@@ -440,6 +440,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_ibot.py

@@ -107,7 +107,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)
@@ -136,7 +136,7 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("ibot", 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})
@@ -144,7 +144,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()
     student = iBOT(
@@ -204,7 +204,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"]
@@ -266,11 +266,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,
@@ -296,7 +296,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
     )
 
@@ -387,7 +387,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:
@@ -507,6 +507,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)
 
@@ -553,12 +557,12 @@ def train(args: argparse.Namespace) -> None:
             with torch.amp.autocast("cuda", enabled=args.amp, dtype=amp_dtype):
                 # Global views
                 with torch.no_grad():
-                    (teacher_embedding, teacher_features) = teacher(torch.concat(images[:2], dim=0), None)
+                    teacher_embedding, teacher_features = teacher(torch.concat(images[:2], dim=0), None)
 
-                (student_embedding, student_features) = student(torch.concat(images[:2], dim=0), masks)
+                student_embedding, student_features = student(torch.concat(images[:2], dim=0), masks)
 
                 # Local views
-                (student_local_embedding, _) = student(torch.concat(images[2:], dim=0), None, return_keys="embedding")
+                student_local_embedding, _ = student(torch.concat(images[2:], dim=0), None, return_keys="embedding")
 
                 loss = ibot_loss(
                     student_embedding,

birder/scripts/train_kd.py

@@ -76,13 +76,13 @@ 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.type != "soft":
         args.temperature = 1.0
 
     # Using the teacher rgb values for the student
-    (teacher, (class_to_idx, signature, rgb_stats, *_)) = fs_ops.load_model(
+    teacher, (class_to_idx, signature, rgb_stats, *_) = fs_ops.load_model(
         device,
         args.teacher,
         config=args.teacher_model_config,
@@ -113,15 +113,15 @@ def train(args: argparse.Namespace) -> None:
         training_wds_path: str | list[str]
         val_wds_path: str | list[str]
         if args.wds_info is not None:
-            (training_wds_path, training_size) = wds_args_from_info(args.wds_info, args.wds_training_split)
-            (val_wds_path, val_size) = wds_args_from_info(args.wds_info, args.wds_val_split)
+            training_wds_path, training_size = wds_args_from_info(args.wds_info, args.wds_training_split)
+            val_wds_path, val_size = wds_args_from_info(args.wds_info, args.wds_val_split)
             if args.wds_train_size is not None:
                 training_size = args.wds_train_size
             if args.wds_val_size is not None:
                 val_size = args.wds_val_size
         else:
-            (training_wds_path, training_size) = prepare_wds_args(args.data_path, args.wds_train_size, device)
-            (val_wds_path, val_size) = prepare_wds_args(args.val_path, args.wds_val_size, device)
+            training_wds_path, training_size = prepare_wds_args(args.data_path, args.wds_train_size, device)
+            val_wds_path, val_size = prepare_wds_args(args.val_path, args.wds_val_size, device)
 
         training_dataset = make_wds_dataset(
             training_wds_path,
@@ -187,7 +187,7 @@ 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
     )
 
@@ -269,7 +269,7 @@ def train(args: argparse.Namespace) -> None:
 
     if args.resume_epoch is not None:
         begin_epoch = args.resume_epoch + 1
-        (student, class_to_idx_saved, training_states) = fs_ops.load_checkpoint(
+        student, class_to_idx_saved, training_states = fs_ops.load_checkpoint(
             device,
             args.student,
             config=args.student_model_config,
@@ -283,8 +283,8 @@ def train(args: argparse.Namespace) -> None:
     else:
         student = registry.net_factory(
             args.student,
-            sample_shape[1],
             num_outputs,
+            sample_shape[1],
             config=args.student_model_config,
             size=args.size,
         )
@@ -383,7 +383,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:
@@ -567,10 +567,16 @@ def train(args: argparse.Namespace) -> None:
     if virtual_epoch_mode is True:
         train_iter = iter(training_loader)
 
+    top_k = args.top_k
     running_loss = training_utils.SmoothedValue(window_size=64)
     running_val_loss = training_utils.SmoothedValue()
     train_accuracy = training_utils.SmoothedValue(window_size=64)
     val_accuracy = training_utils.SmoothedValue()
+    train_topk: Optional[training_utils.SmoothedValue] = None
+    val_topk: Optional[training_utils.SmoothedValue] = None
+    if top_k is not None:
+        train_topk = training_utils.SmoothedValue(window_size=64)
+        val_topk = training_utils.SmoothedValue()
 
     logger.info(f"Starting training with learning rate of {last_lr}")
     for epoch in range(begin_epoch, args.stop_epoch):
@@ -579,6 +585,16 @@ def train(args: argparse.Namespace) -> None:
         if embedding_projection is not None:
             embedding_projection.train()
 
+        # Clear metrics
+        running_loss.clear()
+        running_val_loss.clear()
+        train_accuracy.clear()
+        val_accuracy.clear()
+        if train_topk is not None:
+            train_topk.clear()
+        if val_topk is not None:
+            val_topk.clear()
+
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)
 
@@ -616,7 +632,7 @@ def train(args: argparse.Namespace) -> None:
                         teacher_embedding = teacher.embedding(inputs)
                         teacher_embedding = F.normalize(teacher_embedding, dim=-1)
 
-                    (outputs, student_embedding) = train_student(inputs)
+                    outputs, student_embedding = train_student(inputs)
                     student_embedding = embedding_projection(student_embedding)  # type: ignore[misc]
                     student_embedding = F.normalize(student_embedding, dim=-1)
                     dist_loss = distillation_criterion(student_embedding, teacher_embedding)
@@ -637,7 +653,7 @@ def train(args: argparse.Namespace) -> None:
                         outputs = train_student(inputs)
                         dist_loss = distillation_criterion(outputs, teacher_targets)
                     elif distillation_type == "deit":
-                        (outputs, dist_output) = torch.unbind(train_student(inputs), dim=1)
+                        outputs, dist_output = torch.unbind(train_student(inputs), dim=1)
                         dist_loss = distillation_criterion(dist_output, teacher_targets)
                     else:
                         raise RuntimeError
@@ -693,6 +709,9 @@ def train(args: argparse.Namespace) -> None:
                 targets = targets.argmax(dim=1)
 
             train_accuracy.update(training_utils.accuracy(targets, outputs.detach()))
+            if train_topk is not None:
+                topk_val = training_utils.topk_accuracy(targets, outputs.detach(), topk=(top_k,))[0]
+                train_topk.update(topk_val)
 
             # Write statistics
             if (i % args.log_interval == 0 and i > 0) or i == last_batch_idx:
@@ -711,6 +730,9 @@ def train(args: argparse.Namespace) -> None:
 
                 running_loss.synchronize_between_processes(device)
                 train_accuracy.synchronize_between_processes(device)
+                if train_topk is not None:
+                    train_topk.synchronize_between_processes(device)
+
                 with training_utils.single_handler_logging(logger, file_handler, enabled=not disable_tqdm) as log:
                     log.info(
                         f"[Trn] Epoch {epoch}/{epochs-1}, iter {i+1}/{last_batch_idx+1}  "
@@ -725,8 +747,17 @@ def train(args: argparse.Namespace) -> None:
                         f"[Trn] Epoch {epoch}/{epochs-1}, iter {i+1}/{last_batch_idx+1}  "
                         f"Accuracy: {train_accuracy.avg:.4f}"
                     )
+                    if train_topk is not None:
+                        log.info(
+                            f"[Trn] Epoch {epoch}/{epochs-1}, iter {i+1}/{last_batch_idx+1}  "
+                            f"Accuracy@{top_k}: {train_topk.avg:.4f}"
+                        )
 
                 if training_utils.is_local_primary(args) is True:
+                    performance = {"training_accuracy": train_accuracy.avg}
+                    if train_topk is not None:
+                        performance[f"training_accuracy@{top_k}"] = train_topk.avg
+
                     summary_writer.add_scalars(
                         "loss",
                         {"training": running_loss.avg},
@@ -734,7 +765,7 @@ def train(args: argparse.Namespace) -> None:
                     )
                     summary_writer.add_scalars(
                         "performance",
-                        {"training_accuracy": train_accuracy.avg},
+                        performance,
                         ((epoch - 1) * epoch_samples) + ((i + 1) * batch_size * args.world_size),
                     )
 
@@ -746,6 +777,8 @@ def train(args: argparse.Namespace) -> None:
         # Epoch training metrics
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} training_loss: {running_loss.global_avg:.4f}")
         logger.info(f"[Trn] Epoch {epoch}/{epochs-1} training_accuracy: {train_accuracy.global_avg:.4f}")
+        if train_topk is not None:
+            logger.info(f"[Trn] Epoch {epoch}/{epochs-1} training_accuracy@{top_k}: {train_topk.global_avg:.4f}")
 
         # Validation
         eval_model.eval()
@@ -772,6 +805,9 @@ def train(args: argparse.Namespace) -> None:
                 # Statistics
                 running_val_loss.update(val_loss.detach())
                 val_accuracy.update(training_utils.accuracy(targets, outputs), n=outputs.size(0))
+                if val_topk is not None:
+                    topk_val = training_utils.topk_accuracy(targets, outputs, topk=(top_k,))[0]
+                    val_topk.update(topk_val, n=outputs.size(0))
 
                 # Update progress bar
                 progress.update(n=batch_size * args.world_size)
@@ -789,19 +825,30 @@ def train(args: argparse.Namespace) -> None:
 
         running_val_loss.synchronize_between_processes(device)
         val_accuracy.synchronize_between_processes(device)
+        if val_topk is not None:
+            val_topk.synchronize_between_processes(device)
+
         epoch_val_loss = running_val_loss.global_avg
         epoch_val_accuracy = val_accuracy.global_avg
+        if val_topk is not None:
+            epoch_val_topk = val_topk.global_avg
+        else:
+            epoch_val_topk = None
 
         # Write statistics
         if training_utils.is_local_primary(args) is True:
             summary_writer.add_scalars("loss", {"validation": epoch_val_loss}, epoch * epoch_samples)
-            summary_writer.add_scalars(
-                "performance", {"validation_accuracy": epoch_val_accuracy}, epoch * epoch_samples
-            )
+            performance = {"validation_accuracy": epoch_val_accuracy}
+            if epoch_val_topk is not None:
+                performance[f"validation_accuracy@{top_k}"] = epoch_val_topk
+
+            summary_writer.add_scalars("performance", performance, epoch * epoch_samples)
 
         # Epoch validation metrics
         logger.info(f"[Val] Epoch {epoch}/{epochs-1} validation_loss (target only): {epoch_val_loss:.4f}")
         logger.info(f"[Val] Epoch {epoch}/{epochs-1} validation_accuracy: {epoch_val_accuracy:.4f}")
+        if epoch_val_topk is not None:
+            logger.info(f"[Val] Epoch {epoch}/{epochs-1} validation_accuracy@{top_k}: {epoch_val_topk:.4f}")
 
         # Learning rate scheduler update
         if step_update is False:
@@ -989,7 +1036,7 @@ def get_args_parser() -> argparse.ArgumentParser:
     training_cli.add_compile_args(parser, teacher=True)
     training_cli.add_checkpoint_args(parser, default_save_frequency=5)
     training_cli.add_distributed_args(parser)
-    training_cli.add_logging_and_debug_args(parser)
+    training_cli.add_logging_and_debug_args(parser, classification=True)
     training_cli.add_training_data_args(parser)
 
     return parser

birder/scripts/train_mim.py

@@ -49,7 +49,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
@@ -73,11 +73,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,
@@ -107,7 +107,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
     )
 
@@ -172,7 +172,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_mim_checkpoint(
+        net, training_states = fs_ops.load_mim_checkpoint(
             device,
             args.network,
             config=args.model_config,
@@ -187,7 +187,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, training_states) = fs_ops.load_mim_checkpoint(
+        net, training_states = fs_ops.load_mim_checkpoint(
             device,
             args.network,
             config=args.model_config,
@@ -202,7 +202,7 @@ def train(args: argparse.Namespace) -> None:
 
     else:
         encoder = registry.net_factory(
-            args.encoder, sample_shape[1], 0, config=args.encoder_model_config, size=args.size
+            args.encoder, 0, sample_shape[1], config=args.encoder_model_config, size=args.size
         )
         net = registry.mim_net_factory(
             args.network,
@@ -263,7 +263,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:
@@ -375,6 +375,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_mmcr.py

@@ -74,7 +74,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)
@@ -97,11 +97,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,
@@ -131,7 +131,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
     )
 
@@ -194,12 +194,12 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("mmcr", 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 = MMCR(backbone, config={"projector_dims": args.projector_dims})
 
     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
         )
 
@@ -265,7 +265,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:
@@ -377,6 +377,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)
 
@@ -407,7 +410,7 @@ def train(args: argparse.Namespace) -> None:
 
             # Forward, backward and optimize
             with torch.amp.autocast("cuda", enabled=args.amp, dtype=amp_dtype):
-                (z, z_m) = net(images)
+                z, z_m = net(images)
                 loss = mmcr_loss(z, z_m)
 
             if scaler is not None:

birder/scripts/train_rotnet.py

@@ -83,7 +83,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)
@@ -111,11 +111,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,
@@ -145,7 +145,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
     )
 
@@ -207,12 +207,12 @@ def train(args: argparse.Namespace) -> None:
         network_name = f"{network_name}-{args.tag}"
 
     net = registry.net_factory(
-        args.network, sample_shape[1], len(class_to_idx), config=args.model_config, size=args.size
+        args.network, len(class_to_idx), sample_shape[1], config=args.model_config, size=args.size
     )
 
     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
         )
 
@@ -277,7 +277,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:
@@ -389,6 +389,10 @@ def train(args: argparse.Namespace) -> None:
         tic = time.time()
         net.train()
 
+        # Clear metrics
+        running_loss.clear()
+        train_accuracy.clear()
+
         if args.distributed is True or virtual_epoch_mode is True:
             train_sampler.set_epoch(epoch)