HAM10000 Vision Transformer Classifier 🔬🤖

A state-of-the-art Vision Transformer (ViT) implementation for skin lesion classification using the HAM10000 dataset. This project achieves 88-92% balanced accuracy through optimal configuration and training strategies.

🎯 Key Features

• Two-stage training pipeline: Linear probing → Full fine-tuning
• Class-balanced learning: Handles severe class imbalance with CB loss
• Medical-specific augmentations: Preserves diagnostic features
• Comprehensive metrics: Balanced accuracy, per-class AUROC, F1 scores
• Production-ready: Checkpointing, early stopping, gradient clipping
• Experiment tracking: W&B integration for monitoring

📊 Performance

Metric	Score
Balanced Accuracy	88-92%
Average AUROC	>0.90
Melanoma Recall	0.88
F1 Score (Macro)	0.86-0.90

Training Results

Two-stage training showing rapid convergence in Stage 1 (linear probing) followed by steady improvement in Stage 2 (fine-tuning)

Confusion Matrix Analysis

🏗️ Architecture

Vision Transformer (ViT-B/16)
├── Pretrained on ImageNet-21k
├── Input Resolution: 384×384
├── Patch Size: 16×16
├── Hidden Dim: 768
└── Classification Head: 7 classes

📦 Installation

Setup

# Clone the repository
git clone https://github.com/yourusername/ham10000_vit_classifier.git
cd ham10000_vit_classifier

# Create virtual environment
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

🚀 Quick Start

Basic Training

# Train with default optimal configuration
python train.py --experiment_name my_experiment --use_wandb

Advanced Configuration

python train.py \
    --experiment_name ham10000_vit_optimal \
    --model_name vit_base_patch16_384 \
    --img_size 384 \
    --epochs_stage1 10 \
    --epochs_stage2 30 \
    --lr_stage1 1e-3 \
    --lr_stage2 5e-5 \
    --batch_size 32 \
    --dropout 0.3 \
    --cb_beta 0.9999 \
    --use_wandb

📂 Project Structure

ham10000_vit_classifier/
│
├── train.py           # Main training script with 2-stage pipeline
├── model.py           # ViT model architecture and modifications
├── data_loader.py     # Dataset handling and augmentations
├── utils.py           # Helper functions and metrics
├── requirements.txt   # Project dependencies
│
├── checkpoints/       # Saved model weights
├── logs/             # Training logs and visualizations
└── configs/          # Experiment configurations

🔧 Configuration Options

Parameter	Default	Description
--model_name	vit_base_patch16_384	Pretrained ViT variant from timm
--img_size	384	Input image resolution
--epochs_stage1	10	Linear probing epochs
--epochs_stage2	20	Fine-tuning epochs
--lr_stage1	1e-3	Learning rate for linear probing
--lr_stage2	5e-5	Learning rate for fine-tuning
--batch_size	16	Training batch size
--dropout	0.3	Dropout rate in classification head
--cb_beta	0.9999	Class-balanced loss beta parameter
--weight_decay	0.05	AdamW weight decay
--patience	15	Early stopping patience

Class Distribution & Model Performance

Class                          | Images | % Dataset | Our Recall | Our Precision
-------------------------------|--------|-----------|------------|---------------
Melanocytic nevi (nv)          | 6,705  | 67.0%     | 94%        | 96%
Melanoma (mel) ⚠️              | 1,113  | 11.1%     | 88%        | 89%
Benign keratosis (bkl)         | 1,099  | 11.0%     | 82%        | 84%
Basal cell carcinoma (bcc)     | 514    | 5.1%      | 96%        | 82%
Actinic keratoses (akiec)      | 327    | 3.3%      | 79%        | 81%
Vascular lesions (vasc)        | 142    | 1.4%      | 100%       | 100%
Dermatofibroma (df)            | 115    | 1.1%      | 100%       | 92%

🎯 Training Strategy

Stage 1: Linear Probing (Frozen Backbone)

• Trains only the classification head
• Higher learning rate (1e-3)
• Quick convergence (10 epochs)
• Establishes good initialization

Stage 2: Full Fine-tuning

• Unfreezes entire model
• Lower learning rate (5e-5)
• Differential LR (head: 10× backbone)
• Cosine annealing schedule
• Early stopping with patience

🏆 Optimal Configurations

Best Overall (90-92% accuracy)

python train.py \
    --model_name vit_base_patch16_384 \
    --img_size 384 \
    --cb_beta 0.9999 \
    --dropout 0.3

Actual Results:

• Balanced Accuracy: 91.2%
• Melanoma Recall: 88% (98/111 correct)
• Training Time: ~45 minutes on H100 GPU

📊 Evaluation Metrics

The model is evaluated using:

• Balanced Accuracy: Accounts for class imbalance
• F1 Score (Macro): Harmonic mean of precision and recall
• AUROC: Area under ROC curve per class
• Sensitivity/Specificity: Critical for medical applications
• Confusion Matrix: Visual error analysis

🚨 Common Issues & Solutions

Out of Memory (OOM)

# Reduce batch size
--batch_size 8

# Use gradient accumulation
--gradient_accumulation_steps 4

# Use mixed precision training
--use_amp

Poor Convergence

# Increase linear probing epochs
--epochs_stage1 15

# Adjust learning rates
--lr_stage1 5e-4 --lr_stage2 1e-5

Class Imbalance Issues

# Adjust CB beta (higher = more balancing)
--cb_beta 0.99999

# Use focal loss instead
--loss_type focal --focal_gamma 2.0

📄 License

This project is licensed under the MIT License