Computer Vision

Vision-Transformer

Implementation of Vision Transformer

Jupyter Notebook Created Oct 2024 Updated May 2025

Vision-Transformer

Implementation of Vision Transformer

🔮 Vision Transformer (ViT)

PyTorch Python License

📌 Overview

A PyTorch implementation of the Vision Transformer (ViT) model introduced in the paper "An Image is Worth 16x16 Words". This implementation focuses on clarity, modularity, and ease of use while maintaining high performance.

✨ Key Features

  • 🚀 High-Performance Implementation: Optimized PyTorch code with CUDA support
  • 🎯 Modular Architecture: Easy to modify and extend components
  • 📊 Complete Training Pipeline: Including data loading, training, and evaluation
  • 🔧 Configurable Design: Easily adjustable hyperparameters
  • 📝 Comprehensive Documentation: Well-documented code and usage examples

🛠️ Installation

git clone https://github.com/yourusername/vision-transformer.git
cd vision-transformer
pip install -r requirements.txt

📦 Requirements

  • Python 3.7+
  • PyTorch 1.12.1+
  • torchvision
  • numpy
  • tqdm

🚀 Quick Start

# Initialize the Vision Transformer
model = ViT(
    img_size=224,
    patch_size=16,
    num_classes=3,
    embedding_dim=768,
    num_heads=12,
    num_layers=12,
    mlp_dim=3072,
    dropout=0.1
)

# Train the model
results = train(
    model=model,
    train_dataloader=train_dataloader,
    test_dataloader=test_dataloader,
    epochs=20,
    device=device
)

🏗️ Model Architecture

The Vision Transformer consists of several key components:

1. Patch Embedding

  • Splits input images into fixed-size patches (16x16)
  • Projects patches to embedding dimension (768)
  • Adds positional embeddings

2. Transformer Encoder

  • Multi-head self-attention layers (12 heads)
  • Layer normalization and residual connections
  • MLP blocks with GELU activation
  • Dropout for regularization

3. Classification Head

  • Layer normalization
  • Linear projection to class logits

📈 Training Pipeline

Our implementation includes a robust training pipeline with:

  1. Data Processing

    • Image resizing and normalization
    • Patch creation
    • Batch processing
  2. Training Loop

    • Forward/backward passes
    • Gradient updates
    • Loss calculation
    • Metric tracking
  3. Evaluation

    • Accuracy metrics
    • Loss monitoring
    • Performance validation

🔧 Configuration

config = {
    "model": {
        "img_size": 224,
        "patch_size": 16,
        "num_classes": 3,
        "embedding_dim": 768,
        "num_heads": 12,
        "num_layers": 12,
        "mlp_dim": 3072,
        "dropout": 0.1
    },
    "training": {
        "batch_size": 32,
        "learning_rate": 3e-3,
        "weight_decay": 0.3,
        "epochs": 20
    }
}

📊 Usage Examples

Training on Custom Dataset

# Prepare your data
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

# Create dataloaders
train_data = datasets.ImageFolder("path/to/train", transform=transform)
train_loader = DataLoader(
    train_data,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

# Train model
model = ViT(num_classes=len(train_data.classes))
train(model, train_loader, test_loader, epochs=20)

🤝 Contributing

We welcome contributions! Here's how you can help:

  1. Fork the repository
  2. Create a feature branch (git checkout -b feature/amazing-feature)
  3. Commit your changes (git commit -m 'Add amazing feature')
  4. Push to the branch (git push origin feature/amazing-feature)
  5. Open a Pull Request

📚 Citation

If you use this implementation in your research, please cite:

@article{dosovitskiy2020image,
  title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},
  author={Dosovitskiy, Alexey and others},
  journal={arXiv preprint arXiv:2010.11929},
  year={2020}
}

📄 License

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

📧 Contact

Your Name - @yourtwitter Project Link: https://github.com/yourusername/vision-transformer


Created with 💙 by Ashmit Gupta