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:

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.

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:

(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.

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.

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

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.