Skip to main content

Vision LLMs with Efficient LoRA Fine-Tuning

Project description

langtrain: Vision LLMs (Large Language Models for Vision) with Efficient LoRA Fine-Tuning


Langtrain Logo

PyPI version Downloads License Code style: black

Modular Vision LLMs (Large Language Models for Vision) with Efficient LoRA Fine-Tuning
Build, adapt, and fine-tune vision models with ease and efficiency.


🚀 Quick Links


📚 Table of Contents


✨ Features

  • 🔧 Plug-and-play LoRA adapters for parameter-efficient fine-tuning
  • 🏗️ Modular Vision Transformer (ViT) backbone with customizable components
  • 🎯 Unified model zoo for open-source visual models
  • ⚙️ Easy configuration and extensible codebase
  • 🚀 Production ready with comprehensive testing and documentation
  • 💾 Memory efficient training with gradient checkpointing support
  • 📊 Built-in metrics and visualization tools
  • 🧩 Modular training loop with LoRA support
  • 🎯 Unified CLI for fine-tuning and evaluation
  • 🔌 Extensible callbacks (early stopping, logging, etc.)
  • 📦 Checkpointing and resume
  • 🚀 Mixed precision training
  • 🔧 Easy dataset and model extension
  • Ready for distributed/multi-GPU training

🚀 Showcase

langtrain is a modular, research-friendly framework for building and fine-tuning Vision Large Language Models (LLMs) with efficient Low-Rank Adaptation (LoRA) support. Whether you're working on image classification, visual question answering, or custom vision tasks, langtrain provides the tools you need for parameter-efficient model adaptation.


🏁 Getting Started

Here's a minimal example to get you up and running:

pip install langtrain-ai
import torch
from langtrain_ai.models.vision_transformer import VisionTransformer
from langtrain_ai.utils.config import default_config

# Create model
x = torch.randn(2, 3, 224, 224)
model = VisionTransformer(
    img_size=default_config['img_size'],
    patch_size=default_config['patch_size'],
    in_chans=default_config['in_chans'],
    num_classes=default_config['num_classes'],
    embed_dim=default_config['embed_dim'],
    depth=default_config['depth'],
    num_heads=default_config['num_heads'],
    mlp_ratio=default_config['mlp_ratio'],
    lora_config=default_config['lora'],
)

# Forward pass
with torch.no_grad():
    out = model(x)
    print('Output shape:', out.shape)

For advanced usage, CLI details, and more, see the Documentation and src/langtrain/cli/finetune.py.


🐍 Supported Python Versions

  • Python 3.8+

🧩 Why Langtrain?

  • Parameter-efficient fine-tuning: Plug-and-play LoRA adapters for fast, memory-efficient adaptation with minimal computational overhead
  • Modular ViT backbone: Swap or extend components like patch embedding, attention, or MLP heads with ease
  • Unified model zoo: Access and experiment with open-source visual models through a consistent interface
  • Research & production ready: Clean, extensible codebase with comprehensive configuration options and robust utilities
  • Memory efficient: Fine-tune large models on consumer hardware by updating only a small fraction of parameters

🏗️ Architecture Overview

langtrain is built around a modular Vision Transformer (ViT) backbone, with LoRA adapters strategically injected into attention and MLP layers for efficient fine-tuning. This approach allows you to adapt large pre-trained models using only a fraction of the original parameters.

Model Data Flow

---
config:
  layout: dagre
---
flowchart TD
 subgraph LoRA_Adapters["LoRA Adapters in Attention and MLP"]
        LA1(["LoRA Adapter 1"])
        LA2(["LoRA Adapter 2"])
        LA3(["LoRA Adapter N"])
  end
    A(["Input Image"]) --> B(["Patch Embedding"])
    B --> C(["CLS Token & Positional Encoding"])
    C --> D1(["Encoder Layer 1"])
    D1 --> D2(["Encoder Layer 2"])
    D2 --> D3(["Encoder Layer N"])
    D3 --> E(["LayerNorm"])
    E --> F(["MLP Head"])
    F --> G(["Output Class Logits"])
    LA1 -.-> D1
    LA2 -.-> D2
    LA3 -.-> D3
     LA1:::loraStyle
     LA2:::loraStyle
     LA3:::loraStyle
    classDef loraStyle fill:#e1f5fe,stroke:#0277bd,stroke-width:2px

Architecture Components

Legend:

  • Solid arrows: Main data flow through the Vision Transformer
  • Dashed arrows: LoRA adapter injection points in encoder layers
  • Blue boxes: LoRA adapters for parameter-efficient fine-tuning

Data Flow Steps:

  1. Input Image (224×224×3): Raw image data ready for processing
  2. Patch Embedding: Image split into 16×16 patches and projected to embedding dimension
  3. CLS Token & Positional Encoding: Classification token prepended with learnable position embeddings
  4. Transformer Encoder Stack: Multi-layer transformer with self-attention and MLP blocks
    • LoRA Integration: Low-rank adapters injected into attention and MLP layers
    • Efficient Updates: Only LoRA parameters updated during fine-tuning
  5. LayerNorm: Final normalization of encoder outputs
  6. MLP Head: Task-specific classification or regression head
  7. Output: Final predictions (class probabilities, regression values, etc.)

🧩 Core Modules

Module Description Key Features
PatchEmbedding Image-to-patch conversion and embedding • Configurable patch sizes
• Learnable position embeddings
• Support for different input resolutions
TransformerEncoder Multi-layer transformer backbone • Self-attention mechanisms
• LoRA adapter integration
• Gradient checkpointing support
LoRALinear Low-rank adaptation layers • Configurable rank and scaling
• Memory-efficient updates
• Easy enable/disable functionality
MLPHead Output projection layer • Multi-class classification
• Regression support
• Dropout regularization
Config System Centralized configuration management • YAML/JSON config files
• Command-line overrides
• Validation and defaults
Data Utils Preprocessing and augmentation • Built-in transforms
• Custom dataset loaders
• Efficient data pipelines

📊 Performance & Efficiency

LoRA Benefits

Metric Full Fine-tuning LoRA Fine-tuning Improvement
Trainable Parameters 86M 2.4M 97% reduction
Memory Usage 12GB 4GB 67% reduction
Training Time 4 hours 1.5 hours 62% faster
Storage per Task 344MB 9.6MB 97% smaller

Benchmarks on ViT-Base with CIFAR-100, RTX 3090

Supported Model Sizes

  • ViT-Tiny: 5.7M parameters, perfect for experimentation
  • ViT-Small: 22M parameters, good balance of performance and efficiency
  • ViT-Base: 86M parameters, strong performance across tasks
  • ViT-Large: 307M parameters, state-of-the-art results

🔧 Advanced Configuration

LoRA Configuration

lora_config = {
    "rank": 16,                    # Low-rank dimension
    "alpha": 32,                   # Scaling factor
    "dropout": 0.1,                # Dropout rate
    "target_modules": [            # Modules to adapt
        "attention.qkv",
        "attention.proj", 
        "mlp.fc1",
        "mlp.fc2"
    ],
    "merge_weights": False         # Whether to merge during inference
}

Training Configuration

# config.yaml
model:
  name: "vit_base"
  img_size: 224
  patch_size: 16
  num_classes: 1000

training:
  epochs: 10
  batch_size: 32
  learning_rate: 1e-4
  weight_decay: 0.01
  warmup_steps: 1000

lora:
  rank: 16
  alpha: 32
  dropout: 0.1

📚 Documentation & Resources

Research Papers


🧪 Testing & Quality

Run the comprehensive test suite:

# Unit tests
pytest tests/unit/

# Integration tests  
pytest tests/integration/

# Performance benchmarks
pytest tests/benchmarks/

# All tests with coverage
pytest tests/ --cov=langtrain --cov-report=html

Code Quality Tools

# Linting
flake8 src/
black src/ --check

# Type checking
mypy src/

# Security scanning
bandit -r src/

🚀 Examples & Use Cases

Image Classification

from langtrain_ai import VisionTransformer
from langtrain_ai.datasets import CIFAR10Dataset

# Load pre-trained model
model = VisionTransformer.from_pretrained("vit_base_patch16_224")

# Fine-tune on CIFAR-10
dataset = CIFAR10Dataset(train=True, transform=model.default_transform)
model.finetune(dataset, epochs=10, lora_rank=16)

Custom Dataset

from langtrain_ai.datasets import ImageFolderDataset

# Your custom dataset
dataset = ImageFolderDataset(
    root="/path/to/dataset",
    split="train",
    transform=model.default_transform
)

# Fine-tune with custom configuration
model.finetune(
    dataset, 
    config_path="configs/custom_config.yaml"
)

🧩 Extending the Framework

  • Add new datasets in src/langtrain/data/datasets.py
  • Add new callbacks in src/langtrain/callbacks/
  • Add new models in src/langtrain/models/
  • Add new CLI tools in src/langtrain/cli/

📖 Documentation

  • See code comments and docstrings for details on each module.
  • For advanced usage, see the src/langtrain/cli/finetune.py script.

🤝 Contributing

We welcome contributions from the community! Here's how you can get involved:

Ways to Contribute

  • 🐛 Report bugs by opening issues with detailed reproduction steps
  • 💡 Suggest features through feature requests and discussions
  • 📝 Improve documentation with examples, tutorials, and API docs
  • 🔧 Submit pull requests for bug fixes and new features
  • 🧪 Add tests to improve code coverage and reliability

Development Setup

# Clone and setup development environment
git clone https://github.com/langtrain-ai/langtrain.git
cd langtrain
pip install -e ".[dev]"

# Install pre-commit hooks
pre-commit install

# Run tests
pytest tests/

Community Resources

📄 License & Citation

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

If you use langtrain in your research, please cite:

@software{langtrain2025,
  author = {Pritesh Raj},
  title = {langtrain: Vision LLMs with Efficient LoRA Fine-Tuning},
  url = {https://github.com/langtrain-ai/langtrain},
  year = {2025},
  version = {1.0.0}
}

🌟 Acknowledgements

We thank the following projects and communities:

  • PyTorch - Deep learning framework
  • HuggingFace - Transformers and model hub
  • timm - Vision model implementations
  • PEFT - Parameter-efficient fine-tuning methods

Made in India 🇮🇳 with ❤️ by the langtrain team
Star ⭐ this repo if you find it useful!

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

langtrain_ai-0.0.3.tar.gz (26.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

langtrain_ai-0.0.3-py3-none-any.whl (25.2 kB view details)

Uploaded Python 3

File details

Details for the file langtrain_ai-0.0.3.tar.gz.

File metadata

  • Download URL: langtrain_ai-0.0.3.tar.gz
  • Upload date:
  • Size: 26.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.11

File hashes

Hashes for langtrain_ai-0.0.3.tar.gz
Algorithm Hash digest
SHA256 5f1f8b0c014aa331e9581a8e4ada2cdf38b613cf2c03eb38a1bbcb3d3f8b7152
MD5 d38147b1949609fc55a514558760c2dd
BLAKE2b-256 e0a443b14c86d81e4aa032b2c09548e84cf940f0be0234d21e719b33b0ec40f1

See more details on using hashes here.

File details

Details for the file langtrain_ai-0.0.3-py3-none-any.whl.

File metadata

  • Download URL: langtrain_ai-0.0.3-py3-none-any.whl
  • Upload date:
  • Size: 25.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.11

File hashes

Hashes for langtrain_ai-0.0.3-py3-none-any.whl
Algorithm Hash digest
SHA256 17c40800c42abe7514f99a56e8d47eac26e2e47207d4efb5b748591e9dcc73dc
MD5 30ab87a1014beed1e285536d4ef56bb0
BLAKE2b-256 928e6cc203ba033ec193e914cafbcd94c9618ea085b1d3abc36024d6d0a2ce52

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page