dp-vit 0.1.0__py3-none-any.whl

dp_vit/__init__.py

@@ -0,0 +1,34 @@
+from .checkpoint import load_checkpoint, save_checkpoint
+from .config import DPViTConfig, ICSCalibration, ICSConfig, PruningConfig
+from .ics import (
+    compute_gradient_magnitude,
+    compute_ics,
+    compute_local_variance,
+    compute_patch_size_map_for_image,
+    determine_patch_size_map,
+    fit_ics_calibration,
+)
+from .inference import predict_tensor
+from .model import DPViT
+from .patching import AdaptivePatchEmbedding, adaptive_collate_fn
+from .pruning import prune_mlp_by_magnitude
+
+__all__ = [
+    "AdaptivePatchEmbedding",
+    "DPViT",
+    "DPViTConfig",
+    "ICSCalibration",
+    "ICSConfig",
+    "PruningConfig",
+    "adaptive_collate_fn",
+    "compute_gradient_magnitude",
+    "compute_ics",
+    "compute_local_variance",
+    "compute_patch_size_map_for_image",
+    "determine_patch_size_map",
+    "fit_ics_calibration",
+    "load_checkpoint",
+    "predict_tensor",
+    "prune_mlp_by_magnitude",
+    "save_checkpoint",
+]

dp_vit/checkpoint.py

@@ -0,0 +1,147 @@
+"""Checkpoint saving and loading for DP-ViT."""
+
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, cast
+
+import torch
+from torch import nn
+
+from .config import DPViTConfig, ICSCalibration
+from .model import DPViT
+
+CHECKPOINT_FORMAT_VERSION = 1
+
+def _get_mlp_widths(model: DPViT) -> list[int]:
+    """Return the hidden width of every encoder MLP."""
+    return [
+        int(getattr(getattr(layer, "linear1"), "out_features"))
+        for layer in model.encoder.layers
+    ]
+
+def _resize_mlp_layer(
+    layer: nn.TransformerEncoderLayer,
+    hidden_width: int,
+) -> None:
+    """Create MLP layers with a target width before leading checkpoint weights"""
+    if hidden_width <= 0:
+        raise ValueError("Each MLP hidden width must be greater than zero.")
+
+    old_linear1 = layer.linear1
+    old_linear2 = layer.linear2
+
+    new_linear1 = nn.Linear(
+        in_features=old_linear1.in_features,
+        out_features=hidden_width,
+        bias=old_linear1.bias is not None,
+        device=old_linear1.weight.device,
+        dtype=old_linear1.weight.dtype,
+    )
+    new_linear2 = nn.Linear(
+        in_features=hidden_width,
+        out_features=old_linear2.out_features,
+        bias=old_linear2.bias is not None,
+        device=old_linear2.weight.device,
+        dtype=old_linear2.weight.dtype,
+    )
+
+    layer.linear1 = new_linear1
+    layer.linear2 = new_linear2
+
+def save_checkpoint(
+    path: str | Path,
+    model: DPViT,
+    ics_calibration: ICSCalibration,
+    class_names: list[str] | None = None,
+    pruning_report: dict[str, Any] | None = None,
+) -> Path:
+    """Save a DP-ViT checkpoint and the metadata needed to rebuild it."""
+    if not isinstance(model, DPViT):
+        raise TypeError("model must be an instance of DPViT")
+
+    if not isinstance(ics_calibration, ICSCalibration):
+        raise TypeError("ics_calibration must be an ICSCalibration instance.")
+
+    if class_names is not None:
+        if len(class_names) != model.config.num_classes:
+            raise ValueError(
+                "The number of class names must match num_classes."
+            )
+
+    checkpoint_path = Path(path).expanduser()
+    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+
+    payload = {
+        "format_version": CHECKPOINT_FORMAT_VERSION,
+        "model_config": asdict(model.config),
+        "ics_calibration": asdict(ics_calibration),
+        "class_names": class_names,
+        "mlp_widths": _get_mlp_widths(model),
+        "pruning_report": pruning_report,
+        "model_state_dict": {
+            name: value.detach().cpu()
+            for name, value in model.state_dict().items()
+        },
+    }
+
+    temporary_path = checkpoint_path.with_suffix(
+        checkpoint_path.suffix + ".tmp"
+    )
+    torch.save(payload, temporary_path)
+    temporary_path.replace(checkpoint_path)
+
+    return checkpoint_path
+
+def load_checkpoint(
+    path: str | Path,
+    map_location: str | torch.device = "cpu",
+) -> tuple[DPViT, ICSCalibration, dict[str, Any]]:
+    """
+    Load a DP-ViT checkpoint and reconstruct its architecture.
+
+    Returns:
+        model: The reconstructured DP-ViT model.
+        ics_calibration: The ICS parameters saved with the model.
+        metadata: Class names, MLP widths, and pruning report.
+    """
+    checkpoint_path = Path(path).expanduser()
+
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"Checkpoint file does not exist: {checkpoint_path}")
+
+    payload = torch.load(
+        checkpoint_path,
+        map_location=map_location,
+        weights_only=True,
+    )
+
+    if payload.get("format_version") != CHECKPOINT_FORMAT_VERSION:
+        raise ValueError("Unsupported or missing DP-ViT checkpoint format version.")
+
+    config = DPViTConfig(**payload["model_config"])
+    calibration = ICSCalibration(**payload["ics_calibration"])
+
+    model = DPViT(config)
+
+    mlp_widths = payload["mlp_widths"]
+    encoder_layers = list(model.encoder.layers)
+
+    if len(mlp_widths) != len(encoder_layers):
+        raise ValueError("Checkpoint MLP widths do not match the encoder depth.")
+
+    for layer, hidden_width in zip(encoder_layers, mlp_widths):
+        _resize_mlp_layer(
+            cast(nn.TransformerEncoderLayer, layer),
+            int(hidden_width),
+        )
+
+    model.load_state_dict(payload["model_state_dict"])
+    model.to(map_location)
+
+    metadata = {
+        "class_names": payload.get("class_names"),
+        "mlp_widths": mlp_widths,
+        "pruning_report": payload.get("pruning_report"),
+    }
+
+    return model, calibration, metadata

dp_vit/config.py

@@ -0,0 +1,116 @@
+"""Configuration objects for the DP-ViT Library."""
+
+from dataclasses import dataclass
+import math
+
+@dataclass(frozen=True)
+class DPViTConfig:
+    """Architecture Dynamic Patch Vision Transformer"""
+
+    num_classes: int
+
+    image_size: int = 224
+    region_size: int = 32
+    patch_sizes: tuple[int, int, int] = (8, 16, 32)
+
+    embed_dim: int = 384
+    num_layers: int = 12
+    num_heads: int = 6
+    mlp_ratio: float = 4.0
+    dropout: float = 0.1
+
+    def __post_init__(self) -> None:
+        if self.num_classes <= 0:
+            raise ValueError("num_classes must be bigger than 0")
+
+        if self.image_size <= 0:
+            raise ValueError("image_size must be bigger than 0")
+
+        if self.region_size <= 0:
+            raise ValueError("region_size must be bigger than 0")
+
+        if self.image_size % self.region_size !=0:
+            raise ValueError("image_size must be divisible by region_size")
+
+        if self.patch_sizes != (8, 16, 32):
+            raise ValueError("patch_size must be (8, 16, 32)")
+
+        if any(self.region_size % size != 0 for size in self.patch_sizes):
+            raise ValueError("region_size must be divisible by all patch sizes")
+
+        if self.embed_dim <= 0 or self.num_layers <= 0 or self.num_heads <= 0:
+            raise ValueError("embed_dim, num_layers, and num_heads must be bigger than 0")
+
+        if self.embed_dim % self.num_heads != 0:
+            raise ValueError("embed_dim must be divisible by num_heads")
+
+        if not math.isfinite(self.mlp_ratio) or self.mlp_ratio <= 0:
+            raise ValueError("mlp_ratio must be positive number")
+
+        if not math.isfinite(self.dropout) or not 0 <= self.dropout < 1:
+            raise ValueError("dropout must be in range [0, 1].")
+
+    @property
+    def max_position_cell(self) -> int:
+        """Maximum number of position cells."""
+        smallest_patch = min(self.patch_sizes)
+        return (self.image_size // smallest_patch) ** 2
+
+@dataclass(frozen=True)
+class ICSConfig:
+    "Config percentile for ICS module"
+
+    low_percentile: float = 30.0
+    high_percentile: float = 70.0
+
+    def __post_init__(self) -> None:
+        if not 0 <= self.low_percentile <= 100:
+            raise ValueError("low_percentile must be in range [0, 100]")
+
+        if not 0 <= self.high_percentile <= 100:
+            raise ValueError("high_percentile must be in range [0, 100]")
+
+        if self.low_percentile >= self.high_percentile:
+            raise ValueError("low_percentile must be smaller than high_percentile.")
+
+        if ( self.low_percentile != 30.0 or self.high_percentile != 70.0):
+            raise ValueError("The current ICSCalibration supports only the 30th and 70th percentile.")
+
+@dataclass(frozen=True)
+class ICSCalibration:
+    """Calibration ICS value used for mapping patch."""
+
+    alpha: float
+    beta: float
+    p30: float
+    p70: float
+
+    def __post_init__ (self) -> None:
+        for name, value in (
+            ("alpha", self.alpha),
+            ("beta", self.beta),
+            ("p30", self.p30),
+            ("p70", self.p70),
+        ):
+            if not math.isfinite(value):
+                raise ValueError(f"{name} must be finite number"
+        )
+
+        if self.alpha <= 0 or self.beta <= 0:
+            raise ValueError("alpha and beta must be greater than 0.")
+
+        if self.p30 > self.p70:
+            raise ValueError("p30 must be smaller than p70.")
+
+@dataclass(frozen=True)
+class PruningConfig:
+    """Configuration for neuron MLP structured magnitude-based pruning."""
+
+    target_ratio: float = 0.2
+
+    def __post_init__(self) -> None:
+        if not math.isfinite(self.target_ratio):
+            raise ValueError("target_ratio must be finite number")
+
+        if not 0 < self.target_ratio < 1:
+            raise ValueError("target_ratio must be greater than 0 and smaller than 1.")

dp_vit/ics.py

@@ -0,0 +1,215 @@
+"""INformation Complexity Score (ICS) calculations for DP-ViT."""
+
+from collections.abc import Sequence
+
+import numpy as np
+from scipy.ndimage import sobel
+
+from .config import ICSCalibration, ICSConfig
+
+def _to_grayscale(image: np.ndarray) -> np.ndarray:
+    """Convert a grayscale or RGB image to a 2D float64 intensity array."""
+    image = np.asarray(image)
+
+    if image.ndim == 2:
+        gray = image.astype(np.float64)
+    elif image.ndim == 3 and image.shape[2] == 3:
+        image = image.astype(np.float64)
+        gray = (
+            0.299 * image[..., 0]
+            + 0.587 * image[..., 1]
+            + 0.114 * image[..., 2]
+        )
+    else:
+        raise ValueError("image must have shape (height, width) or (height, width, 3).")
+
+    if not np.isfinite(gray).all():
+        raise ValueError("image must contain only finite values.")
+
+    return gray
+
+def _make_region_blocks(
+    image: np.ndarray,
+    region_size: int,
+) -> np.ndarray:
+    """Split a 2D image into non-overlapping square regions."""
+    if isinstance(region_size, bool) or not isinstance(region_size, int):
+        raise TypeError("region_size must be an integer.")
+
+    if region_size <= 0:
+        raise ValueError("region_size must be greater than zero.")
+
+    height, width = image.shape
+
+    if height % region_size != 0 or width % region_size != 0:
+        raise ValueError("Image height and width must be divisible by region_size.")
+
+    rows = height // region_size
+    cols = width // region_size
+
+    blocks = image.reshape(
+        rows, region_size, cols, region_size
+    )
+    blocks = blocks.transpose(0, 2, 1, 3)
+
+    return blocks.reshape(rows, cols, region_size * region_size)
+
+def compute_local_variance(
+    image: np.ndarray,
+    region_size: int = 32,
+) -> np.ndarray:
+    """
+    Compute the population variance of grayscale intensity in each region.
+    
+    The variance for each region is
+        Var(R) = mean((pixel - region_mean) ** 2)
+    """
+    gray = _to_grayscale(image)
+    blocks = _make_region_blocks(gray, region_size)
+
+    region_means = blocks.mean(axis=-1, keepdims=True)
+    variance_map = np.mean((blocks - region_means) ** 2, axis=-1)
+
+    return variance_map
+
+def compute_gradient_magnitude(
+    image: np.ndarray,
+    region_size: int = 32,
+) -> np.ndarray:
+    """ 
+    Compute the mean Sobel gradient magnitude in each region.
+
+    Sobel gradients are calculated on full image before region averaging.
+    """
+    gray = _to_grayscale(image)
+
+    gradient_x = sobel(gray, axis=1)
+    gradient_y = sobel(gray, axis=0)
+    magnitude = np.sqrt(gradient_x ** 2 + gradient_y ** 2)
+
+    blocks = _make_region_blocks(magnitude, region_size)
+    return blocks.mean(axis=-1)
+
+def compute_ics(
+    variance_map: np.ndarray,
+    gradient_map: np.ndarray,
+    alpha: float,
+    beta: float,
+) -> np.ndarray:
+    """Combine local variance and gradient magnitude into ICS map."""
+    variance_map = np.asarray(variance_map, dtype=np.float64)
+    gradient_map = np.asarray(gradient_map, dtype=np.float64)
+
+    if variance_map.shape != gradient_map.shape:
+        raise ValueError("variance_map and gradient_map must have the same shape.")
+
+    if not np.isfinite(alpha) or alpha <= 0:
+        raise ValueError("alpha must be a finite positive number.")
+
+    if not np.isfinite(beta) or beta <= 0:
+        raise ValueError("beta must be a finite positive number.")
+
+    return alpha * variance_map + beta * gradient_map
+
+def fit_ics_calibration(
+    calibration_images: Sequence[np.ndarray],
+    region_size: int = 32,
+    config: ICSConfig | None = None,
+) -> ICSCalibration:
+    """
+    Fit ICS weights and global percentile thresholds using calibration images.
+
+    All images must use the same intensity range and preprocessing convention.
+    """
+    if not calibration_images:
+        raise ValueError("calibration_images must not be empty.")
+
+    if config is None:
+        config = ICSConfig()
+
+    all_variances = []
+    all_gradients = []
+
+    for image in calibration_images:
+        all_variances.append(
+            compute_local_variance(image, region_size).ravel()
+        )
+        all_gradients.append(
+            compute_gradient_magnitude(image, region_size).ravel()
+        )
+
+    all_variances = np.concatenate(all_variances)
+    all_gradients = np.concatenate(all_gradients)
+
+    epsilon = 1e-8
+    alpha = 1.0 / (all_variances.std() + epsilon)
+    beta = 1.0 / (all_gradients.std() + epsilon)
+
+    all_ics = alpha * all_variances + beta * all_gradients
+
+    p30 = float(np.percentile(all_ics, config.low_percentile))
+    p70 = float(np.percentile(all_ics, config.high_percentile))
+
+    return ICSCalibration(
+        alpha=float(alpha),
+        beta=float(beta),
+        p30=p30,
+        p70=p70,
+    )
+
+def determine_patch_size_map(
+    ics_map: np.ndarray,
+    calibration: ICSCalibration,
+    patch_sizes: tuple[int, int, int] = (8, 16, 32),
+) -> np.ndarray:
+    """
+    Map each ICS region to a patch size.
+
+    The patch_sizes tuple is ordewred as (small, medium, large).
+    """
+    if len(patch_sizes) != 3:
+        raise ValueError("patch_sizes must contain exactly three values.")
+
+    small_patch, medium_patch, large_patch = patch_sizes
+
+    if not small_patch < medium_patch < large_patch:
+        raise ValueError("patch_sizes must be ordered from smallest to largest.")
+
+    size_map = np.full(ics_map.shape, large_patch, dtype=np.int32)
+
+    medium_mask = (
+        (ics_map >= calibration.p30)
+        & (ics_map < calibration.p70)
+    )
+
+    small_mask = ics_map >= calibration.p70
+
+    size_map[medium_mask] = medium_patch
+    size_map[small_mask] = small_patch
+
+    return size_map
+
+def compute_patch_size_map_for_image(
+    image: np.ndarray,
+    calibration: ICSCalibration,
+    region_size: int = 32,
+    patch_sizes: tuple[int, int, int] = (8, 16, 32),
+) -> np.ndarray:
+    """
+    Compute an iamge's ICS map and convert it to an adaptive patch map.
+    """
+    variance_map = compute_local_variance(image, region_size)
+    gradient_map = compute_gradient_magnitude(image, region_size)
+
+    ics_map = compute_ics(
+        variance_map,
+        gradient_map,
+        alpha=calibration.alpha,
+        beta=calibration.beta,
+    )
+
+    return determine_patch_size_map(
+        ics_map,
+        calibration=calibration,
+        patch_sizes=patch_sizes
+    )

dp_vit/inference.py

@@ -0,0 +1,132 @@
+"""Single-image inference utilities for DP-ViT."""
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import numpy as np
+import torch 
+from torch import Tensor
+
+from .config import ICSCalibration
+from .ics import compute_patch_size_map_for_image
+from .model import DPViT
+from .patching import adaptive_collate_fn
+
+def predict_tensor(
+    model: DPViT,
+    image_before_normalization: Tensor,
+    ics_calibration: ICSCalibration,
+    class_names: Sequence[str] | None = None,
+    normalize_fn : Callable[[Tensor], Tensor] | None = None,
+) -> dict[str, Any]:
+    """
+    Predict the class of one preprocessed image tensor.
+
+    The input must be an RGB tensor in [0, 1], before channel normalization.
+    Resizing and other user-selected preprocessing must be applied beforehand.
+    """
+    if not isinstance(model, DPViT):
+        raise TypeError("model must be an instance of DPViT.")
+
+    if not isinstance(image_before_normalization, Tensor):
+        raise TypeError("image_before_normalization must be a torch.Tensor.")
+
+    if image_before_normalization.ndim != 3:
+        raise ValueError("The input image must have shape (channels, height, width).")
+
+    channels, height, width = image_before_normalization.shape
+
+    if channels != 3:
+        raise ValueError("DP-ViT currently expects an RGB image with 3 channels.")
+
+    if (height, width) != (model.config.image_size, model.config.image_size):
+        raise ValueError(
+            f"Expected an image of size {model.config.image_size}x"
+            f"{model.config.image_size}, got {height}x{width}."
+        )
+
+    if not torch.isfinite(image_before_normalization).all():
+        raise ValueError("The input image must contain only finite values.")
+
+    if (
+        image_before_normalization.min() <0
+        or image_before_normalization.max() > 1
+    ):
+        raise ValueError(
+            "ICS input must be in [0, 1]. Apply channel normalization"
+            "after ICS patch-map generation."
+        )
+
+    if not isinstance(ics_calibration, ICSCalibration):
+        raise TypeError("ics_calibration must be an ICSCalibration instance.")
+
+    if class_names is not None and len(class_names) != model.config.num_classes:
+        raise ValueError("class_names length must match model.num_classes.")
+
+    image = image_before_normalization.detach().to(
+        device="cpu",
+        dtype=torch.float32,
+    )
+
+    image_np = image.permute(1, 2, 0).numpy()
+
+    patch_size_map = compute_patch_size_map_for_image(
+        image=image_np,
+        calibration=ics_calibration,
+        region_size=model.config.region_size,
+        patch_sizes=model.config.patch_sizes,
+    )
+
+    if normalize_fn is None:
+        model_image = image
+    else: 
+        model_image = normalize_fn(image)
+
+    if not isinstance(model_image, Tensor):
+        raise TypeError("normalize_fn must return a torch.Tensor.")
+
+    if model_image.shape != image.shape:
+        raise ValueError("normalize_fn must preserve the input image shape.")
+
+    batch = adaptive_collate_fn(
+        [(model_image, patch_size_map, 0)],
+        region_size=model.config.region_size,
+        patch_sizes=model.config.patch_sizes,
+    )
+
+    device = next(model.parameters()).device
+    batch = {
+        key: value.to(device) if isinstance(value, Tensor) else value
+        for key, value in batch.items()
+    }
+
+    was_training = model.training
+    model.eval()
+
+    try:
+        with torch.inference_mode():
+            logits = model(batch)
+            probabilities = torch.softmax(logits, dim=1)[0]
+            predicted_index = int(probabilities.argmax())
+    finally:
+        model.train(was_training)
+
+    patch_counts = {
+        int(size): int(np.count_nonzero(patch_size_map == size))
+        for size in model.config.patch_sizes
+    }
+
+    predicted_name = (
+        class_names[predicted_index]
+        if class_names is not None
+        else None
+    )
+
+    return {
+        "class_index": predicted_index,
+        "class_name": predicted_name,
+        "confidence": float(probabilities[predicted_index].cpu()),
+        "probabilities": probabilities.cpu().tolist(),
+        "patch_size_map": patch_size_map.tolist(),
+        "patch_counts": patch_counts,
+    }