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

birder/scripts/benchmark.py

@@ -13,6 +13,7 @@ from birder.common import cli
 from birder.conf import settings
 from birder.model_registry import Task
 from birder.model_registry import registry
+from birder.net.base import DetectorBackbone
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +28,23 @@ def prepare_model(net: torch.nn.Module) -> None:
         param.requires_grad_(False)
 
 
+def init_plain_model(
+    model_name: str, sample_shape: tuple[int, ...], device: torch.device, args: argparse.Namespace
+) -> torch.nn.Module:
+    size = (sample_shape[2], sample_shape[3])
+    input_channels = sample_shape[1]
+    if args.backbone is not None:
+        backbone = registry.net_factory(args.backbone, args.num_classes, input_channels, size=size)
+        net = registry.detection_net_factory(model_name, args.num_classes, backbone, size=size)
+    else:
+        net = registry.net_factory(model_name, args.num_classes, input_channels, size=size)
+
+    net.to(device)
+    prepare_model(net)
+
+    return net
+
+
 def throughput_benchmark(
     net: torch.nn.Module, device: torch.device, sample_shape: tuple[int, ...], model_name: str, args: argparse.Namespace
 ) -> tuple[float, int]:
@@ -110,14 +128,10 @@ def memory_benchmark(
     )
 
     if args.plain is True:
-        size = (sample_shape[2], sample_shape[3])
-        input_channels = sample_shape[1]
-        net = registry.net_factory(model_name, input_channels, 0, size=size)
-        net.to(device)
-        prepare_model(net)
+        net = init_plain_model(model_name, sample_shape, device, args)
 
     else:
-        (net, _) = birder.load_pretrained_model(model_name, inference=True, device=device)
+        net, _ = birder.load_pretrained_model(model_name, inference=True, device=device)
         if args.size is not None:
             size = (sample_shape[2], sample_shape[3])
             net.adjust_size(size)
@@ -182,7 +196,8 @@ def benchmark(args: argparse.Namespace) -> None:
     if args.plain is True:
         model_list = args.models or []
         if len(model_list) == 0:
-            model_list = registry.list_models(include_filter=args.filter, task=Task.IMAGE_CLASSIFICATION)
+            task = Task.OBJECT_DETECTION if args.backbone is not None else Task.IMAGE_CLASSIFICATION
+            model_list = registry.list_models(include_filter=args.filter, task=task)
 
     else:
         model_list = birder.list_pretrained_models(args.filter)
@@ -234,11 +249,9 @@ def benchmark(args: argparse.Namespace) -> None:
         else:
             # Initialize model
             if args.plain is True:
-                net = registry.net_factory(model_name, input_channels, 0, size=size)
-                net.to(device)
-                prepare_model(net)
+                net = init_plain_model(model_name, sample_shape, device, args)
             else:
-                (net, _) = birder.load_pretrained_model(model_name, inference=True, device=device)
+                net, _ = birder.load_pretrained_model(model_name, inference=True, device=device)
                 if args.size is not None:
                     net.adjust_size(size)
 
@@ -247,7 +260,7 @@ def benchmark(args: argparse.Namespace) -> None:
                 net = torch.compile(net)
 
             peak_memory = None
-            (t_elapsed, batch_size) = throughput_benchmark(net, device, sample_shape, model_name, args)
+            t_elapsed, batch_size = throughput_benchmark(net, device, sample_shape, model_name, args)
             if t_elapsed < 0.0:
                 continue
 
@@ -305,12 +318,18 @@ def get_args_parser() -> argparse.ArgumentParser:
             "--compile --suffix il-common --append\n"
             "python -m birder.scripts.benchmark --plain --models rdnet_t convnext_v1_tiny --bench-iter 50 --repeats 1 "
             "--gpu --size 416 --dry-run\n"
+            "python -m birder.scripts.benchmark --plain --models retinanet --backbone resnet_v1_50 --num-classes 91 "
+            "--size 640 --gpu --dry-run\n"
         ),
         formatter_class=cli.ArgumentHelpFormatter,
     )
     parser.add_argument("--filter", type=str, help="models to benchmark (fnmatch type filter)")
     parser.add_argument("--models", nargs="+", help="plain network names to benchmark")
     parser.add_argument("--plain", default=False, action="store_true", help="benchmark plain networks without weights")
+    parser.add_argument("--backbone", type=str, help="backbone name for plain detection benchmarks")
+    parser.add_argument(
+        "--num-classes", type=int, default=0, metavar="N", help="number of classes for plain benchmarks"
+    )
     parser.add_argument("--compile", default=False, action="store_true", help="enable compilation")
     parser.add_argument(
         "--amp", default=False, action="store_true", help="use torch.amp.autocast for mixed precision inference"
@@ -353,6 +372,12 @@ def validate_args(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--memory cannot be used with --compile")
     if args.plain is False and args.models is not None:
         raise cli.ValidationError("--models can only be used with --plain")
+    if args.backbone is not None and args.plain is False:
+        raise cli.ValidationError("--backbone can only be used with --plain")
+    if args.backbone is not None and registry.exists(args.backbone, net_type=DetectorBackbone) is False:
+        raise cli.ValidationError(
+            f"--backbone {args.backbone} not supported, see list-models tool for available options"
+        )
 
 
 def args_from_dict(**kwargs: Any) -> argparse.Namespace:

birder/scripts/evaluate.py

@@ -37,7 +37,7 @@ def evaluate(args: argparse.Namespace) -> None:
     amp_dtype: torch.dtype = getattr(torch, args.amp_dtype)
     model_list = birder.list_pretrained_models(args.filter)
     for model_name in model_list:
-        (net, (class_to_idx, signature, rgb_stats, *_)) = birder.load_pretrained_model(
+        net, (class_to_idx, signature, rgb_stats, *_) = birder.load_pretrained_model(
             model_name, inference=True, device=device, dtype=model_dtype
         )
         if args.parallel is True and torch.cuda.device_count() > 1:

birder/scripts/predict.py

@@ -204,7 +204,7 @@ def predict(args: argparse.Namespace) -> None:
             raise RuntimeError("'pip install torchao' to load quantization operators") from exc
 
     network_name = lib.get_network_name(args.network, tag=args.tag)
-    (net, (class_to_idx, signature, rgb_stats, *_)) = fs_ops.load_model(
+    net, (class_to_idx, signature, rgb_stats, *_) = fs_ops.load_model(
         device,
         args.network,
         config=args.model_config,
@@ -261,11 +261,11 @@ def predict(args: argparse.Namespace) -> None:
     if 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)
 
         num_samples = dataset_size
         dataset = make_wds_dataset(

birder/scripts/predict_detection.py

@@ -60,7 +60,7 @@ def predict(args: argparse.Namespace) -> None:
     network_name = lib.get_detection_network_name(
         args.network, tag=args.tag, backbone=args.backbone, backbone_tag=args.backbone_tag
     )
-    (net, (class_to_idx, signature, rgb_stats, *_)) = fs_ops.load_detection_model(
+    net, (class_to_idx, signature, rgb_stats, *_) = fs_ops.load_detection_model(
         device,
         args.network,
         config=args.model_config,
@@ -197,7 +197,7 @@ def predict(args: argparse.Namespace) -> None:
     # Inference
     tic = time.time()
     with torch.inference_mode():
-        (sample_paths, detections, targets) = infer_dataloader(
+        sample_paths, detections, targets = infer_dataloader(
             device,
             net,
             inference_loader,

birder/scripts/train.py

@@ -7,6 +7,7 @@ import time
 from collections.abc import Iterator
 from pathlib import Path
 from typing import Any
+from typing import Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -52,7 +53,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)
@@ -77,15 +78,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,
@@ -149,7 +150,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
     )
 
@@ -231,7 +232,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_checkpoint(
+        net, class_to_idx_saved, training_states = fs_ops.load_checkpoint(
             device,
             args.network,
             config=args.model_config,
@@ -247,7 +248,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_checkpoint(
+        net, class_to_idx_saved, training_states = fs_ops.load_checkpoint(
             device,
             args.network,
             config=args.model_config,
@@ -262,7 +263,7 @@ def train(args: argparse.Namespace) -> None:
             assert class_to_idx == class_to_idx_saved
 
     else:
-        net = registry.net_factory(args.network, sample_shape[1], num_outputs, config=args.model_config, size=args.size)
+        net = registry.net_factory(args.network, num_outputs, sample_shape[1], config=args.model_config, size=args.size)
         training_states = fs_ops.TrainingStates.empty()
 
     net.to(device, dtype=model_dtype)
@@ -328,7 +329,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:
@@ -474,16 +475,32 @@ 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):
         tic = time.time()
         net.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)
 
@@ -565,6 +582,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:
@@ -583,6 +603,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}  "
@@ -597,8 +620,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},
@@ -606,7 +638,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),
                     )
 
@@ -618,6 +650,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()
@@ -649,6 +683,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)
@@ -666,19 +703,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: {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:
@@ -849,7 +897,7 @@ def get_args_parser() -> argparse.ArgumentParser:
     training_cli.add_compile_args(parser)
     training_cli.add_checkpoint_args(parser, default_save_frequency=5, pretrained=True)
     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_barlow_twins.py

@@ -69,7 +69,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)
@@ -92,11 +92,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,
@@ -126,7 +126,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
     )
 
@@ -189,12 +189,12 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("barlow_twins", 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 = BarlowTwins(backbone, config={"projector_sizes": args.projector_dims, "off_lambda": args.off_lambda})
 
     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
         )
 
@@ -253,7 +253,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:
@@ -365,6 +365,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_byol.py

@@ -70,7 +70,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
@@ -94,11 +94,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,
@@ -128,7 +128,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
     )
 
@@ -191,7 +191,7 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("byol", 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 = BYOL(
         backbone,
         config={
@@ -202,7 +202,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
         )
 
@@ -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)
 

birder/scripts/train_capi.py

@@ -79,7 +79,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)
@@ -108,8 +108,8 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("capi", encoder=args.network, tag=args.tag)
 
-    student_backbone = registry.net_factory(args.network, sample_shape[1], 0, config=args.model_config, size=args.size)
-    teacher_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)
+    teacher_backbone = registry.net_factory(args.network, 0, sample_shape[1], config=args.model_config, size=args.size)
 
     teacher_backbone.load_state_dict(student_backbone.state_dict())
 
@@ -144,7 +144,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"]
@@ -194,11 +194,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,
@@ -224,7 +224,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
     )
 
@@ -326,8 +326,8 @@ def train(args: argparse.Namespace) -> None:
     student_temp = 0.12
 
     # Gradient scaler and AMP related tasks
-    (scaler, amp_dtype) = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
-    (clustering_scaler, _) = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
+    scaler, amp_dtype = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
+    clustering_scaler, _ = training_utils.get_amp_scaler(args.amp, args.amp_dtype)
 
     # Load states
     if args.load_states is True:
@@ -453,6 +453,11 @@ def train(args: argparse.Namespace) -> None:
         tic = time.time()
         net.train()
 
+        # Clear metrics
+        running_loss.clear()
+        running_clustering_loss.clear()
+        running_target_entropy.clear()
+
         if args.sinkhorn_queue_size is not None:
             queue_active = epoch > args.sinkhorn_queue_warmup_epochs
             teacher_without_ddp.head.set_queue_active(queue_active)
@@ -499,7 +504,7 @@ def train(args: argparse.Namespace) -> None:
 
             # Forward, backward and optimize
             with torch.amp.autocast("cuda", enabled=args.amp, dtype=amp_dtype):
-                (selected_assignments, clustering_loss) = teacher(images, None, predict_indices)
+                selected_assignments, clustering_loss = teacher(images, None, predict_indices)
 
             if clustering_scaler is not None:
                 clustering_scaler.scale(clustering_loss).backward()

birder/scripts/train_data2vec.py

@@ -69,7 +69,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
@@ -99,7 +99,7 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("data2vec", 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 = Data2Vec(
         backbone,
         config={
@@ -112,7 +112,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
         )
 
@@ -160,11 +160,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,
@@ -190,7 +190,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
     )
 
@@ -279,7 +279,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:
@@ -391,6 +391,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)