sleap-nn 0.1.0a2__py3-none-any.whl → 0.1.0a4__py3-none-any.whl

sleap_nn/export/exporters/__init__.py

@@ -0,0 +1,51 @@
+"""Exporters for serialized model formats."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable, Optional
+
+import torch
+
+from sleap_nn.export.exporters.onnx_exporter import export_to_onnx
+from sleap_nn.export.exporters.tensorrt_exporter import export_to_tensorrt
+
+
+def export_model(
+    model: torch.nn.Module,
+    save_path: str | Path,
+    fmt: str = "onnx",
+    input_shape: Iterable[int] = (1, 1, 512, 512),
+    opset_version: int = 17,
+    output_names: Optional[list] = None,
+    verify: bool = True,
+    **kwargs,
+) -> Path:
+    """Export a model to the requested format."""
+    fmt = fmt.lower()
+    if fmt == "onnx":
+        return export_to_onnx(
+            model,
+            save_path,
+            input_shape=input_shape,
+            opset_version=opset_version,
+            output_names=output_names,
+            verify=verify,
+        )
+    if fmt == "tensorrt":
+        return export_to_tensorrt(model, save_path, input_shape=input_shape, **kwargs)
+    if fmt == "both":
+        export_to_onnx(
+            model,
+            save_path,
+            input_shape=input_shape,
+            opset_version=opset_version,
+            output_names=output_names,
+            verify=verify,
+        )
+        return export_to_tensorrt(model, save_path, input_shape=input_shape, **kwargs)
+
+    raise ValueError(f"Unknown export format: {fmt}")
+
+
+__all__ = ["export_model", "export_to_onnx", "export_to_tensorrt"]

sleap_nn/export/exporters/onnx_exporter.py

@@ -0,0 +1,80 @@
+"""ONNX export utilities."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional
+
+import torch
+
+
+def export_to_onnx(
+    model: torch.nn.Module,
+    save_path: str | Path,
+    input_shape: Iterable[int] = (1, 1, 512, 512),
+    input_dtype: torch.dtype = torch.uint8,
+    opset_version: int = 17,
+    dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
+    input_names: Optional[List[str]] = None,
+    output_names: Optional[List[str]] = None,
+    do_constant_folding: bool = True,
+    verify: bool = True,
+) -> Path:
+    """Export a PyTorch model to ONNX."""
+    save_path = Path(save_path)
+    model.eval()
+
+    if input_names is None:
+        input_names = ["image"]
+    if dynamic_axes is None:
+        dynamic_axes = {"image": {0: "batch", 2: "height", 3: "width"}}
+
+    device = None
+    try:
+        device = next(model.parameters()).device
+    except StopIteration:
+        device = torch.device("cpu")
+
+    if input_dtype.is_floating_point:
+        dummy_input = torch.randn(*input_shape, device=device, dtype=input_dtype)
+    else:
+        dummy_input = torch.randint(
+            0, 256, input_shape, device=device, dtype=input_dtype
+        )
+
+    if output_names is None:
+        with torch.no_grad():
+            test_out = model(dummy_input)
+        output_names = _infer_output_names(test_out)
+
+    torch.onnx.export(
+        model,
+        dummy_input,
+        save_path.as_posix(),
+        opset_version=opset_version,
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes=dynamic_axes,
+        do_constant_folding=do_constant_folding,
+        dynamo=False,
+    )
+
+    if verify:
+        _verify_onnx(save_path)
+
+    return save_path
+
+
+def _infer_output_names(output) -> List[str]:
+    if isinstance(output, dict):
+        return list(output.keys())
+    if isinstance(output, (list, tuple)):
+        return [f"output_{idx}" for idx in range(len(output))]
+    return ["output_0"]
+
+
+def _verify_onnx(path: Path) -> None:
+    import onnx
+
+    model = onnx.load(path.as_posix())
+    onnx.checker.check_model(model)

sleap_nn/export/exporters/tensorrt_exporter.py

@@ -0,0 +1,291 @@
+"""TensorRT export utilities."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable, Optional, Tuple
+
+import torch
+from torch import nn
+
+
+def export_to_tensorrt(
+    model: nn.Module,
+    save_path: str | Path,
+    input_shape: Tuple[int, int, int, int] = (1, 1, 512, 512),
+    input_dtype: torch.dtype = torch.uint8,
+    precision: str = "fp16",
+    min_shape: Optional[Tuple[int, int, int, int]] = None,
+    opt_shape: Optional[Tuple[int, int, int, int]] = None,
+    max_shape: Optional[Tuple[int, int, int, int]] = None,
+    workspace_size: int = 2 << 30,  # 2GB default
+    method: str = "onnx",
+    verbose: bool = True,
+) -> Path:
+    """Export a PyTorch model to TensorRT format.
+
+    This function supports multiple compilation methods:
+    - "onnx": Exports to ONNX first, then compiles with TensorRT (most reliable)
+    - "jit": Uses torch.jit.trace + torch_tensorrt.compile (alternative)
+
+    Args:
+        model: The PyTorch model to export (typically an ONNX wrapper).
+        save_path: Path to save the TensorRT engine (.trt file).
+        input_shape: (B, C, H, W) optimal input tensor shape.
+        input_dtype: Input tensor dtype (torch.uint8 or torch.float32).
+        precision: Model precision - "fp32" or "fp16".
+        min_shape: Minimum input shape for dynamic shapes (default: batch=1, H/W halved).
+        opt_shape: Optimal input shape (default: same as input_shape).
+        max_shape: Maximum input shape (default: batch=16, H/W doubled).
+        workspace_size: TensorRT workspace size in bytes (default 2GB).
+        method: Compilation method - "onnx" or "jit".
+        verbose: Print export info.
+
+    Returns:
+        Path to the exported TensorRT engine.
+
+    Note:
+        TensorRT models are NOT cross-platform. The exported model will only
+        work on the same GPU architecture and TensorRT version used for export.
+    """
+    import tensorrt as trt
+
+    model.eval()
+    device = next(model.parameters()).device
+
+    save_path = Path(save_path)
+    if not save_path.suffix:
+        save_path = save_path.with_suffix(".trt")
+
+    B, C, H, W = input_shape
+
+    if min_shape is None:
+        min_shape = (1, C, H // 2, W // 2)
+    if opt_shape is None:
+        opt_shape = input_shape
+    if max_shape is None:
+        max_shape = (min(16, B * 4), C, H * 2, W * 2)
+
+    if verbose:
+        print(f"Exporting model to TensorRT...")
+        print(f"  Input shape: {input_shape}")
+        print(f"  Min/Opt/Max: {min_shape} / {opt_shape} / {max_shape}")
+        print(f"  Precision: {precision}")
+        print(f"  Workspace: {workspace_size / 1e9:.1f} GB")
+        print(f"  Method: {method}")
+
+    if method == "onnx":
+        return _export_tensorrt_onnx(
+            model,
+            save_path,
+            input_shape,
+            input_dtype,
+            min_shape,
+            opt_shape,
+            max_shape,
+            precision,
+            workspace_size,
+            verbose,
+        )
+    elif method == "jit":
+        return _export_tensorrt_jit(
+            model,
+            save_path,
+            input_shape,
+            input_dtype,
+            min_shape,
+            opt_shape,
+            max_shape,
+            precision,
+            workspace_size,
+            verbose,
+        )
+    else:
+        raise ValueError(f"Unknown method: {method}. Use 'onnx' or 'jit'.")
+
+
+def _export_tensorrt_onnx(
+    model: nn.Module,
+    save_path: Path,
+    input_shape: Tuple[int, int, int, int],
+    input_dtype: torch.dtype,
+    min_shape: Tuple[int, int, int, int],
+    opt_shape: Tuple[int, int, int, int],
+    max_shape: Tuple[int, int, int, int],
+    precision: str,
+    workspace_size: int,
+    verbose: bool,
+) -> Path:
+    """Export via ONNX, then compile to TensorRT engine."""
+    import tensorrt as trt
+
+    # Check if ONNX file already exists (from prior export step)
+    onnx_path = save_path.with_suffix(".onnx")
+
+    if onnx_path.exists():
+        if verbose:
+            print(f"  Using existing ONNX file: {onnx_path}")
+    else:
+        # Need to export to ONNX first
+        device = next(model.parameters()).device
+
+        if verbose:
+            print("  Exporting to ONNX first...")
+
+        # Create example input with correct dtype
+        if input_dtype == torch.uint8:
+            example_input = torch.randint(
+                0, 255, input_shape, dtype=torch.uint8, device=device
+            )
+        else:
+            example_input = torch.randn(*input_shape, dtype=input_dtype, device=device)
+
+        # Export to ONNX
+        torch.onnx.export(
+            model,
+            example_input,
+            onnx_path,
+            opset_version=17,
+            input_names=["images"],
+            dynamic_axes={"images": {0: "batch", 2: "height", 3: "width"}},
+            do_constant_folding=True,
+        )
+
+        if verbose:
+            print(f"  ONNX export complete: {onnx_path}")
+
+    if verbose:
+        print(f"  Building TensorRT engine (this may take a while)...")
+
+    # Create TensorRT builder
+    logger = trt.Logger(trt.Logger.WARNING if not verbose else trt.Logger.INFO)
+    builder = trt.Builder(logger)
+    network = builder.create_network(
+        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+    )
+    parser = trt.OnnxParser(network, logger)
+
+    # Parse ONNX model
+    with open(onnx_path, "rb") as f:
+        if not parser.parse(f.read()):
+            errors = []
+            for i in range(parser.num_errors):
+                errors.append(str(parser.get_error(i)))
+            raise RuntimeError(f"ONNX parsing failed: {errors}")
+
+    # Configure builder
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_size)
+
+    if precision == "fp16":
+        if builder.platform_has_fast_fp16:
+            config.set_flag(trt.BuilderFlag.FP16)
+            if verbose:
+                print("  Enabled FP16 mode")
+        else:
+            if verbose:
+                print("  WARNING: Platform does not have fast FP16, using FP32")
+
+    # Set up optimization profile for dynamic shapes
+    profile = builder.create_optimization_profile()
+    input_name = network.get_input(0).name
+
+    profile.set_shape(input_name, min_shape, opt_shape, max_shape)
+    config.add_optimization_profile(profile)
+
+    # Build engine
+    serialized_engine = builder.build_serialized_network(network, config)
+    if serialized_engine is None:
+        raise RuntimeError("Failed to build TensorRT engine")
+
+    # Save engine
+    engine_path = save_path.with_suffix(".trt")
+    with open(engine_path, "wb") as f:
+        f.write(serialized_engine)
+
+    if verbose:
+        import os
+
+        print(f"  Exported TensorRT engine to: {engine_path}")
+        print(f"  Engine size: {os.path.getsize(engine_path) / 1e6:.2f} MB")
+
+    return engine_path
+
+
+def _export_tensorrt_jit(
+    model: nn.Module,
+    save_path: Path,
+    input_shape: Tuple[int, int, int, int],
+    input_dtype: torch.dtype,
+    min_shape: Tuple[int, int, int, int],
+    opt_shape: Tuple[int, int, int, int],
+    max_shape: Tuple[int, int, int, int],
+    precision: str,
+    workspace_size: int,
+    verbose: bool,
+) -> Path:
+    """Export using torch.jit.trace + torch_tensorrt.compile."""
+    import torch_tensorrt
+
+    device = next(model.parameters()).device
+
+    if verbose:
+        print("  Tracing model with torch.jit...")
+
+    # Create example input (float32 for tracing)
+    if input_dtype == torch.uint8:
+        example_input = torch.randint(
+            0, 255, input_shape, dtype=torch.uint8, device=device
+        )
+    else:
+        example_input = torch.randn(*input_shape, dtype=input_dtype, device=device)
+
+    # Trace the model
+    with torch.no_grad():
+        traced_model = torch.jit.trace(model, example_input)
+
+    # Map precision to torch dtype
+    precision_map = {
+        "fp32": torch.float32,
+        "fp16": torch.float16,
+    }
+    if precision not in precision_map:
+        raise ValueError(f"Unknown precision: {precision}. Use 'fp32' or 'fp16'")
+
+    enabled_precisions = {precision_map[precision]}
+    if precision == "fp16":
+        enabled_precisions.add(torch.float32)
+
+    # Create input specs for TensorRT
+    trt_inputs = [
+        torch_tensorrt.Input(
+            min_shape=min_shape,
+            opt_shape=opt_shape,
+            max_shape=max_shape,
+            dtype=torch.float32,  # TRT internally uses float32 input spec
+        )
+    ]
+
+    if verbose:
+        print("  Compiling with TensorRT...")
+
+    # Compile with TensorRT
+    trt_model = torch_tensorrt.compile(
+        traced_model,
+        inputs=trt_inputs,
+        enabled_precisions=enabled_precisions,
+        workspace_size=workspace_size,
+        truncate_long_and_double=True,
+    )
+
+    # Save as TorchScript
+    ts_path = save_path.with_suffix(".ts")
+    torch.jit.save(trt_model, ts_path)
+
+    if verbose:
+        import os
+
+        print(f"  Exported TensorRT-accelerated model to: {ts_path}")
+        print(f"  Model size: {os.path.getsize(ts_path) / 1e6:.2f} MB")
+
+    return ts_path

sleap_nn/export/metadata.py

@@ -0,0 +1,225 @@
+"""Metadata helpers for exported models."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, asdict
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+import hashlib
+import json
+
+from sleap_nn import __version__
+
+
+@dataclass
+class ExportMetadata:
+    """Metadata embedded or saved alongside exported models."""
+
+    # Version info
+    sleap_nn_version: str
+    export_timestamp: str
+    export_format: str  # "onnx" or "tensorrt"
+
+    # Model info
+    model_type: str  # "centroid", "centered_instance", "topdown", "bottomup"
+    model_name: str
+    checkpoint_path: str
+
+    # Architecture
+    backbone: str
+    n_nodes: int
+    n_edges: int
+    node_names: List[str]
+    edge_inds: List[Tuple[int, int]]
+
+    # Input/output spec
+    input_scale: float
+    input_channels: int
+    output_stride: int
+    crop_size: Optional[Tuple[int, int]] = None
+
+    # Export parameters
+    max_instances: Optional[int] = None
+    max_peaks_per_node: Optional[int] = None
+    max_batch_size: int = 1
+    precision: str = "fp32"
+
+    # Preprocessing - input is uint8 [0,255], normalized internally to float32 [0,1]
+    input_dtype: str = "uint8"
+    normalization: str = "0_to_1"
+
+    # Multiclass model fields (optional)
+    n_classes: Optional[int] = None
+    class_names: Optional[List[str]] = None
+
+    # Training config reference
+    training_config_embedded: bool = False
+    training_config_hash: str = ""
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to JSON-serializable dict."""
+        data = asdict(self)
+        data["edge_inds"] = [list(pair) for pair in self.edge_inds]
+        if self.crop_size is not None:
+            data["crop_size"] = list(self.crop_size)
+        return data
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ExportMetadata":
+        """Load from dict."""
+        edge_inds = [tuple(pair) for pair in data.get("edge_inds", [])]
+        crop_size = data.get("crop_size")
+        if crop_size is not None:
+            crop_size = tuple(crop_size)
+        return cls(
+            sleap_nn_version=data.get("sleap_nn_version", ""),
+            export_timestamp=data.get("export_timestamp", ""),
+            export_format=data.get("export_format", ""),
+            model_type=data.get("model_type", ""),
+            model_name=data.get("model_name", ""),
+            checkpoint_path=data.get("checkpoint_path", ""),
+            backbone=data.get("backbone", ""),
+            n_nodes=int(data.get("n_nodes", 0)),
+            n_edges=int(data.get("n_edges", 0)),
+            node_names=list(data.get("node_names", [])),
+            edge_inds=edge_inds,
+            input_scale=float(data.get("input_scale", 1.0)),
+            input_channels=int(data.get("input_channels", 1)),
+            output_stride=int(data.get("output_stride", 1)),
+            crop_size=crop_size,
+            max_instances=data.get("max_instances"),
+            max_peaks_per_node=data.get("max_peaks_per_node"),
+            max_batch_size=int(data.get("max_batch_size", 1)),
+            precision=data.get("precision", "fp32"),
+            input_dtype=data.get("input_dtype", "uint8"),
+            normalization=data.get("normalization", "0_to_1"),
+            n_classes=data.get("n_classes"),
+            class_names=data.get("class_names"),
+            training_config_embedded=bool(data.get("training_config_embedded", False)),
+            training_config_hash=data.get("training_config_hash", ""),
+        )
+
+    def save(self, path: str | Path) -> None:
+        """Save to JSON file."""
+        path = Path(path)
+        path.write_text(json.dumps(self.to_dict(), indent=2, sort_keys=True))
+
+    @classmethod
+    def load(cls, path: str | Path) -> "ExportMetadata":
+        """Load from JSON file."""
+        path = Path(path)
+        data = json.loads(path.read_text())
+        return cls.from_dict(data)
+
+    @classmethod
+    def default_timestamp(cls) -> str:
+        """Return an ISO timestamp for export."""
+        return datetime.now().isoformat()
+
+
+def hash_file(path: str | Path) -> str:
+    """Compute SHA256 hash for a file."""
+    path = Path(path)
+    hasher = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(8192), b""):
+            hasher.update(chunk)
+    return hasher.hexdigest()
+
+
+def build_base_metadata(
+    *,
+    export_format: str,
+    model_type: str,
+    model_name: str,
+    checkpoint_path: str,
+    backbone: str,
+    n_nodes: int,
+    n_edges: int,
+    node_names: List[str],
+    edge_inds: List[Tuple[int, int]],
+    input_scale: float,
+    input_channels: int,
+    output_stride: int,
+    crop_size: Optional[Tuple[int, int]] = None,
+    max_instances: Optional[int] = None,
+    max_peaks_per_node: Optional[int] = None,
+    max_batch_size: int = 1,
+    precision: str = "fp32",
+    training_config_hash: str = "",
+    training_config_embedded: bool = False,
+    input_dtype: str = "uint8",
+    normalization: str = "0_to_1",
+    n_classes: Optional[int] = None,
+    class_names: Optional[List[str]] = None,
+) -> ExportMetadata:
+    """Create an ExportMetadata instance with standard defaults."""
+    return ExportMetadata(
+        sleap_nn_version=__version__,
+        export_timestamp=ExportMetadata.default_timestamp(),
+        export_format=export_format,
+        model_type=model_type,
+        model_name=model_name,
+        checkpoint_path=checkpoint_path,
+        backbone=backbone,
+        n_nodes=n_nodes,
+        n_edges=n_edges,
+        node_names=node_names,
+        edge_inds=edge_inds,
+        input_scale=input_scale,
+        input_channels=input_channels,
+        output_stride=output_stride,
+        crop_size=crop_size,
+        max_instances=max_instances,
+        max_peaks_per_node=max_peaks_per_node,
+        max_batch_size=max_batch_size,
+        precision=precision,
+        input_dtype=input_dtype,
+        normalization=normalization,
+        n_classes=n_classes,
+        class_names=class_names,
+        training_config_embedded=training_config_embedded,
+        training_config_hash=training_config_hash,
+    )
+
+
+def embed_metadata_in_onnx(
+    model_path: str | Path,
+    metadata: ExportMetadata,
+    training_config_text: Optional[str] = None,
+) -> None:
+    """Embed metadata into an ONNX model file.
+
+    Raises ImportError if onnx is unavailable.
+    """
+    import onnx  # local import to keep dependency optional
+
+    model_path = Path(model_path)
+    model = onnx.load(model_path.as_posix())
+    model.metadata_props.append(
+        onnx.StringStringEntryProto(
+            key="sleap_nn_metadata", value=json.dumps(metadata.to_dict())
+        )
+    )
+    if training_config_text:
+        model.metadata_props.append(
+            onnx.StringStringEntryProto(
+                key="training_config", value=training_config_text
+            )
+        )
+    onnx.save(model, model_path.as_posix())
+
+
+def extract_metadata_from_onnx(model_path: str | Path) -> ExportMetadata:
+    """Extract metadata from an ONNX model file.
+
+    Raises ValueError if metadata is missing.
+    """
+    import onnx  # local import to keep dependency optional
+
+    model = onnx.load(Path(model_path).as_posix())
+    for prop in model.metadata_props:
+        if prop.key == "sleap_nn_metadata":
+            return ExportMetadata.from_dict(json.loads(prop.value))
+    raise ValueError("No sleap_nn metadata found in ONNX model")