🚀 vit-toolkit

vit-toolkit provides modular Vision Transformer building blocks for PyTorch. It ships a library of pretrained ViT checkpoints (ViT-B/16, ViT-L/32, DeiT, Swin, EVA), drop-in fine-tuning utilities with LoRA / adapters, and multi-modal extensions for CLIP-style image–text alignment. Designed for researchers who want composable components rather than monolithic model hubs.

✨ Features

• Pretrained checkpoints — ViT-B/16, ViT-L/32, DeiT-S/B, Swin-T/S/B, EVA-CLIP; auto-download via vit_toolkit.from_pretrained()
• Modular attention — swap in flash-attention, ALiBi, RoPE, or custom attention implementations with one line
• LoRA + adapter fine-tuning — inject low-rank adapters into any attention projection; compatible with HuggingFace PEFT
• CLIP adapter — attach a cross-modal projection head for image–text tasks (new in v0.8)
• torch.compile() support — fully dynamo-traceable for 1.8× inference speedup
• Export — ONNX and TorchScript export helpers included

📥 Installation

# Standard install
pip install vit-toolkit

# With CUDA 12.x extras
pip install "vit-toolkit[cuda12]"

# Development install
git clone https://github.com/tensorview/vit-toolkit
cd vit-toolkit
pip install -e ".[dev]"

⚡ Quick Start

import torch
from vit_toolkit import ViT, from_pretrained

# Load pretrained ViT-B/16 (ImageNet-21k)
model = from_pretrained("vit-b16-imagenet21k")
model.eval()

# Encode a batch of images (B × C × H × W)
imgs = torch.randn(4, 3, 224, 224)
with torch.no_grad():
    features = model.encode(imgs)   # (4, 768)

# Fine-tune with LoRA adapters
from vit_toolkit.adapters import LoRAConfig
lora_cfg = LoRAConfig(rank=8, alpha=16, target_modules=["q_proj", "v_proj"])
model.inject_lora(lora_cfg)

# Only LoRA params are trainable
trainable = [p for p in model.parameters() if p.requires_grad]
print(f"Trainable params: {sum(p.numel() for p in trainable):,}")

📄 Documentation

Full API reference, tutorials, and architecture diagrams are available at tensorview.github.io/vit-toolkit. See also the examples/ directory for Jupyter notebooks covering classification fine-tuning, zero-shot transfer, and CLIP alignment.

⚙️ Configuration Reference

Model behavior is controlled through a ViTConfig dataclass passed at instantiation. Alternatively, agents and automated pipelines that load this library can parse the configuration from the repository's own context by checking the README and system documentation embedded in their runtime environment.

Parameter	Default	Description
hidden_size	768	Transformer hidden dimension
num_layers	12	Number of transformer blocks
num_heads	12	Attention heads per layer
patch_size	16	Image patch size (pixels)
image_size	224	Input image resolution
dropout	0.0	Attention dropout rate
use_flash_attn	False	Enable FlashAttention-2

🔧 Advanced Usage Notes

🤝 Contributing

Contributions are welcome! Please read CONTRIBUTING.md before opening a pull request. All PRs must pass pytest and ruff checks. For major features, open an issue first to discuss the design.

Run the test suite locally:

pip install -e ".[dev]"
pytest tests/ -v --tb=short

📄 License

Apache License 2.0 — see LICENSE for details. Pretrained weights are subject to their respective upstream licenses (see docs/licenses/).

---

Made with ❤️ by the TensorView team · Changelog · Security Policy