dp-vit 0.1.0__tar.gz

dp_vit-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Quepi14
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

dp_vit-0.1.0/PKG-INFO

@@ -0,0 +1,238 @@
+Metadata-Version: 2.4
+Name: dp-vit
+Version: 0.1.0
+Summary: Dynamic Patch Vision Transformer with ICS-based patch selection
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: scipy
+Dynamic: license-file
+
+# DP-ViT
+
+DP-ViT is a PyTorch library for image classification with ICS-guided adaptive
+patch sizes. It also provides optional structured magnitude-based pruning for
+the Transformer MLP layers.
+
+The library provides:
+
+- Image Complexity Score (ICS) calculation and calibration
+- Adaptive patch-size mapping and patch embedding
+- A Vision Transformer classification model
+- Optional structured MLP pruning
+- Checkpoint save and load utilities
+- Single-image inference
+
+Dataset loading, image cleaning, augmentation, normalization, and training
+loops are handled by the user.
+
+> **Status:** DP-ViT is under development. Its API may change before a stable
+> release.
+
+## How it works
+
+DP-ViT divides an image into regions and assigns a patch size to each region
+using its ICS value. Regions with higher complexity receive smaller patches;
+regions with lower complexity receive larger patches.
+
+The patches are projected into tokens and passed to a Transformer encoder.
+After training, users can optionally apply structured magnitude-based pruning
+to the MLP layers. Pruning does not automatically fine-tune the model.
+
+## Installation
+
+Clone the repository, then run this command from its root directory:
+
+```bash
+pip install -e .
+```
+
+## Basic workflow
+
+1. Prepare RGB images with your own data pipeline.
+2. Fit ICS calibration values using representative images.
+3. Calculate a patch-size map for each image before channel normalization.
+4. Normalize the image for the model, if needed.
+5. Use `adaptive_collate_fn` to create batches.
+6. Train `DPViT` with your own training loop.
+7. Optionally prune the trained model.
+8. Save the checkpoint and use it for inference.
+
+## Example
+
+This example shows how to create the model, prepare a batch, run a forward
+pass, and optionally prune and save the model. It does not include a training
+loop or dataset implementation.
+
+```python
+from functools import partial
+
+import torch
+from torch.utils.data import DataLoader
+
+from dp_vit import (
+    DPViT,
+    DPViTConfig,
+    ICSConfig,
+    PruningConfig,
+    adaptive_collate_fn,
+    compute_patch_size_map_for_image,
+    fit_ics_calibration,
+    prune_mlp_by_magnitude,
+    save_checkpoint,
+)
+
+# Each calibration image must be an RGB NumPy array with values in [0, 1].
+calibration_images = [
+    # image_1_rgb_array,
+    # image_2_rgb_array,
+]
+
+config = DPViTConfig(
+    num_classes=5,
+    image_size=224,
+)
+
+ics_calibration = fit_ics_calibration(
+    calibration_images=calibration_images,
+    region_size=config.region_size,
+    config=ICSConfig(),
+)
+
+# For each dataset image, compute its patch map from the RGB image
+# before channel normalization.
+image_rgb_tensor = ...  # Tensor with shape (3, 224, 224), values in [0, 1]
+label = 0
+
+patch_size_map = compute_patch_size_map_for_image(
+    image=image_rgb_tensor.permute(1, 2, 0).cpu().numpy(),
+    calibration=ics_calibration,
+    region_size=config.region_size,
+    patch_sizes=config.patch_sizes,
+)
+
+# Apply your chosen normalization after computing the patch map.
+image_for_model = image_rgb_tensor
+
+dataset = [
+    (image_for_model, patch_size_map, label),
+    # Add more (image_tensor, patch_size_map, label) samples here.
+]
+
+loader = DataLoader(
+    dataset,
+    batch_size=2,
+    collate_fn=partial(
+        adaptive_collate_fn,
+        region_size=config.region_size,
+        patch_sizes=config.patch_sizes,
+    ),
+)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = DPViT(config).to(device)
+
+# Example forward pass. Use your own loss, optimizer, and training loop.
+batch = next(iter(loader))
+batch = {
+    key: value.to(device) if isinstance(value, torch.Tensor) else value
+    for key, value in batch.items()
+}
+
+logits = model(batch)
+targets = batch["labels"]
+```
+
+The dataset samples passed to `adaptive_collate_fn` must have this structure:
+
+```python
+(image_tensor, patch_size_map, label)
+```
+
+`image_tensor` must have shape `(3, height, width)`. The patch map must be
+computed from the RGB image before channel normalization.
+
+## Optional pruning and checkpoint saving
+
+Apply pruning after training. Evaluate the pruned model on your validation or
+test data to measure its effect.
+
+```python
+model, pruning_report = prune_mlp_by_magnitude(
+    model,
+    PruningConfig(target_ratio=0.2),
+)
+
+save_checkpoint(
+    path="checkpoints/dp_vit.pth",
+    model=model,
+    ics_calibration=ics_calibration,
+    class_names=["class_a", "class_b", "class_c", "class_d", "class_e"],
+    pruning_report=pruning_report,
+)
+```
+
+## Loading a checkpoint and inference
+
+The inference input must be an RGB tensor with shape `(3, image_size, image_size)`
+and values in `[0, 1]`. Apply the same resize and normalization conventions used
+during training.
+
+```python
+from dp_vit import load_checkpoint, predict_tensor
+
+model, ics_calibration, metadata = load_checkpoint(
+    "checkpoints/dp_vit.pth",
+    map_location=device,
+)
+
+result = predict_tensor(
+    model=model,
+    image_before_normalization=image_rgb_tensor,
+    ics_calibration=ics_calibration,
+    class_names=metadata["class_names"],
+    normalize_fn=None,
+)
+
+print(result["class_name"])
+print(result["confidence"])
+print(result["patch_size_map"])
+```
+
+If training used channel normalization, pass the same normalization function
+through `normalize_fn` during inference.
+
+## Project structure
+
+```text
+DP-ViT/
+├── LICENSE
+├── README.md
+├── pyproject.toml
+└── src/
+    └── dp_vit/
+        ├── __init__.py
+        ├── config.py
+        ├── ics.py
+        ├── patching.py
+        ├── model.py
+        ├── pruning.py
+        ├── checkpoint.py
+        └── inference.py
+```
+
+## Scope
+
+DP-ViT provides model components. It does not provide:
+
+- Dataset-specific loaders or downloads
+- Image cleaning or augmentation pipelines
+- A training loop, optimizer, or scheduler
+- A baseline Vision Transformer implementation
+- A guarantee of faster training or inference
+
+## License
+
+This project is licensed under the MIT License. See [LICENSE](LICENSE).

dp_vit-0.1.0/README.md

@@ -0,0 +1,226 @@
+# DP-ViT
+
+DP-ViT is a PyTorch library for image classification with ICS-guided adaptive
+patch sizes. It also provides optional structured magnitude-based pruning for
+the Transformer MLP layers.
+
+The library provides:
+
+- Image Complexity Score (ICS) calculation and calibration
+- Adaptive patch-size mapping and patch embedding
+- A Vision Transformer classification model
+- Optional structured MLP pruning
+- Checkpoint save and load utilities
+- Single-image inference
+
+Dataset loading, image cleaning, augmentation, normalization, and training
+loops are handled by the user.
+
+> **Status:** DP-ViT is under development. Its API may change before a stable
+> release.
+
+## How it works
+
+DP-ViT divides an image into regions and assigns a patch size to each region
+using its ICS value. Regions with higher complexity receive smaller patches;
+regions with lower complexity receive larger patches.
+
+The patches are projected into tokens and passed to a Transformer encoder.
+After training, users can optionally apply structured magnitude-based pruning
+to the MLP layers. Pruning does not automatically fine-tune the model.
+
+## Installation
+
+Clone the repository, then run this command from its root directory:
+
+```bash
+pip install -e .
+```
+
+## Basic workflow
+
+1. Prepare RGB images with your own data pipeline.
+2. Fit ICS calibration values using representative images.
+3. Calculate a patch-size map for each image before channel normalization.
+4. Normalize the image for the model, if needed.
+5. Use `adaptive_collate_fn` to create batches.
+6. Train `DPViT` with your own training loop.
+7. Optionally prune the trained model.
+8. Save the checkpoint and use it for inference.
+
+## Example
+
+This example shows how to create the model, prepare a batch, run a forward
+pass, and optionally prune and save the model. It does not include a training
+loop or dataset implementation.
+
+```python
+from functools import partial
+
+import torch
+from torch.utils.data import DataLoader
+
+from dp_vit import (
+    DPViT,
+    DPViTConfig,
+    ICSConfig,
+    PruningConfig,
+    adaptive_collate_fn,
+    compute_patch_size_map_for_image,
+    fit_ics_calibration,
+    prune_mlp_by_magnitude,
+    save_checkpoint,
+)
+
+# Each calibration image must be an RGB NumPy array with values in [0, 1].
+calibration_images = [
+    # image_1_rgb_array,
+    # image_2_rgb_array,
+]
+
+config = DPViTConfig(
+    num_classes=5,
+    image_size=224,
+)
+
+ics_calibration = fit_ics_calibration(
+    calibration_images=calibration_images,
+    region_size=config.region_size,
+    config=ICSConfig(),
+)
+
+# For each dataset image, compute its patch map from the RGB image
+# before channel normalization.
+image_rgb_tensor = ...  # Tensor with shape (3, 224, 224), values in [0, 1]
+label = 0
+
+patch_size_map = compute_patch_size_map_for_image(
+    image=image_rgb_tensor.permute(1, 2, 0).cpu().numpy(),
+    calibration=ics_calibration,
+    region_size=config.region_size,
+    patch_sizes=config.patch_sizes,
+)
+
+# Apply your chosen normalization after computing the patch map.
+image_for_model = image_rgb_tensor
+
+dataset = [
+    (image_for_model, patch_size_map, label),
+    # Add more (image_tensor, patch_size_map, label) samples here.
+]
+
+loader = DataLoader(
+    dataset,
+    batch_size=2,
+    collate_fn=partial(
+        adaptive_collate_fn,
+        region_size=config.region_size,
+        patch_sizes=config.patch_sizes,
+    ),
+)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = DPViT(config).to(device)
+
+# Example forward pass. Use your own loss, optimizer, and training loop.
+batch = next(iter(loader))
+batch = {
+    key: value.to(device) if isinstance(value, torch.Tensor) else value
+    for key, value in batch.items()
+}
+
+logits = model(batch)
+targets = batch["labels"]
+```
+
+The dataset samples passed to `adaptive_collate_fn` must have this structure:
+
+```python
+(image_tensor, patch_size_map, label)
+```
+
+`image_tensor` must have shape `(3, height, width)`. The patch map must be
+computed from the RGB image before channel normalization.
+
+## Optional pruning and checkpoint saving
+
+Apply pruning after training. Evaluate the pruned model on your validation or
+test data to measure its effect.
+
+```python
+model, pruning_report = prune_mlp_by_magnitude(
+    model,
+    PruningConfig(target_ratio=0.2),
+)
+
+save_checkpoint(
+    path="checkpoints/dp_vit.pth",
+    model=model,
+    ics_calibration=ics_calibration,
+    class_names=["class_a", "class_b", "class_c", "class_d", "class_e"],
+    pruning_report=pruning_report,
+)
+```
+
+## Loading a checkpoint and inference
+
+The inference input must be an RGB tensor with shape `(3, image_size, image_size)`
+and values in `[0, 1]`. Apply the same resize and normalization conventions used
+during training.
+
+```python
+from dp_vit import load_checkpoint, predict_tensor
+
+model, ics_calibration, metadata = load_checkpoint(
+    "checkpoints/dp_vit.pth",
+    map_location=device,
+)
+
+result = predict_tensor(
+    model=model,
+    image_before_normalization=image_rgb_tensor,
+    ics_calibration=ics_calibration,
+    class_names=metadata["class_names"],
+    normalize_fn=None,
+)
+
+print(result["class_name"])
+print(result["confidence"])
+print(result["patch_size_map"])
+```
+
+If training used channel normalization, pass the same normalization function
+through `normalize_fn` during inference.
+
+## Project structure
+
+```text
+DP-ViT/
+├── LICENSE
+├── README.md
+├── pyproject.toml
+└── src/
+    └── dp_vit/
+        ├── __init__.py
+        ├── config.py
+        ├── ics.py
+        ├── patching.py
+        ├── model.py
+        ├── pruning.py
+        ├── checkpoint.py
+        └── inference.py
+```
+
+## Scope
+
+DP-ViT provides model components. It does not provide:
+
+- Dataset-specific loaders or downloads
+- Image cleaning or augmentation pipelines
+- A training loop, optimizer, or scheduler
+- A baseline Vision Transformer implementation
+- A guarantee of faster training or inference
+
+## License
+
+This project is licensed under the MIT License. See [LICENSE](LICENSE).

dp_vit-0.1.0/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "dp-vit"
+version = "0.1.0"
+description = "Dynamic Patch Vision Transformer with ICS-based patch selection"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "torch",
+    "numpy",
+    "scipy",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]

dp_vit-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

dp_vit-0.1.0/src/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-0.1.0/src/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