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

birder/scripts/train_simclr.py

@@ -67,7 +67,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)
@@ -90,11 +90,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,
@@ -124,7 +124,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
     )
 
@@ -187,7 +187,7 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("simclr", 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 = SimCLR(
         backbone,
         config={
@@ -199,7 +199,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
         )
 
@@ -258,7 +258,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:
@@ -370,6 +370,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_vicreg.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:
         args.size = registry.get_default_size(args.network)
@@ -93,11 +93,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,
@@ -127,7 +127,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
     )
 
@@ -190,7 +190,7 @@ def train(args: argparse.Namespace) -> None:
 
     network_name = get_mim_network_name("vicreg", 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 = VICReg(
         backbone,
         config={
@@ -205,7 +205,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
         )
 
@@ -264,7 +264,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:
@@ -376,6 +376,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/tools/adversarial.py

@@ -29,7 +29,7 @@ logger = logging.getLogger(__name__)
 def _load_model_and_transform(
     args: argparse.Namespace, device: torch.device
 ) -> tuple[torch.nn.Module, dict[str, int], RGBType, Callable[..., torch.Tensor], Callable[..., torch.Tensor]]:
-    (net, model_info) = fs_ops.load_model(
+    net, model_info = fs_ops.load_model(
         device, args.network, tag=args.tag, epoch=args.epoch, inference=True, reparameterized=args.reparameterized
     )
 
@@ -105,8 +105,8 @@ def _display_results(
     success: Optional[bool],
     result: AttackResult,
 ) -> None:
-    (orig_label, orig_prob) = original_pred
-    (adv_label, adv_prob) = adv_pred
+    orig_label, orig_prob = original_pred
+    adv_label, adv_prob = adv_pred
 
     # Log results
     logger.info(f"Original: {orig_label} ({orig_prob * 100:.2f}%)")
@@ -139,7 +139,7 @@ def run_attack(args: argparse.Namespace) -> None:
 
     logger.info(f"Using device {device}")
 
-    (net, class_to_idx, rgb_stats, transform, reverse_transform) = _load_model_and_transform(args, device)
+    net, class_to_idx, rgb_stats, transform, reverse_transform = _load_model_and_transform(args, device)
     label_names = [name for name, _idx in sorted(class_to_idx.items(), key=lambda item: item[1])]
     img = Image.open(args.image_path)
     input_tensor = transform(img).unsqueeze(dim=0).to(device)

birder/tools/auto_anchors.py

@@ -92,7 +92,7 @@ def _load_coco_boxes(
             stats["missing_images"] += 1
             continue
 
-        (img_w, img_h, file_name) = images[image_id]
+        img_w, img_h, file_name = images[image_id]
         if file_name in ignore_list:
             stats["ignored_images"] += 1
             continue
@@ -219,7 +219,7 @@ def _validate_args(
     output_format = args.format if args.format is not None else (preset["format"] if preset else None)
     if num_scales is None or num_anchors is None or output_format is None:
         raise cli.ValidationError(
-            "Missing configuration. Provide --num-scales, --num-anchors, and --format or use a --preset"
+            "Missing configuration. Provide --num-scales, --num-anchors and --format or use a --preset"
         )
     if num_scales < 1:
         raise cli.ValidationError("--num-scales must be >= 1")
@@ -244,10 +244,10 @@ def _validate_args(
 
 # pylint: disable=too-many-locals
 def auto_anchors(args: argparse.Namespace) -> None:
-    (size, num_scales, num_anchors, output_format, strides) = _validate_args(args)
+    size, num_scales, num_anchors, output_format, strides = _validate_args(args)
 
     ignore_list = _load_ignore_list(args.ignore_file)
-    (boxes, stats) = _load_coco_boxes(
+    boxes, stats = _load_coco_boxes(
         args.coco_json_path, size, ignore_list, args.min_size, ignore_crowd=not args.include_crowd
     )
 
@@ -262,7 +262,7 @@ def auto_anchors(args: argparse.Namespace) -> None:
         f"missing_size={stats['missing_size']}, too_small={stats['too_small']}"
     )
 
-    (anchors, _assignments) = _kmeans_anchors(boxes, num_anchors, args.seed, args.max_iter)
+    anchors, _assignments = _kmeans_anchors(boxes, num_anchors, args.seed, args.max_iter)
     areas = anchors.prod(dim=1)
     anchors = anchors[torch.argsort(areas)]
     anchors_per_scale = num_anchors // num_scales

birder/tools/avg_model.py

@@ -44,7 +44,7 @@ def avg_models(
             num_classes = lib.get_num_labels_from_signature(signature)
             size = lib.get_size_from_signature(signature)
 
-            net = registry.net_factory(network, input_channels, num_classes, size=size)
+            net = registry.net_factory(network, num_classes, input_channels, size=size)
             if reparameterized is True:
                 net.reparameterize_model()
 

birder/tools/convert_model.py

@@ -74,6 +74,7 @@ def onnx_export(
     net: torch.nn.Module,
     signature: SignatureType | DetectionSignatureType,
     class_to_idx: dict[str, int],
+    rgb_stats: RGBType,
     model_path: str | Path,
     dynamo: bool,
     trace: bool,
@@ -117,9 +118,19 @@ def onnx_export(
 
     signature["inputs"][0]["data_shape"][0] = 0
 
-    logger.info("Saving class to index json...")
-    with open(f"{model_path}_class_to_idx.json", "w", encoding="utf-8") as handle:
-        json.dump(class_to_idx, handle, indent=2)
+    logger.info("Saving model data json...")
+    with open(f"{model_path}_data.json", "w", encoding="utf-8") as handle:
+        json.dump(
+            {
+                "birder_version": __version__,
+                "task": net.task,
+                "class_to_idx": class_to_idx,
+                "signature": signature,
+                "rgb_stats": rgb_stats,
+            },
+            handle,
+            indent=2,
+        )
 
     # Test exported model
     onnx_model = onnx.load(str(model_path))
@@ -238,7 +249,7 @@ def main(args: argparse.Namespace) -> None:
     signature: SignatureType | DetectionSignatureType
     backbone_custom_config = None
     if args.backbone is None:
-        (net, (class_to_idx, signature, rgb_stats, custom_config)) = fs_ops.load_model(
+        net, (class_to_idx, signature, rgb_stats, custom_config) = fs_ops.load_model(
             device,
             args.network,
             config=args.model_config,
@@ -251,22 +262,20 @@ def main(args: argparse.Namespace) -> None:
         network_name = lib.get_network_name(args.network, tag=args.tag)
 
     else:
-        (net, (class_to_idx, signature, rgb_stats, custom_config, backbone_custom_config)) = (
-            fs_ops.load_detection_model(
-                device,
-                args.network,
-                config=args.model_config,
-                tag=args.tag,
-                reparameterized=args.reparameterized,
-                backbone=args.backbone,
-                backbone_config=args.backbone_model_config,
-                backbone_tag=args.backbone_tag,
-                backbone_reparameterized=args.backbone_reparameterized,
-                epoch=args.epoch,
-                new_size=args.resize,
-                inference=True,
-                export_mode=True,
-            )
+        net, (class_to_idx, signature, rgb_stats, custom_config, backbone_custom_config) = fs_ops.load_detection_model(
+            device,
+            args.network,
+            config=args.model_config,
+            tag=args.tag,
+            reparameterized=args.reparameterized,
+            backbone=args.backbone,
+            backbone_config=args.backbone_model_config,
+            backbone_tag=args.backbone_tag,
+            backbone_reparameterized=args.backbone_reparameterized,
+            epoch=args.epoch,
+            new_size=args.resize,
+            inference=True,
+            export_mode=True,
         )
         network_name = lib.get_detection_network_name(
             args.network, tag=args.tag, backbone=args.backbone, backbone_tag=args.backbone_tag
@@ -407,8 +416,7 @@ def main(args: argparse.Namespace) -> None:
         )
 
     elif args.onnx is True or args.onnx_dynamo is True:
-        config_export(net, signature, rgb_stats, model_path)
-        onnx_export(net, signature, class_to_idx, model_path, args.onnx_dynamo, args.trace)
+        onnx_export(net, signature, class_to_idx, rgb_stats, model_path, args.onnx_dynamo, args.trace)
 
     elif args.config is True:
         config_export(net, signature, rgb_stats, model_path)

birder/tools/det_results.py

@@ -239,7 +239,7 @@ def main(args: argparse.Namespace) -> None:
             logger.warning("Cannot compare confusion matrix, processing only the first file")
 
         results = next(iter(results_dict.values()))
-        (cnf_matrix, label_names) = confusion_matrix_data(
+        cnf_matrix, label_names = confusion_matrix_data(
             results, args.cnf_score_threshold, args.cnf_iou_threshold, args.classes, args.cnf_errors_only
         )
         title = f"Confusion matrix (score >= {args.cnf_score_threshold:.2f}, IoU >= {args.cnf_iou_threshold:.2f})"

birder/tools/download_model.py

@@ -52,7 +52,7 @@ def main(args: argparse.Namespace) -> None:
         )
         raise SystemExit(1)
 
-    (model_file, url) = get_pretrained_model_url(args.model_name, args.format)
+    model_file, url = get_pretrained_model_url(args.model_name, args.format)
     dst = settings.MODELS_DIR.joinpath(model_file)
     if dst.exists() is True and args.force is False:
         logger.warning(f"File {model_file} already exists... aborting")

birder/tools/ensemble_model.py

@@ -58,7 +58,7 @@ def main(args: argparse.Namespace) -> None:
     signature_list = []
     rgb_stats_list = []
     for network in args.networks:
-        (net, model_info) = fs_ops.load_model(device, network, inference=True, pts=args.pts, pt2=args.pt2)
+        net, model_info = fs_ops.load_model(device, network, inference=True, pts=args.pts, pt2=args.pt2)
         nets.append(net)
         class_to_idx_list.append(model_info.class_to_idx)
         signature_list.append(model_info.signature)

birder/tools/introspection.py

@@ -126,6 +126,14 @@ def set_parser(subparsers: Any) -> None:
         formatter_class=cli.ArgumentHelpFormatter,
     )
     subparser.add_argument("-n", "--network", type=str, required=True, help="the neural network to use")
+    subparser.add_argument(
+        "--model-config",
+        action=cli.FlexibleDictAction,
+        help=(
+            "override the model default configuration, accepts key-value pairs or JSON "
+            "('drop_path_rate=0.2' or '{\"units\": [3, 24, 36, 3], \"dropout\": 0.2}'"
+        ),
+    )
     subparser.add_argument("-e", "--epoch", type=int, metavar="N", help="model checkpoint to load")
     subparser.add_argument("-t", "--tag", type=str, help="model tag (from the training phase)")
     subparser.add_argument(
@@ -145,7 +153,7 @@ def set_parser(subparsers: Any) -> None:
     subparser.add_argument(
         "--target",
         type=str,
-        help="target class, leave empty to use predicted class (gradcam, guided-backprop, and transformer-attribution)",
+        help="target class, leave empty to use predicted class (gradcam, guided-backprop and transformer-attribution)",
     )
     subparser.add_argument("--block-name", type=str, default="body", help="target block (gradcam only)")
     subparser.add_argument(
@@ -203,9 +211,10 @@ def main(args: argparse.Namespace) -> None:
 
     logger.info(f"Using device {device}")
 
-    (net, model_info) = fs_ops.load_model(
+    net, model_info = fs_ops.load_model(
         device,
         args.network,
+        config=args.model_config,
         tag=args.tag,
         epoch=args.epoch,
         new_size=args.size,

birder/tools/labelme_to_coco.py

@@ -25,8 +25,8 @@ def _create_annotation(
     annotation["image_id"] = image_id
 
     # Bounding box in (x, y, w, h) format
-    (x0, y0) = points[0]
-    (x1, y1) = points[1]
+    x0, y0 = points[0]
+    x1, y1 = points[1]
     x = min(x0, x1)
     y = min(y0, y1)
     w = abs(x0 - x1)

birder/tools/model_info.py

@@ -73,7 +73,7 @@ def main(args: argparse.Namespace) -> None:
     signature: SignatureType | DetectionSignatureType
     backbone_custom_config = None
     if args.backbone is None:
-        (net, (class_to_idx, signature, rgb_stats, custom_config)) = fs_ops.load_model(
+        net, (class_to_idx, signature, rgb_stats, custom_config) = fs_ops.load_model(
             device,
             args.network,
             tag=args.tag,
@@ -86,19 +86,17 @@ def main(args: argparse.Namespace) -> None:
         )
 
     else:
-        (net, (class_to_idx, signature, rgb_stats, custom_config, backbone_custom_config)) = (
-            fs_ops.load_detection_model(
-                device,
-                args.network,
-                tag=args.tag,
-                backbone=args.backbone,
-                backbone_tag=args.backbone_tag,
-                epoch=args.epoch,
-                inference=True,
-                pts=args.pts,
-                pt2=args.pt2,
-                st=args.st,
-            )
+        net, (class_to_idx, signature, rgb_stats, custom_config, backbone_custom_config) = fs_ops.load_detection_model(
+            device,
+            args.network,
+            tag=args.tag,
+            backbone=args.backbone,
+            backbone_tag=args.backbone_tag,
+            epoch=args.epoch,
+            inference=True,
+            pts=args.pts,
+            pt2=args.pt2,
+            st=args.st,
         )
 
     model_info = get_model_info(net)

birder/tools/pack.py

@@ -114,7 +114,7 @@ def read_worker(q_in: Any, q_out: Any, error_event: Any, size: Optional[int], fi
             break
 
         try:
-            (idx, path, target) = deq
+            idx, path, target = deq
             if size is None:
                 suffix = Path(path).suffix[1:]
                 if file_format != suffix:
@@ -172,7 +172,7 @@ def wds_write_worker(
             while more:
                 deq: Optional[tuple[int, bytes, str, int]] = q_out.get()
                 if deq is not None:
-                    (idx, sample, suffix, target) = deq
+                    idx, sample, suffix, target = deq
                     buf[idx] = (sample, suffix, target)
 
                 else:
@@ -180,7 +180,7 @@ def wds_write_worker(
 
                 # Ensures ordered write
                 while count in buf:
-                    (sample, suffix, target) = buf[count]
+                    sample, suffix, target = buf[count]
                     del buf[count]
 
                     if args.no_cls is True:
@@ -238,7 +238,7 @@ def directory_write_worker(
             while more:
                 deq: Optional[tuple[int, bytes, str, int]] = q_out.get()
                 if deq is not None:
-                    (idx, sample, suffix, target) = deq
+                    idx, sample, suffix, target = deq
                     buf[idx] = (sample, suffix, target)
 
                 else:
@@ -246,7 +246,7 @@ def directory_write_worker(
 
                 # Ensures ordered write
                 while count in buf:
-                    (sample, suffix, target) = buf[count]
+                    sample, suffix, target = buf[count]
                     del buf[count]
                     with open(
                         pack_path.joinpath(idx_to_class[target]).joinpath(f"{count:06d}.{suffix}"), "wb"
@@ -274,7 +274,7 @@ def pack(args: argparse.Namespace, pack_path: Path) -> None:
             if len(line.strip()) == 0 or line.strip().startswith("#") is True:
                 continue
 
-            (data_path, r) = line.split()
+            data_path, r = line.split()
             data_path = os.path.expanduser(data_path)
             repeats = int(r)
             for _ in range(repeats):
@@ -391,7 +391,7 @@ def pack(args: argparse.Namespace, pack_path: Path) -> None:
                     cleanup_processes()
                     raise RuntimeError()
 
-            (path, target) = dataset[sample_idx]
+            path, target = dataset[sample_idx]
 
             while True:
                 try:
@@ -430,7 +430,7 @@ def pack(args: argparse.Namespace, pack_path: Path) -> None:
             raise RuntimeError()
 
         if args.type == "wds":
-            (wds_path, num_shards) = fs_ops.wds_braces_from_path(pack_path, prefix=f"{args.suffix}-{args.split}")
+            wds_path, num_shards = fs_ops.wds_braces_from_path(pack_path, prefix=f"{args.suffix}-{args.split}")
             logger.info(f"Packed {len(dataset):,} samples into {num_shards} shards at {wds_path}")
         elif args.type == "directory":
             logger.info(f"Packed {len(dataset):,} samples")

birder/tools/quantize_model.py

@@ -1,7 +1,9 @@
 import argparse
 import itertools
+import json
 import logging
 import time
+from pathlib import Path
 from typing import Any
 
 import torch
@@ -15,7 +17,11 @@ from birder.common import fs_ops
 from birder.common import lib
 from birder.common.lib import get_network_name
 from birder.conf import settings
+from birder.data.transforms.classification import RGBType
 from birder.data.transforms.classification import inference_preset
+from birder.net.base import SignatureType
+from birder.net.detection.base import DetectionSignatureType
+from birder.version import __version__
 
 try:
     from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e
@@ -28,8 +34,10 @@ except ImportError:
     _HAS_TORCHAO = False
 
 try:
+    from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
     from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import XNNPACKQuantizer
     from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import get_symmetric_quantization_config
+    from executorch.exir import to_edge_transform_and_lower
 
     _HAS_EXECUTORCH = True
 except ImportError:
@@ -54,6 +62,33 @@ def _build_quantizer(backend: str) -> Any:
     raise ValueError(f"Unsupported backend: {backend}")
 
 
+def _save_pte(
+    exported_net: torch.export.ExportedProgram,
+    dst: str | Path,
+    task: str,
+    class_to_idx: dict[str, int],
+    signature: SignatureType | DetectionSignatureType,
+    rgb_stats: RGBType,
+) -> None:
+    edge_program = to_edge_transform_and_lower(exported_net, partitioner=[XnnpackPartitioner()])
+    executorch_program = edge_program.to_executorch()
+    with open(dst, "wb") as f:
+        f.write(executorch_program.buffer)
+
+    with open(f"{dst}_data.json", "w", encoding="utf-8") as handle:
+        json.dump(
+            {
+                "birder_version": __version__,
+                "task": task,
+                "class_to_idx": class_to_idx,
+                "signature": signature,
+                "rgb_stats": rgb_stats,
+            },
+            handle,
+            indent=2,
+        )
+
+
 def set_parser(subparsers: Any) -> None:
     subparser = subparsers.add_parser(
         "quantize-model",
@@ -65,6 +100,7 @@ def set_parser(subparsers: Any) -> None:
             "python -m birder.tools quantize-model -n convnext_v2_tiny -t eu-common\n"
             "python -m birder.tools quantize-model --network densenet_121 -e 100 --num-calibration-batches 256\n"
             "python -m birder.tools quantize-model -n efficientnet_v2_s -e 200 --qbackend xnnpack --batch-size 1\n"
+            "python -m birder.tools quantize-model -n hgnet_v2_b4 --qbackend xnnpack --pte\n"
         ),
         formatter_class=cli.ArgumentHelpFormatter,
     )
@@ -81,6 +117,9 @@ def set_parser(subparsers: Any) -> None:
     subparser.add_argument(
         "--qbackend", type=str, choices=["x86", "xnnpack"], default="x86", help="quantization backend"
     )
+    subparser.add_argument(
+        "--pte", default=False, action="store_true", help="lower quantized model to ExecuTorch PTE format"
+    )
     subparser.add_argument("--batch-size", type=int, default=1, metavar="N", help="the batch size")
     subparser.add_argument(
         "--num-calibration-batches",
@@ -96,8 +135,13 @@ def set_parser(subparsers: Any) -> None:
 
 # pylint: disable=too-many-locals
 def main(args: argparse.Namespace) -> None:
+    if args.pte is True and args.qbackend != "xnnpack":
+        raise cli.ValidationError("--pte requires --qbackend xnnpack")
+
     network_name = get_network_name(args.network, tag=args.tag)
     model_path = fs_ops.model_path(network_name, epoch=args.epoch, quantized=True, pt2=True)
+    if args.pte is True:
+        model_path = model_path.with_suffix(".pte")
     if model_path.exists() is True and args.force is False:
         logger.warning("Quantized model already exists... aborting")
         raise SystemExit(1)
@@ -105,7 +149,7 @@ def main(args: argparse.Namespace) -> None:
     device = torch.device("cpu")
 
     # Load model
-    (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, tag=args.tag, epoch=args.epoch, inference=True, reparameterized=args.reparameterized
     )
     net.eval()
@@ -154,9 +198,14 @@ def main(args: argparse.Namespace) -> None:
         exported_quantized_net = torch.export.export(quantized_net, example_inputs)
 
     toc = time.time()
-    (minutes, seconds) = divmod(toc - tic, 60)
+    minutes, seconds = divmod(toc - tic, 60)
     logger.info(f"{int(minutes):0>2}m{seconds:04.1f}s to quantize model")
 
     model_path = fs_ops.model_path(network_name, epoch=args.epoch, quantized=True, pt2=True)
-    logger.info(f"Saving quantized PT2 model {model_path}...")
-    fs_ops.save_pt2(exported_quantized_net, model_path, task, class_to_idx, signature, rgb_stats)
+    if args.pte is True:
+        model_path = model_path.with_suffix(".pte")
+        logger.info(f"Lowering quantized model to PTE {model_path}...")
+        _save_pte(exported_quantized_net, model_path, task, class_to_idx, signature, rgb_stats)
+    else:
+        logger.info(f"Saving quantized PT2 model {model_path}...")
+        fs_ops.save_pt2(exported_quantized_net, model_path, task, class_to_idx, signature, rgb_stats)

birder/tools/results.py

@@ -125,7 +125,7 @@ def print_most_confused_pairs(most_confused_df: pl.DataFrame) -> None:
 
 def convert_to_sparse(results_file: str, sparse_k: int) -> None:
     logger.info(f"Converting {results_file} to sparse format (k={sparse_k})...")
-    (_, detected_sparse_k) = detect_file_format(results_file)
+    _, detected_sparse_k = detect_file_format(results_file)
 
     if detected_sparse_k is not None:
         logger.info(f"File is already in sparse format (with k={detected_sparse_k}). Skipping conversion.")
@@ -233,7 +233,7 @@ def main(args: argparse.Namespace) -> None:
             logger.warning("Cannot print mistakes in compare mode. processing only the first file")
 
         if args.imperfect_only is True:
-            (result_name, results) = next(iter(results_dict.items()))
+            result_name, results = next(iter(results_dict.items()))
             mistake_prediction_indices = results.mistakes["prediction"].unique().to_numpy().tolist()
             mistake_label_indices = results.mistakes["label"].unique().to_numpy().tolist()
             imperfect_class_indices = np.unique(mistake_prediction_indices + mistake_label_indices).tolist()