diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index cb61c3db..7a41aef1 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -1,16 +1,16 @@ ## Description - + ## Related Issue - + ## Type of Change -- [ ] 📚 Examples / docs / tutorials / dependencies update +- [ ] 📚 Documentation / examples / dependencies update - [ ] 🔧 Bug fix (non-breaking change which fixes an issue) - [ ] 🥂 Improvement (non-breaking change which improves an existing feature) - [ ] 🚀 New feature (non-breaking change which adds functionality) @@ -21,8 +21,9 @@ -- [ ] I've read the [`CODE_OF_CONDUCT.md`](https://github.com/lighter/lighter/blob/master/CODE_OF_CONDUCT.md) document. -- [ ] I've read the [`CONTRIBUTING.md`](https://github.com/lighter/lighter/blob/master/CONTRIBUTING.md) guide. -- [ ] I've updated the code style using `make codestyle`. -- [ ] I've written tests for all new methods and classes that I created. -- [ ] I've written the docstring in Google format for all the methods and classes that I used. +- [ ] I've read the [`CONTRIBUTING.md`](https://github.com/project-lighter/lighter/blob/main/CONTRIBUTING.md) guide +- [ ] Code passes linting: `just lint` +- [ ] Code passes type checking: `just types` +- [ ] All tests pass with coverage: `just coverage` +- [ ] I've added tests for new functionality +- [ ] I've updated documentation if needed diff --git a/.gitignore b/.gitignore index 3755fab6..bd0303e6 100644 --- a/.gitignore +++ b/.gitignore @@ -145,29 +145,21 @@ tensorboard/ # ---- Our ignores ---- **/.DS_Store -**/predictions/ -test_dir/ -checkpoints/ -prototyping* -.aider* *.code-workspace .scale_batch_size* **/*.txt .ruff_cache -cifar10/ - -# Paper -paper/jats/ - # Coding agents CLAUDE.md GEMINI.md AGENTS.md - **/.claude/ +# Outputs, logs, dataset files **/lightning_logs/ **/outputs/ **/.datasets/ **/*.zip +**/predictions/ +checkpoints/ diff --git a/docs/examples/image-classification.md b/docs/examples/image-classification.md deleted file mode 100644 index fc0cfe5f..00000000 --- a/docs/examples/image-classification.md +++ /dev/null @@ -1,685 +0,0 @@ ---- -title: Image Classification Example ---- - -# Complete Image Classification Example - -Train a CIFAR-10 classifier from scratch with all the features. - -This example shows a complete, production-ready setup including: - -- Both LightningModule and LighterModule approaches -- Data augmentation -- Learning rate scheduling -- Multiple metrics -- Checkpointing -- Early stopping -- TensorBoard logging -- Multi-GPU support - -## Project Setup - -### Directory Structure - -``` -cifar10/ -├── __lighter__.py # Marker file (enables project.* imports) -├── __init__.py -├── models.py # Model definitions -├── data.py # Data utilities (optional) -├── configs/ -│ ├── resnet18.yaml # ResNet-18 config -│ ├── resnet50.yaml # ResNet-50 config -│ └── efficientnet.yaml # EfficientNet config -└── outputs/ # Generated by Lighter -``` - -### Installation - -```bash -pip install lighter torch torchvision pytorch-lightning torchmetrics -``` - -## Approach 1: Using LightningModule - -### Step 1: Create the Module - -`models.py`: - -```python -import pytorch_lightning as pl -import torch -import torch.nn.functional as F -import torchmetrics - - -class CIFAR10Classifier(pl.LightningModule): - """CIFAR-10 image classifier with custom training logic.""" - - def __init__( - self, - network, - learning_rate=0.001, - weight_decay=0.0001, - max_epochs=100, - ): - super().__init__() - self.save_hyperparameters(ignore=['network']) - self.network = network - - # Metrics - self.train_acc = torchmetrics.Accuracy(task='multiclass', num_classes=10) - self.val_acc = torchmetrics.Accuracy(task='multiclass', num_classes=10) - self.val_f1 = torchmetrics.F1Score(task='multiclass', num_classes=10) - - def forward(self, x): - return self.network(x) - - def training_step(self, batch, batch_idx): - x, y = batch - logits = self(x) - loss = F.cross_entropy(logits, y) - - # Update metrics - self.train_acc(logits, y) - - # Log - self.log('train/loss', loss, on_step=True, on_epoch=True) - self.log('train/acc', self.train_acc, on_step=False, on_epoch=True) - - return loss - - def validation_step(self, batch, batch_idx): - x, y = batch - logits = self(x) - loss = F.cross_entropy(logits, y) - - # Update metrics - self.val_acc(logits, y) - self.val_f1(logits, y) - - # Log - self.log('val/loss', loss) - self.log('val/acc', self.val_acc, on_step=False, on_epoch=True) - self.log('val/f1', self.val_f1, on_step=False, on_epoch=True) - - def configure_optimizers(self): - optimizer = torch.optim.AdamW( - self.parameters(), - lr=self.hparams.learning_rate, - weight_decay=self.hparams.weight_decay, - ) - - scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( - optimizer, - T_max=self.hparams.max_epochs, - eta_min=1e-6, - ) - - return { - 'optimizer': optimizer, - 'lr_scheduler': { - 'scheduler': scheduler, - 'interval': 'epoch', - } - } -``` - -### Step 2: Create the Config - -`configs/resnet18.yaml`: - -```yaml -# CIFAR-10 Classification with ResNet-18 - -vars: - num_classes: 10 - batch_size: 128 - num_workers: 4 - base_lr: 0.001 - max_epochs: 100 - -trainer: - _target_: pytorch_lightning.Trainer - max_epochs: "%vars::max_epochs" - accelerator: auto - devices: 1 - - callbacks: - # Save best models - - _target_: pytorch_lightning.callbacks.ModelCheckpoint - monitor: val/acc - mode: max - save_top_k: 3 - filename: 'best-acc-{epoch:02d}-{val/acc:.4f}' - - # Early stopping - - _target_: pytorch_lightning.callbacks.EarlyStopping - monitor: val/loss - patience: 15 - mode: min - verbose: true - - # Log learning rate - - _target_: pytorch_lightning.callbacks.LearningRateMonitor - logging_interval: epoch - - logger: - _target_: pytorch_lightning.loggers.TensorBoardLogger - save_dir: logs - name: cifar10_resnet18 - -model: - _target_: project.models.CIFAR10Classifier - learning_rate: "%vars::base_lr" - weight_decay: 0.0001 - max_epochs: "%vars::max_epochs" - - network: - _target_: torchvision.models.resnet18 - num_classes: "%vars::num_classes" - weights: null # Train from scratch - -data: - _target_: lighter.LighterDataModule - - train_dataloader: - _target_: torch.utils.data.DataLoader - batch_size: "%vars::batch_size" - shuffle: true - num_workers: "%vars::num_workers" - pin_memory: true - persistent_workers: true - - dataset: - _target_: torchvision.datasets.CIFAR10 - root: ./data - train: true - download: true - - transform: - _target_: torchvision.transforms.Compose - transforms: - # Augmentation - - _target_: torchvision.transforms.RandomCrop - size: 32 - padding: 4 - - _target_: torchvision.transforms.RandomHorizontalFlip - p: 0.5 - - # Normalization - - _target_: torchvision.transforms.ToTensor - - _target_: torchvision.transforms.Normalize - mean: [0.4914, 0.4822, 0.4465] - std: [0.2470, 0.2435, 0.2616] - - val_dataloader: - _target_: torch.utils.data.DataLoader - batch_size: "%vars::batch_size" - num_workers: "%vars::num_workers" - pin_memory: true - - dataset: - _target_: torchvision.datasets.CIFAR10 - root: ./data - train: false - download: true - - transform: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ToTensor - - _target_: torchvision.transforms.Normalize - mean: [0.4914, 0.4822, 0.4465] - std: [0.2470, 0.2435, 0.2616] -``` - -### Step 3: Run Training - -```bash -cd cifar10 -lighter fit configs/resnet18.yaml -``` - -**Expected output:** - -``` -Epoch 0: 100%|████████| 391/391 [00:45<00:00, 8.58it/s, loss=1.82, train/acc=0.335, v_num=0] -Validation: 100%|████████| 79/79 [00:05<00:00, 14.23it/s] -Epoch 1: 100%|████████| 391/391 [00:44<00:00, 8.76it/s, loss=1.45, train/acc=0.472, v_num=0] -... -Epoch 99: 100%|████████| 391/391 [00:43<00:00, 8.98it/s, loss=0.23, train/acc=0.921, v_num=0] -``` - -**Expected accuracy:** ~90-92% on validation set after 100 epochs. - -## Approach 2: Using LighterModule - -### Step 1: Create the Module - -`models.py`: - -```python -from lighter import LighterModule - - -class CIFAR10Model(LighterModule): - """CIFAR-10 classifier using LighterModule.""" - - def training_step(self, batch, batch_idx): - x, y = batch - logits = self(x) - loss = self.criterion(logits, y) - - # Update metrics - if self.train_metrics: - self.train_metrics(logits, y) - - return {'loss': loss} - - def validation_step(self, batch, batch_idx): - x, y = batch - logits = self(x) - loss = self.criterion(logits, y) - - # Update metrics - if self.val_metrics: - self.val_metrics(logits, y) - - return {'loss': loss} -``` - -### Step 2: Create the Config - -`configs/resnet18.yaml`: - -```yaml -# CIFAR-10 with LighterModule - -vars: - num_classes: 10 - batch_size: 128 - num_workers: 4 - base_lr: 0.001 - max_epochs: 100 - -trainer: - _target_: pytorch_lightning.Trainer - max_epochs: "%vars::max_epochs" - accelerator: auto - devices: 1 - - callbacks: - - _target_: pytorch_lightning.callbacks.ModelCheckpoint - monitor: val/Accuracy - mode: max - save_top_k: 3 - filename: 'best-acc-{epoch:02d}-{val/Accuracy:.4f}' - - - _target_: pytorch_lightning.callbacks.EarlyStopping - monitor: val/loss - patience: 15 - mode: min - - - _target_: pytorch_lightning.callbacks.LearningRateMonitor - logging_interval: epoch - - logger: - _target_: pytorch_lightning.loggers.TensorBoardLogger - save_dir: logs - name: cifar10_resnet18 - -model: - _target_: project.models.CIFAR10Model - - network: - _target_: torchvision.models.resnet18 - num_classes: "%vars::num_classes" - weights: null - - criterion: - _target_: torch.nn.CrossEntropyLoss - label_smoothing: 0.1 # Regularization - - optimizer: - _target_: torch.optim.AdamW - params: "$@model::network.parameters()" - lr: "%vars::base_lr" - weight_decay: 0.0001 - - scheduler: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - optimizer: "@model::optimizer" - T_max: "%vars::max_epochs" - eta_min: 0.000001 - - train_metrics: - - _target_: torchmetrics.Accuracy - task: multiclass - num_classes: "%vars::num_classes" - - _target_: torchmetrics.F1Score - task: multiclass - num_classes: "%vars::num_classes" - average: macro - - val_metrics: "%model::train_metrics" - -data: - # Same as above... - _target_: lighter.LighterDataModule - train_dataloader: - _target_: torch.utils.data.DataLoader - batch_size: "%vars::batch_size" - shuffle: true - num_workers: "%vars::num_workers" - pin_memory: true - persistent_workers: true - dataset: - _target_: torchvision.datasets.CIFAR10 - root: ./data - train: true - download: true - transform: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.RandomCrop - size: 32 - padding: 4 - - _target_: torchvision.transforms.RandomHorizontalFlip - - _target_: torchvision.transforms.ToTensor - - _target_: torchvision.transforms.Normalize - mean: [0.4914, 0.4822, 0.4465] - std: [0.2470, 0.2435, 0.2616] - - val_dataloader: - _target_: torch.utils.data.DataLoader - batch_size: "%vars::batch_size" - num_workers: "%vars::num_workers" - pin_memory: true - dataset: - _target_: torchvision.datasets.CIFAR10 - root: ./data - train: false - transform: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ToTensor - - _target_: torchvision.transforms.Normalize - mean: [0.4914, 0.4822, 0.4465] - std: [0.2470, 0.2435, 0.2616] -``` - -### Step 3: Run Training - -```bash -lighter fit configs/resnet18.yaml -``` - -Same results, less code! - -## Experiments - -### Try Different Architectures - -Override the network directly from CLI: - -```bash -# ResNet-50 -lighter fit configs/resnet18.yaml \ - model::network::_target_=torchvision.models.resnet50 \ - vars::batch_size=64 \ - vars::base_lr=0.0005 - -# EfficientNet -lighter fit configs/resnet18.yaml \ - model::network::_target_=torchvision.models.efficientnet_b0 \ - vars::batch_size=64 -``` - -### Hyperparameter Tuning - -Try different learning rates: - -```bash -lighter fit configs/resnet18.yaml vars::base_lr=0.0001 -lighter fit configs/resnet18.yaml vars::base_lr=0.001 -lighter fit configs/resnet18.yaml vars::base_lr=0.01 -``` - -Try different batch sizes: - -```bash -lighter fit configs/resnet18.yaml vars::batch_size=64 -lighter fit configs/resnet18.yaml vars::batch_size=256 -``` - -## Multi-GPU Training - -Use all GPUs: - -```bash -lighter fit configs/resnet18.yaml \ - trainer::devices=-1 \ - trainer::strategy=ddp -``` - -Or specific number: - -```bash -lighter fit configs/resnet18.yaml \ - trainer::devices=4 \ - trainer::strategy=ddp -``` - -**Note:** With DDP, effective batch size = `batch_size × num_gpus`. - -## Mixed Precision - -Train faster with 16-bit precision: - -```bash -lighter fit configs/resnet18.yaml trainer::precision=16 -``` - -Or BFloat16 (on A100/H100): - -```bash -lighter fit configs/resnet18.yaml trainer::precision="bf16-mixed" -``` - -## Transfer Learning - -Use pretrained ImageNet weights: - -`configs/pretrained.yaml`: - -```yaml -model: - network: - weights: IMAGENET1K_V2 # Pretrained weights - -# Lower LR for finetuning -vars: - base_lr: 0.0001 -``` - -Run by composing with base config: - -```bash -lighter fit configs/resnet18.yaml configs/pretrained.yaml -``` - -**Expected:** ~94-95% accuracy (better than training from scratch). - -## Monitoring Training - -### TensorBoard - -View logs: - -```bash -tensorboard --logdir logs -``` - -Open browser to `http://localhost:6006`. - -You'll see: -- Train/val loss curves -- Accuracy curves -- Learning rate schedule -- Model graph - -### Weights & Biases - -Use W&B for experiment tracking: - -```yaml -trainer: - logger: - _target_: pytorch_lightning.loggers.WandbLogger - project: cifar10 - name: resnet18_experiment -``` - -Run: - -```bash -wandb login -lighter fit configs/resnet18.yaml -``` - -## Testing - -After training, test on the test set: - -Add test dataloader to config: - -```yaml -data: - test_dataloader: - _target_: torch.utils.data.DataLoader - batch_size: 128 - num_workers: 4 - dataset: - _target_: torchvision.datasets.CIFAR10 - root: ./data - train: false - transform: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ToTensor - - _target_: torchvision.transforms.Normalize - mean: [0.4914, 0.4822, 0.4465] - std: [0.2470, 0.2435, 0.2616] -``` - -Add test_metrics to model: - -```yaml -model: - test_metrics: "%model::train_metrics" -``` - -Run test: - -```bash -lighter test configs/resnet18.yaml \ - args::test::ckpt_path=outputs/.../checkpoints/best-acc-*.ckpt -``` - -## Complete Working Example - -The complete code is available in the Lighter repository: - -[View on GitHub →](https://github.com/project-lighter/lighter/tree/main/projects/cifar10) - -## Common Issues - -### Out of Memory - -**Problem:** CUDA out of memory error - -**Solutions:** - -1. Reduce batch size: - ```bash - lighter fit config.yaml vars::batch_size=64 - ``` - -2. Use gradient accumulation: - ```yaml - trainer: - accumulate_grad_batches: 2 - ``` - -3. Use mixed precision: - ```yaml - trainer: - precision: 16 - ``` - -### Slow Training - -**Problem:** Training too slow - -**Solutions:** - -1. Increase num_workers: - ```yaml - data: - train_dataloader: - num_workers: 8 - ``` - -2. Enable pin_memory and persistent_workers: - ```yaml - data: - train_dataloader: - pin_memory: true - persistent_workers: true - ``` - -3. Use mixed precision: - ```yaml - trainer: - precision: 16 - ``` - -### Low Accuracy - -**Problem:** Model not learning well - -**Solutions:** - -1. Check data normalization matches pretrained weights -2. Try different learning rate: - ```bash - lighter fit config.yaml vars::base_lr=0.01 - ``` -3. Increase max_epochs -4. Add more augmentation -5. Use pretrained weights - -## Next Steps - -- [Multi-GPU Example](multi-gpu.md) - Distributed training setup -- [Training Guide](../guides/training.md) - More training tips -- [Best Practices](../guides/best-practices.md) - Production patterns - -## Summary - -This example showed: - -- ✅ Complete CIFAR-10 classifier setup -- ✅ Both LightningModule and LighterModule approaches -- ✅ Data augmentation and normalization -- ✅ Learning rate scheduling -- ✅ Multiple metrics -- ✅ Checkpointing and early stopping -- ✅ TensorBoard logging -- ✅ Hyperparameter tuning from CLI -- ✅ Multi-GPU support -- ✅ Mixed precision training -- ✅ Transfer learning - -**Key takeaway:** One config file controls everything. Iterate fast without code changes. diff --git a/docs/examples/index.md b/docs/examples/index.md new file mode 100644 index 00000000..c560aa6f --- /dev/null +++ b/docs/examples/index.md @@ -0,0 +1,107 @@ +--- +title: Example Projects +--- + +# Example Projects + +Lighter includes example projects demonstrating real-world applications across various domains. Each project is self-contained with its own config files and documentation. + +## Available Projects + +| Project | Domain | Description | Extra Dependencies | +|---------|--------|-------------|-------------------| +| [cifar10](https://github.com/project-lighter/lighter/tree/main/projects/cifar10) | Image Classification | Basic image classification with ResNet | - | +| [eeg](https://github.com/project-lighter/lighter/tree/main/projects/eeg) | EEG Analysis | Brain signal classification | braindecode, eegdash, mne | +| [huggingface_llm](https://github.com/project-lighter/lighter/tree/main/projects/huggingface_llm) | Text Classification | Sentiment analysis with HuggingFace Transformers | transformers, datasets | +| [lora](https://github.com/project-lighter/lighter/tree/main/projects/lora) | Parameter-Efficient Fine-Tuning | LoRA fine-tuning of language models | peft | +| [medical_segmentation](https://github.com/project-lighter/lighter/tree/main/projects/medical_segmentation) | Medical Imaging | 3D medical image segmentation | monai, itk | +| [self_supervised](https://github.com/project-lighter/lighter/tree/main/projects/self_supervised) | Self-Supervised Learning | SimCLR contrastive learning | lightly | +| [video_recognition](https://github.com/project-lighter/lighter/tree/main/projects/video_recognition) | Video Understanding | Video classification with SlowFast | pytorchvideo, av | +| [vision_language](https://github.com/project-lighter/lighter/tree/main/projects/vision_language) | Vision-Language | CLIP-style image-text contrastive learning | transformers | + +## Running an Example + +Each project follows the same structure: + +``` +projects// +├── __lighter__.py # Project marker (enables project.* imports) +├── __init__.py +├── *.py # Custom modules +├── configs/ +│ └── *.yaml # Experiment configs +└── README.md # Project-specific documentation +``` + +To run any example: + +```bash +# Clone the repo +git clone https://github.com/project-lighter/lighter.git +cd lighter + +# Install Lighter +pip install -e . + +# Navigate to a project +cd projects/cifar10 + +# Install extra dependencies if needed (check the README) +pip install + +# Run training +lighter fit configs/example.yaml +``` + +## Project Highlights + +### cifar10 + +The simplest starting point. Demonstrates: + +- Basic `LighterModule` usage +- Data augmentation in config +- Standard training workflow + +### medical_segmentation + +Shows how to use MONAI with Lighter for 3D medical imaging: + +- 3D UNet architecture +- Medical imaging transforms +- Dice loss and metrics + +### self_supervised + +Contrastive learning with SimCLR: + +- Custom projection heads +- Multi-view data augmentation +- NT-Xent loss + +### huggingface_llm + +Integrates HuggingFace Transformers: + +- Tokenizer configuration +- Pre-trained model loading +- Text classification + +### lora + +Parameter-efficient fine-tuning: + +- LoRA adapters via PEFT +- Freezing base model layers +- Memory-efficient training + +## Creating Your Own Project + +Use any example as a template: + +1. Copy the project directory +2. Add a `__lighter__.py` marker file +3. Modify configs to point to your data and model +4. Run with `lighter fit configs/your_config.yaml` + +See the [Custom Code Guide](../guides/custom-code.md) for details on project structure. diff --git a/docs/examples/multi-gpu.md b/docs/examples/multi-gpu.md deleted file mode 100644 index 1e45f18d..00000000 --- a/docs/examples/multi-gpu.md +++ /dev/null @@ -1,704 +0,0 @@ ---- -title: Multi-GPU Training ---- - -# Multi-GPU Training - -Scale your training across multiple GPUs with Distributed Data Parallel (DDP). - -This guide shows how to train on multiple GPUs using PyTorch Lightning's DDP strategy through Lighter configs. - -## Quick Start - -Train on all available GPUs: - -```bash -lighter fit config.yaml trainer::devices=-1 trainer::strategy=ddp -``` - -Train on specific number of GPUs: - -```bash -lighter fit config.yaml trainer::devices=4 trainer::strategy=ddp -``` - -That's it! Your code works unchanged. - -## Configuration - -### In Config File - -```yaml -trainer: - _target_: pytorch_lightning.Trainer - devices: 4 # Use 4 GPUs - strategy: ddp # Distributed Data Parallel - accelerator: auto # Automatically use CUDA if available -``` - -### From CLI - -Override devices: - -```bash -# All GPUs -lighter fit config.yaml trainer::devices=-1 - -# Specific GPUs (0, 1, 2, 3) -lighter fit config.yaml trainer::devices=4 - -# Specific GPU IDs -lighter fit config.yaml 'trainer::devices=[0,2,3]' -``` - -## DDP Strategy - -### Basic DDP - -Recommended for most cases: - -```yaml -trainer: - strategy: ddp -``` - -Features: -- Each GPU gets own process -- Gradients synchronized across GPUs -- Model replicated on each GPU -- Data split across GPUs - -### DDP Spawn - -Alternative that spawns subprocesses: - -```yaml -trainer: - strategy: ddp_spawn -``` - -Use when: -- DDP doesn't work on your system -- Debugging (easier to see errors) - -**Note:** Slightly slower than DDP. - -### DDP Find Unused Parameters - -If you get "unused parameters" error: - -```yaml -trainer: - strategy: - _target_: pytorch_lightning.strategies.DDPStrategy - find_unused_parameters: true -``` - -## Batch Size Considerations - -### Per-GPU Batch Size - -Each GPU processes `batch_size` samples: - -```yaml -data: - train_dataloader: - batch_size: 32 # Each GPU: 32 samples -``` - -**Effective batch size** with 4 GPUs = 32 × 4 = 128 - -### Keep Total Batch Size - -To keep same total batch size across different GPU counts: - -```yaml -vars: - num_gpus: 4 - total_batch_size: 128 - -data: - train_dataloader: - batch_size: "$%vars::total_batch_size // %vars::num_gpus" -``` - -Override for different GPU counts: - -```bash -# 1 GPU: batch_size = 128 -lighter fit config.yaml vars::num_gpus=1 - -# 4 GPUs: batch_size = 32 per GPU -lighter fit config.yaml vars::num_gpus=4 - -# 8 GPUs: batch_size = 16 per GPU -lighter fit config.yaml vars::num_gpus=8 -``` - -## Learning Rate Scaling - -### Linear Scaling Rule - -When increasing batch size, scale LR proportionally: - -```yaml -vars: - num_gpus: 4 - base_lr: 0.001 - -model: - optimizer: - lr: "$%vars::base_lr * %vars::num_gpus" -``` - -**Example:** -- 1 GPU: LR = 0.001, batch = 32 -- 4 GPUs: LR = 0.004, batch = 128 (32×4) - -### Square Root Scaling - -Alternative for very large batch sizes: - -```yaml -model: - optimizer: - lr: "$%vars::base_lr * (%vars::num_gpus ** 0.5)" -``` - -## Complete Multi-GPU Example - -`experiments/multi_gpu.yaml`: - -```yaml -vars: - # Hardware - num_gpus: 4 - - # Dataset - num_classes: 10 - - # Hyperparameters - base_lr: 0.001 - total_batch_size: 512 # Total across all GPUs - max_epochs: 100 - - # Computed - per_gpu_batch_size: "$%vars::total_batch_size // %vars::num_gpus" - scaled_lr: "$%vars::base_lr * %vars::num_gpus" - -trainer: - _target_: pytorch_lightning.Trainer - max_epochs: "%vars::max_epochs" - devices: "%vars::num_gpus" - strategy: ddp - accelerator: auto - - # Recommended settings - sync_batchnorm: true # Sync batch normalization - precision: 16 # Mixed precision - - callbacks: - - _target_: pytorch_lightning.callbacks.ModelCheckpoint - monitor: val/acc - mode: max - save_top_k: 3 - - - _target_: pytorch_lightning.callbacks.LearningRateMonitor - logging_interval: epoch - - logger: - _target_: pytorch_lightning.loggers.TensorBoardLogger - save_dir: logs - name: multi_gpu_experiment - -model: - _target_: lighter.LighterModule - - network: - _target_: torchvision.models.resnet50 - num_classes: "%vars::num_classes" - - criterion: - _target_: torch.nn.CrossEntropyLoss - - optimizer: - _target_: torch.optim.AdamW - params: "$@model::network.parameters()" - lr: "%vars::scaled_lr" - weight_decay: 0.0001 - - scheduler: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - optimizer: "@model::optimizer" - T_max: "%vars::max_epochs" - - train_metrics: - - _target_: torchmetrics.Accuracy - task: multiclass - num_classes: "%vars::num_classes" - - val_metrics: "%model::train_metrics" - -data: - _target_: lighter.LighterDataModule - - train_dataloader: - _target_: torch.utils.data.DataLoader - batch_size: "%vars::per_gpu_batch_size" - shuffle: true - num_workers: 8 # Increase for multi-GPU - pin_memory: true - persistent_workers: true - dataset: - _target_: torchvision.datasets.CIFAR10 - root: ./data - train: true - download: true - transform: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.RandomCrop - size: 32 - padding: 4 - - _target_: torchvision.transforms.RandomHorizontalFlip - - _target_: torchvision.transforms.ToTensor - - _target_: torchvision.transforms.Normalize - mean: [0.4914, 0.4822, 0.4465] - std: [0.2470, 0.2435, 0.2616] - - val_dataloader: - _target_: torch.utils.data.DataLoader - batch_size: "%vars::per_gpu_batch_size" - num_workers: 8 - pin_memory: true - dataset: - _target_: torchvision.datasets.CIFAR10 - root: ./data - train: false - transform: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ToTensor - - _target_: torchvision.transforms.Normalize - mean: [0.4914, 0.4822, 0.4465] - std: [0.2470, 0.2435, 0.2616] -``` - -Run: - -```bash -lighter fit experiments/multi_gpu.yaml -``` - -## Advanced Strategies - -### FSDP (Fully Sharded Data Parallel) - -For very large models that don't fit on single GPU: - -```yaml -trainer: - strategy: fsdp - devices: 4 -``` - -FSDP shards: -- Model parameters -- Gradients -- Optimizer states - -Across GPUs, saving memory. - -### DeepSpeed - -For even larger models: - -```yaml -trainer: - strategy: - _target_: pytorch_lightning.strategies.DeepSpeedStrategy - stage: 2 # ZeRO Stage 2 - devices: 4 - precision: 16 -``` - -Stages: -- **Stage 1**: Shard optimizer states -- **Stage 2**: Shard gradients + Stage 1 -- **Stage 3**: Shard parameters + Stage 2 - -### DDP with Static Graph - -For maximum performance (PyTorch 1.11+): - -```yaml -trainer: - strategy: - _target_: pytorch_lightning.strategies.DDPStrategy - static_graph: true -``` - -**Requirements:** -- Model structure doesn't change between steps -- No dynamic control flow - -**Benefit:** ~10% speedup. - -## Data Loading Optimization - -### Increase num_workers - -More workers for multi-GPU: - -```yaml -data: - train_dataloader: - num_workers: "$%vars::num_gpus * 4" # 4 workers per GPU -``` - -### Use Persistent Workers - -Avoid worker respawning: - -```yaml -data: - train_dataloader: - persistent_workers: true -``` - -### Pin Memory - -Faster GPU transfer: - -```yaml -data: - train_dataloader: - pin_memory: true -``` - -## Gradient Accumulation - -Simulate even larger batch sizes: - -```yaml -trainer: - accumulate_grad_batches: 4 -``` - -**Effective batch size** = `batch_size × num_gpus × accumulate_grad_batches` - -Example with 4 GPUs: -- `batch_size = 32` -- `num_gpus = 4` -- `accumulate_grad_batches = 4` -- **Effective = 32 × 4 × 4 = 512** - -## Sync Batch Normalization - -Important for small per-GPU batch sizes: - -```yaml -trainer: - sync_batchnorm: true -``` - -Synchronizes batch norm statistics across GPUs. - -**Use when:** Per-GPU batch size < 8. - -## Mixed Precision - -Combine with multi-GPU for maximum speed: - -```yaml -trainer: - precision: 16 # or "bf16-mixed" - devices: 4 - strategy: ddp -``` - -**Speedup:** ~2-3× faster than FP32. - -## Monitoring Multi-GPU Training - -### TensorBoard - -Same as single GPU: - -```bash -tensorboard --logdir logs -``` - -Metrics automatically aggregated across GPUs. - -### Weights & Biases - -Works out of the box: - -```yaml -trainer: - logger: - _target_: pytorch_lightning.loggers.WandbLogger - project: my_project -``` - -Only rank 0 process logs to avoid duplicates. - -## Checkpointing - -Same as single GPU: - -```yaml -trainer: - callbacks: - - _target_: pytorch_lightning.callbacks.ModelCheckpoint - save_top_k: 3 -``` - -Only rank 0 saves checkpoints automatically. - -## Testing Locally - -Test DDP on single machine with multiple physical GPUs: - -```bash -# Run DDP on 2 physical GPUs -lighter fit config.yaml trainer::devices=2 trainer::strategy=ddp -``` - -**Important:** `devices=k` selects `k` physical GPUs per node (equivalent to `list(range(k))`). No GPU virtualization or simulation is performed. You must have at least as many physical GPUs as specified (e.g., at least 2 physical GPUs for the example above). - -## Common Issues - -### Out of Memory - -**Solutions:** - -1. Reduce per-GPU batch size: - ```bash - lighter fit config.yaml data::train_dataloader::batch_size=16 - ``` - -2. Use gradient accumulation: - ```yaml - trainer: - accumulate_grad_batches: 2 - ``` - -3. Use FSDP or DeepSpeed for large models - -### Slow Startup - -**Problem:** Long startup time with DDP - -**Cause:** Dataset download or preprocessing on each rank - -**Solution:** Download data before training: - -```bash -# Download once -python -c "from torchvision.datasets import CIFAR10; CIFAR10('./data', download=True)" - -# Then train -lighter fit config.yaml -``` - -### Hanging at Initialization - -**Problem:** Process hangs at "Initializing distributed" - -**Solutions:** - -1. Check firewall settings -2. Try different DDP backend: - ```yaml - trainer: - strategy: - _target_: pytorch_lightning.strategies.DDPStrategy - process_group_backend: gloo # Instead of nccl - ``` - -### Different Results Across GPUs - -**Problem:** Metrics differ between runs - -**Cause:** Random seed not set or data shuffling - -**Solution:** - -```python -# In __lighter__.py -import pytorch_lightning as pl -pl.seed_everything(42, workers=True) -``` - -```yaml -data: - train_dataloader: - shuffle: true # Ensure shuffling -``` - -### Unused Parameters Error - -**Problem:** "RuntimeError: Expected to have finished reduction in the prior iteration" - -**Solution:** - -```yaml -trainer: - strategy: - _target_: pytorch_lightning.strategies.DDPStrategy - find_unused_parameters: true -``` - -## Performance Tips - -### 1. Use All CPU Cores - -```yaml -data: - train_dataloader: - num_workers: "$%vars::num_gpus * 4" -``` - -### 2. Prefetch Data - -```yaml -data: - train_dataloader: - prefetch_factor: 2 -``` - -### 3. Mixed Precision - -```yaml -trainer: - precision: 16 -``` - -### 4. Compile Model (PyTorch 2.0+) - -In your module: - -```python -def __init__(self, network, ...): - super().__init__() - self.network = torch.compile(network) -``` - -### 5. Optimize Data Loading - -- Cache dataset if it fits in RAM -- Preprocess data offline -- Use fast storage (SSD > HDD) - -## Scaling Example - -Compare different GPU counts: - -```bash -# 1 GPU baseline -lighter fit config.yaml vars::num_gpus=1 - -# 2 GPUs (~1.8× speedup) -lighter fit config.yaml vars::num_gpus=2 - -# 4 GPUs (~3.5× speedup) -lighter fit config.yaml vars::num_gpus=4 - -# 8 GPUs (~6.5× speedup) -lighter fit config.yaml vars::num_gpus=8 -``` - -**Expected scaling:** ~85-90% efficiency (linear would be 100%). - -## Multi-Node Training - -For training across multiple machines: - -```yaml -trainer: - strategy: ddp - devices: 4 # GPUs per node - num_nodes: 2 # Number of machines -``` - -Run on each node: - -```bash -# Node 0 -MASTER_ADDR=node0_address MASTER_PORT=12345 \ - lighter fit config.yaml \ - trainer::num_nodes=2 \ - trainer::devices=4 - -# Node 1 -MASTER_ADDR=node0_address MASTER_PORT=12345 NODE_RANK=1 \ - lighter fit config.yaml \ - trainer::num_nodes=2 \ - trainer::devices=4 -``` - -Requires: -- Shared filesystem for checkpoints -- Network connectivity between nodes -- Matching software environment - -## Quick Reference - -```yaml -# Basic multi-GPU -trainer: - devices: 4 - strategy: ddp - -# All GPUs -trainer: - devices: -1 - strategy: ddp - -# Large models -trainer: - strategy: fsdp - devices: 4 - -# Very large models -trainer: - strategy: - _target_: pytorch_lightning.strategies.DeepSpeedStrategy - stage: 3 - devices: 4 - -# Batch size scaling -vars: - num_gpus: 4 - total_batch: 512 - -data: - train_dataloader: - batch_size: "$%vars::total_batch // %vars::num_gpus" - -# LR scaling -model: - optimizer: - lr: "$%vars::base_lr * %vars::num_gpus" -``` - -## Next Steps - -- [Training Guide](../guides/training.md) - More training strategies -- [Best Practices](../guides/best-practices.md) - Production optimization -- [Image Classification Example](image-classification.md) - Complete example - -## Summary - -Multi-GPU training with Lighter: - -- ✅ Simple config changes only -- ✅ Code works unchanged -- ✅ Automatic gradient synchronization -- ✅ Linear scaling with proper settings -- ✅ Multiple strategies (DDP, FSDP, DeepSpeed) -- ✅ Works with all Lightning features - -**Key takeaway:** Add `trainer::devices=4 trainer::strategy=ddp` to use 4 GPUs. That's it! diff --git a/docs/guides/best-practices.md b/docs/guides/best-practices.md index dc184247..ecb509a4 100644 --- a/docs/guides/best-practices.md +++ b/docs/guides/best-practices.md @@ -974,7 +974,7 @@ Before production: ## Next Steps - [Training Guide](training.md) - Run experiments -- [Examples](../examples/image-classification.md) - Complete working code +- [Example Projects](../examples/index.md) - Complete working code - [FAQ](../faq.md) - Common questions ## Quick Reference diff --git a/docs/guides/custom-code.md b/docs/guides/custom-code.md index eef268d9..38f224eb 100644 --- a/docs/guides/custom-code.md +++ b/docs/guides/custom-code.md @@ -807,7 +807,7 @@ lighter fit configs/baseline.yaml - [Training Guide](training.md) - Run experiments, save outputs - [Best Practices](best-practices.md) - Production patterns -- [Complete Example](../examples/image-classification.md) - Full CIFAR-10 with all features +- [Example Projects](../examples/index.md) - Complete working examples ## Quick Reference diff --git a/docs/guides/training.md b/docs/guides/training.md index 8d676f5f..0bebd935 100644 --- a/docs/guides/training.md +++ b/docs/guides/training.md @@ -994,7 +994,7 @@ trainer: ## Next Steps - [Best Practices](best-practices.md) - Production patterns -- [Examples](../examples/image-classification.md) - Complete working examples +- [Example Projects](../examples/index.md) - Complete working examples - [CLI Reference](../reference/cli.md) - Full command documentation ## Quick Reference diff --git a/docs/index.md b/docs/index.md index 59d901f2..d45e734f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -340,7 +340,7 @@ Ready to try it? Pick your path: Full, working code you can copy-paste. - [:octicons-arrow-right-24: Examples](examples/image-classification.md) + [:octicons-arrow-right-24: Examples](examples/index.md) - :material-school:{ .lg .middle } **Guides** @@ -352,6 +352,23 @@ Ready to try it? Pick your path: +## Example Projects + +Ready-to-run projects demonstrating Lighter across domains: + +| Project | Domain | Features | +|---------|--------|----------| +| [cifar10](https://github.com/project-lighter/lighter/tree/main/projects/cifar10) | Image Classification | Basic setup, MetricCollection, FileWriter | +| [eeg](https://github.com/project-lighter/lighter/tree/main/projects/eeg) | EEG Analysis | Braindecode integration, regression | +| [huggingface_llm](https://github.com/project-lighter/lighter/tree/main/projects/huggingface_llm) | Sentiment Classification | Transformers, datasets, model-computed loss | +| [lora](https://github.com/project-lighter/lighter/tree/main/projects/lora) | Fine-Tuning | PEFT/LoRA, parameter filtering | +| [medical_segmentation](https://github.com/project-lighter/lighter/tree/main/projects/medical_segmentation) | Medical Imaging | MONAI, 3D volumes, sliding window | +| [self_supervised](https://github.com/project-lighter/lighter/tree/main/projects/self_supervised) | SSL Computer Vision | SimCLR, lightly library | +| [video_recognition](https://github.com/project-lighter/lighter/tree/main/projects/video_recognition) | Video | 3D CNNs, PytorchVideo | +| [vision_language](https://github.com/project-lighter/lighter/tree/main/projects/vision_language) | Vision-Language | CLIP-style dual encoders | + +Each project includes a README with setup instructions and demonstrates different Lighter features. + ## Community - [:fontawesome-brands-discord: Discord](https://discord.gg/zJcnp6KrUp) - Get help, share configs diff --git a/docs/quickstart.md b/docs/quickstart.md index 7fc2e236..8a97ad63 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -365,7 +365,7 @@ val_metrics: "%model::train_metrics" # New instance with same config ### See Complete Examples -[Image Classification](examples/image-classification.md) - Full CIFAR-10 example with all the bells and whistles +[Example Projects](examples/index.md) - Complete examples across various domains (image classification, medical imaging, NLP, and more) ### Organize Your Project diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 0a5104cb..60c63f67 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -527,7 +527,7 @@ def __init__(self, ...): - [Configuration Guide](../guides/configuration.md) - Learn config syntax - [Training Guide](../guides/training.md) - Training workflows -- [Examples](../examples/image-classification.md) - Complete examples +- [Example Projects](../examples/index.md) - Complete examples ## Quick Reference diff --git a/mkdocs.yml b/mkdocs.yml index 601741fb..f8d0dd93 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -79,9 +79,7 @@ nav: - LighterModule: guides/lighter-module.md - Training: guides/training.md - Best Practices: guides/best-practices.md - - Examples: - - Image Classification: examples/image-classification.md - - Multi-GPU Training: examples/multi-gpu.md + - Examples: examples/index.md - Reference: - CLI: reference/cli.md - API: reference/ diff --git a/projects/README.md b/projects/README.md index b03ae6a2..45e5d1cc 100644 --- a/projects/README.md +++ b/projects/README.md @@ -1,3 +1,14 @@ -# Projects Readme +# Example Projects -By default, Git will not track any projects or files added to this folder. If you wish to track them, please add `!projects/` to Lighter's gitignore. +Example projects demonstrating Lighter across domains. + +| Project | Domain | Extra Dependencies | +|---------|--------|-----------| +| [cifar10](./cifar10/) | Image Classification | - | +| [eeg](./eeg/) | EEG Analysis | braindecode, eegdash, mne | +| [huggingface_llm](./huggingface_llm/) | Sentiment Classification | transformers, datasets | +| [lora](./lora/) | Parameter-Efficient Fine-Tuning (LoRA) | peft | +| [medical_segmentation](./medical_segmentation/) | Medical Imaging | monai, itk | +| [self_supervised](./self_supervised/) | SSL Computer Vision | lightly | +| [video_recognition](./video_recognition/) | Video Understanding | pytorchvideo, av | +| [vision_language](./vision_language/) | Vision-Language | transformers | diff --git a/projects/cifar10/README.md b/projects/cifar10/README.md new file mode 100644 index 00000000..714b02c3 --- /dev/null +++ b/projects/cifar10/README.md @@ -0,0 +1,38 @@ +# CIFAR-10 Image Classification + +Reference example - simple CNN for CIFAR-10 classification. + +## Dataset + +**CIFAR-10** - 60,000 32x32 color images, 10 classes. Auto-downloaded. + +## Architecture + +3-layer CNN (~62K parameters): conv layers with ReLU/pooling, 2 FC layers. + +## Lighter Features Demonstrated + +- **`$` expressions** - `$@model::network.parameters()` +- **`%` raw references** - `%::train_metrics` for config reuse +- **MetricCollection** - Accuracy, F1, Precision, Recall +- **FileWriter callback** - saves predictions as tensors + +## Usage + +```bash +pip install lighter +cd projects/cifar10 + +# Train +lighter fit configs/example.yaml + +# Quick test +lighter fit configs/example.yaml trainer::fast_dev_run=true + +# Longer training +lighter fit configs/example.yaml trainer::max_epochs=50 +``` + +## References + +- [CIFAR-10 Dataset](https://www.cs.toronto.edu/~kriz/cifar.html) diff --git a/projects/cifar10/configs/example.yaml b/projects/cifar10/configs/example.yaml index ddc48c30..b5a08cd5 100644 --- a/projects/cifar10/configs/example.yaml +++ b/projects/cifar10/configs/example.yaml @@ -13,10 +13,10 @@ trainer: writer_fn: tensor model: - _target_: project.model.CIFAR10Model + _target_: project.models.model.CIFAR10Model network: - _target_: project.models.net.Net + _target_: project.networks.net.Net criterion: _target_: torch.nn.CrossEntropyLoss diff --git a/projects/cifar10/model.py b/projects/cifar10/models/model.py similarity index 86% rename from projects/cifar10/model.py rename to projects/cifar10/models/model.py index 1f6690d1..d68bee15 100644 --- a/projects/cifar10/model.py +++ b/projects/cifar10/models/model.py @@ -19,8 +19,6 @@ def training_step(self, batch, batch_idx): pred = self(x) # Compute loss using criterion from config - if self.criterion is None: - raise RuntimeError("criterion is required for training but was not set in config") loss = self.criterion(pred, y) # Update metrics (user calls them explicitly) @@ -34,8 +32,6 @@ def validation_step(self, batch, batch_idx): """Validation step with user-defined logic.""" x, y = batch pred = self(x) - if self.criterion is None: - raise RuntimeError("criterion is required for validation but was not set in config") loss = self.criterion(pred, y) if self.val_metrics is not None: diff --git a/projects/cifar10/networks/__init__.py b/projects/cifar10/networks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/cifar10/models/net.py b/projects/cifar10/networks/net.py similarity index 100% rename from projects/cifar10/models/net.py rename to projects/cifar10/networks/net.py diff --git a/projects/eeg/README.md b/projects/eeg/README.md new file mode 100644 index 00000000..ea719975 --- /dev/null +++ b/projects/eeg/README.md @@ -0,0 +1,57 @@ +# EEG 2025 Challenge + +Implementation for the [NeurIPS 2025 EEG Foundation Challenge](https://eeg2025.github.io/). + +## Challenges + +**Challenge 1**: Predict response time from CCD EEG data (RMSE metric) + +**Challenge 2**: Predict externalizing factor from EEG (MAE metric) + +Both tasks: regression on 129-channel EEG, 2s windows @ 100Hz. + +## Dataset + +**HBN-EEG** - ~3,000 participants, auto-downloaded via [EEGDash](https://github.com/eeg2025/eegdash). + +## Models + +Default: **EEGNeX** (state-of-the-art). Alternatives: EEGNet, EEGConformer (from Braindecode). + +## Lighter Features Demonstrated + +- **Braindecode integration** - EEGNeX, EEGNet models +- **EEGDash data loading** - automatic download and preprocessing +- **`%` raw references** - shared config values via `vars::` +- **MetricCollection** - MAE, RMSE tracking + +## Usage + +```bash +pip install lighter braindecode eegdash mne +cd projects/eeg + +# Challenge 1 +lighter fit configs/challenge1.yaml + +# Challenge 2 +lighter fit configs/challenge2.yaml + +# Quick test +lighter fit configs/challenge1.yaml trainer::fast_dev_run=true +``` + +## Submission + +```bash +python -m projects.eeg.submission --export --ckpt1 path/to/c1.ckpt --ckpt2 path/to/c2.ckpt +python -m projects.eeg.submission --zip +``` + +Submit to [CodaBench](https://www.codabench.org/competitions/9975/). + +## References + +- [EEG 2025 Challenge](https://eeg2025.github.io/) +- [Braindecode](https://braindecode.org/) +- [EEGDash](https://github.com/eeg2025/eegdash) diff --git a/projects/eeg/__init__.py b/projects/eeg/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/eeg/__lighter__.py b/projects/eeg/__lighter__.py new file mode 100644 index 00000000..ca1adb14 --- /dev/null +++ b/projects/eeg/__lighter__.py @@ -0,0 +1,9 @@ +# This file marks the directory as a Lighter project. +# EEG 2025 Challenge: Cross-Task and Cross-Subject EEG Decoding +# +# Features demonstrated: +# - EEG regression tasks (response time, psychopathology scores) +# - HBN-EEG dataset via EEGDash +# - Braindecode models (EEGNeX, EEGNet, EEGConformer) +# - Subject-based train/val/test splits +# - Competition submission generation diff --git a/projects/eeg/configs/challenge1.yaml b/projects/eeg/configs/challenge1.yaml new file mode 100644 index 00000000..d717bd99 --- /dev/null +++ b/projects/eeg/configs/challenge1.yaml @@ -0,0 +1,143 @@ +# EEG 2025 Challenge 1: Cross-Task Transfer Learning +# +# Predict response time from Contrast Change Detection (CCD) EEG data. +# This is a regression task using the HBN-EEG dataset. +# +# Challenge: https://eeg2025.github.io/ +# Starter Kit: https://github.com/eeg2025/startkit +# +# Requirements: +# pip install eegdash braindecode mne torch torchmetrics pytorch-lightning +# +# Run: +# lighter fit projects/eeg/configs/challenge1.yaml +# lighter fit projects/eeg/configs/challenge1.yaml trainer::fast_dev_run=false + +vars: + n_chans: 129 + n_outputs: 1 + n_times: 200 # 2 seconds at 100 Hz + sfreq: 100 + batch_size: 128 + num_workers: 4 + data_dir: .datasets/eeg2025 + release: R5 + mini: true # Set to false for full dataset + seed: 2025 + +trainer: + _target_: pytorch_lightning.Trainer + max_epochs: 100 + accelerator: auto + log_every_n_steps: 10 + + callbacks: + - _target_: pytorch_lightning.callbacks.EarlyStopping + monitor: val/loss/epoch + mode: min + patience: 10 + min_delta: 0.0001 + + - _target_: pytorch_lightning.callbacks.ModelCheckpoint + monitor: val/loss/epoch + mode: min + save_top_k: 3 + filename: challenge1-{epoch}-{val/loss/epoch:.4f} + + +model: + _target_: project.models.eeg_model.EEGRegressionModel + + # EEGNeX - recommended for EEG 2025 Challenge + network: + _target_: braindecode.models.EEGNeX + n_chans: "%vars::n_chans" + n_outputs: "%vars::n_outputs" + n_times: "%vars::n_times" + sfreq: "%vars::sfreq" + drop_prob: 0.5 + + # Alternative: EEGNet (uncomment below) + # network: + # _target_: braindecode.models.EEGNetv4 + # n_chans: "%vars::n_chans" + # n_outputs: "%vars::n_outputs" + # n_times: "%vars::n_times" + # drop_prob: 0.5 + + # Alternative: EEGConformer (uncomment below) + # network: + # _target_: braindecode.models.EEGConformer + # n_chans: "%vars::n_chans" + # n_outputs: "%vars::n_outputs" + # n_times: "%vars::n_times" + # drop_prob: 0.5 + + # MSE Loss for regression + criterion: + _target_: torch.nn.MSELoss + + optimizer: + _target_: torch.optim.AdamW + params: $@model::network.parameters() + lr: 0.001 + weight_decay: 0.00001 + + scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + optimizer: "@model::optimizer" + T_max: 100 + + train_metrics: + _target_: torchmetrics.MetricCollection + metrics: + - _target_: torchmetrics.regression.MeanAbsoluteError + - _target_: torchmetrics.regression.MeanSquaredError + squared: false # This gives RMSE + + val_metrics: "%model::train_metrics" + test_metrics: "%model::train_metrics" + +data: + _target_: lighter.LighterDataModule + + train_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: true + dataset: + _target_: project.data.hbn_dataset.HBNDatasetChallenge1 + data_dir: "%vars::data_dir" + release: "%vars::release" + mini: "%vars::mini" + split: train + seed: "%vars::seed" + + val_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: false + dataset: + _target_: project.data.hbn_dataset.HBNDatasetChallenge1 + data_dir: "%vars::data_dir" + release: "%vars::release" + mini: "%vars::mini" + split: val + seed: "%vars::seed" + + test_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + shuffle: false + dataset: + _target_: project.data.hbn_dataset.HBNDatasetChallenge1 + data_dir: "%vars::data_dir" + release: "%vars::release" + mini: "%vars::mini" + split: test + seed: "%vars::seed" diff --git a/projects/eeg/configs/challenge2.yaml b/projects/eeg/configs/challenge2.yaml new file mode 100644 index 00000000..06c133d9 --- /dev/null +++ b/projects/eeg/configs/challenge2.yaml @@ -0,0 +1,132 @@ +# EEG 2025 Challenge 2: Externalizing Factor Prediction +# +# Predict psychopathology scores (externalizing factor) from EEG recordings. +# This is a regression task requiring subject-invariant representations. +# +# Challenge: https://eeg2025.github.io/ +# Starter Kit: https://github.com/eeg2025/startkit +# +# Requirements: +# pip install eegdash braindecode mne torch torchmetrics pytorch-lightning +# +# Run: +# lighter fit projects/eeg/configs/challenge2.yaml +# lighter fit projects/eeg/configs/challenge2.yaml trainer::fast_dev_run=false + +vars: + n_chans: 129 + n_outputs: 1 + n_times: 200 # 2 seconds at 100 Hz (after cropping from 4s windows) + sfreq: 100 + batch_size: 128 + num_workers: 4 + data_dir: .datasets/eeg2025 + releases: + - R5 + task: contrastChangeDetection + mini: true # Set to false for full dataset + seed: 2025 + +trainer: + _target_: pytorch_lightning.Trainer + max_epochs: 100 + accelerator: auto + log_every_n_steps: 10 + + callbacks: + - _target_: pytorch_lightning.callbacks.EarlyStopping + monitor: val/loss/epoch + mode: min + patience: 15 + min_delta: 0.0001 + + - _target_: pytorch_lightning.callbacks.ModelCheckpoint + monitor: val/loss/epoch + mode: min + save_top_k: 3 + filename: challenge2-{epoch}-{val/loss/epoch:.4f} + + +model: + _target_: project.models.eeg_model.EEGRegressionModel + + # EEGNeX - recommended for EEG 2025 Challenge + network: + _target_: braindecode.models.EEGNeX + n_chans: "%vars::n_chans" + n_outputs: "%vars::n_outputs" + n_times: "%vars::n_times" + sfreq: "%vars::sfreq" + drop_prob: 0.5 + + # L1 Loss (MAE) - more robust for psychopathology prediction + criterion: + _target_: torch.nn.L1Loss + + optimizer: + _target_: torch.optim.Adamax + params: $@model::network.parameters() + lr: 0.002 + + scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + optimizer: "@model::optimizer" + T_max: 100 + + train_metrics: + _target_: torchmetrics.MetricCollection + metrics: + - _target_: torchmetrics.regression.MeanAbsoluteError + - _target_: torchmetrics.regression.MeanSquaredError + squared: false # RMSE + - _target_: torchmetrics.regression.R2Score + + val_metrics: "%model::train_metrics" + test_metrics: "%model::train_metrics" + +data: + _target_: lighter.LighterDataModule + + train_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: true + dataset: + _target_: project.data.hbn_dataset.HBNDatasetChallenge2 + data_dir: "%vars::data_dir" + releases: "%vars::releases" + task: "%vars::task" + mini: "%vars::mini" + split: train + seed: "%vars::seed" + + val_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: false + dataset: + _target_: project.data.hbn_dataset.HBNDatasetChallenge2 + data_dir: "%vars::data_dir" + releases: "%vars::releases" + task: "%vars::task" + mini: "%vars::mini" + split: val + seed: "%vars::seed" + + test_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + shuffle: false + dataset: + _target_: project.data.hbn_dataset.HBNDatasetChallenge2 + data_dir: "%vars::data_dir" + releases: "%vars::releases" + task: "%vars::task" + mini: "%vars::mini" + split: test + seed: "%vars::seed" diff --git a/projects/eeg/data/__init__.py b/projects/eeg/data/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/eeg/data/hbn_dataset.py b/projects/eeg/data/hbn_dataset.py new file mode 100644 index 00000000..25cd0a54 --- /dev/null +++ b/projects/eeg/data/hbn_dataset.py @@ -0,0 +1,396 @@ +"""HBN-EEG Dataset loading for EEG 2025 Challenge. + +This module provides dataset classes for both challenges using EEGDash. + +Challenge 1: Cross-Task Transfer Learning +- Task: Predict response time from Contrast Change Detection (CCD) EEG +- Target: rt_from_stimulus (response time) +- Preprocessing: Window extraction around stimulus events + +Challenge 2: Externalizing Factor Prediction +- Task: Predict psychopathology scores from EEG +- Target: externalizing factor (p_factor in dataset) +- Preprocessing: Fixed-length windows with random cropping + +Dataset: +- HBN-EEG: ~3000 subjects with EEG recordings +- 129 EEG channels (EGI system) +- Multiple paradigms: CCD, Resting State, etc. + +Requirements: + pip install eegdash braindecode mne +""" + +import math +import random +from pathlib import Path + +import numpy as np +from braindecode.datasets.base import BaseConcatDataset +from braindecode.preprocessing import ( + Preprocessor, + create_fixed_length_windows, + create_windows_from_events, + preprocess, +) +from eegdash.dataset import EEGChallengeDataset +from eegdash.hbn.windows import ( + add_aux_anchors, + add_extras_columns, + annotate_trials_with_target, + keep_only_recordings_with, +) +from sklearn.model_selection import train_test_split +from sklearn.utils import check_random_state +from torch.utils.data import Dataset + +# Subjects to exclude (corrupted/incomplete data) +EXCLUDED_SUBJECTS = [ + "NDARWV769JM7", + "NDARME789TD2", + "NDARUA442ZVF", + "NDARJP304NK1", + "NDARTY128YLU", + "NDARDW550GU6", + "NDARLD243KRE", + "NDARUJ292JXV", + "NDARBA381JGH", +] + +# Default parameters +DEFAULT_SFREQ = 100 # Sampling frequency after resampling +DEFAULT_N_CHANS = 129 # Number of EEG channels + + +class HBNDatasetChallenge1(Dataset): + """HBN-EEG Dataset for Challenge 1: Cross-Task Transfer Learning. + + Predicts response time from Contrast Change Detection (CCD) EEG data. + + Args: + data_dir: Directory for dataset cache. + release: Data release version ('R5', 'R6', etc.). + mini: Whether to use mini version for testing. + split: Data split ('train', 'val', 'test'). + seed: Random seed for reproducibility. + epoch_len_s: Length of each epoch in seconds. + shift_after_stim: Time shift after stimulus onset. + window_len: Window length in seconds. + """ + + def __init__( + self, + data_dir: str = "data", + release: str = "R5", + mini: bool = False, + split: str = "train", + seed: int = 2025, + epoch_len_s: float = 2.0, + shift_after_stim: float = 0.5, + window_len: float = 2.0, + valid_frac: float = 0.1, + test_frac: float = 0.1, + ) -> None: + super().__init__() + if valid_frac + test_frac <= 0: + raise ValueError("valid_frac + test_frac must be > 0") + if valid_frac + test_frac >= 1.0: + raise ValueError("valid_frac + test_frac must be < 1.0") + self.data_dir = Path(data_dir) + self.data_dir.mkdir(parents=True, exist_ok=True) + self.split = split + self.sfreq = DEFAULT_SFREQ + + # Load dataset + dataset = EEGChallengeDataset( + task="contrastChangeDetection", + release=release, + cache_dir=self.data_dir, + mini=mini, + ) + + # Preprocessing: annotate trials with response time target + transformation = [ + Preprocessor( + annotate_trials_with_target, + target_field="rt_from_stimulus", + epoch_length=epoch_len_s, + require_stimulus=True, + require_response=True, + apply_on_array=False, + ), + Preprocessor(add_aux_anchors, apply_on_array=False), + ] + preprocess(dataset, transformation, n_jobs=1) + + # Create windows around stimulus events + anchor = "stimulus_anchor" + dataset = keep_only_recordings_with(anchor, dataset) + + self.windows_dataset = create_windows_from_events( + dataset, + mapping={anchor: 0}, + trial_start_offset_samples=int(shift_after_stim * self.sfreq), + trial_stop_offset_samples=int((shift_after_stim + window_len) * self.sfreq), + window_size_samples=int(epoch_len_s * self.sfreq), + window_stride_samples=self.sfreq, + preload=True, + ) + + # Add extra columns with target values + self.windows_dataset = add_extras_columns( + self.windows_dataset, + dataset, + desc=anchor, + keys=( + "target", + "rt_from_stimulus", + "rt_from_trialstart", + "stimulus_onset", + "response_onset", + "correct", + "response_type", + ), + ) + + # Split by subject + meta = self.windows_dataset.get_metadata() + subjects = [s for s in meta["subject"].unique() if s not in EXCLUDED_SUBJECTS] + + train_subj, valid_test_subj = train_test_split( + subjects, + test_size=(valid_frac + test_frac), + random_state=check_random_state(seed), + shuffle=True, + ) + valid_subj, test_subj = train_test_split( + valid_test_subj, + test_size=test_frac / (valid_frac + test_frac), + random_state=check_random_state(seed + 1), + shuffle=True, + ) + + # Select appropriate split + split_map = {"train": train_subj, "val": valid_subj, "test": test_subj} + target_subjects = set(split_map[split]) + + subject_split = self.windows_dataset.split("subject") + split_datasets = [subject_split[s] for s in subject_split if s in target_subjects] + + self.dataset = BaseConcatDataset(split_datasets) if split_datasets else None + + # Cache metadata for efficient access + self._metadata = self.dataset.get_metadata() if self.dataset else None + + def __len__(self) -> int: + return len(self.dataset) if self.dataset else 0 + + def __getitem__(self, idx: int) -> tuple: + X, y, window_info = self.dataset[idx] + # Get target from cached metadata + target = float(self._metadata.iloc[idx].get("rt_from_stimulus", y)) + # Convert to float32 for MPS compatibility (MPS doesn't support float64) + X = X.astype(np.float32) + return X, np.float32(target) + + +class HBNDatasetChallenge2(Dataset): + """HBN-EEG Dataset for Challenge 2: Externalizing Factor Prediction. + + Predicts psychopathology scores (externalizing factor) from EEG. + + Args: + data_dir: Directory for dataset cache. + releases: List of data release versions. + task: EEG paradigm to use. + mini: Whether to use mini version for testing. + split: Data split ('train', 'val', 'test'). + seed: Random seed for reproducibility. + window_size_s: Window size in seconds. + window_stride_s: Window stride in seconds. + crop_size_s: Crop size for random cropping. + target_name: Name of target variable in metadata. + """ + + def __init__( + self, + data_dir: str = "data", + releases: list[str] | None = None, + task: str = "contrastChangeDetection", + mini: bool = False, + split: str = "train", + seed: int = 2025, + window_size_s: float = 4.0, + window_stride_s: float = 2.0, + crop_size_s: float = 2.0, + target_name: str = "p_factor", + valid_frac: float = 0.1, + test_frac: float = 0.1, + ) -> None: + super().__init__() + if valid_frac + test_frac <= 0: + raise ValueError("valid_frac + test_frac must be > 0") + if valid_frac + test_frac >= 1.0: + raise ValueError("valid_frac + test_frac must be < 1.0") + self.data_dir = Path(data_dir) + self.data_dir.mkdir(parents=True, exist_ok=True) + self.split = split + self.sfreq = DEFAULT_SFREQ + self.crop_size_samples = int(crop_size_s * self.sfreq) + self.target_name = target_name + self.rng = random.Random(seed) + + if releases is None: + releases = ["R5"] + + # Load datasets from multiple releases + all_datasets = [] + for release in releases: + ds = EEGChallengeDataset( + release=release, + task=task, + mini=mini, + description_fields=[ + "subject", + "session", + "run", + "task", + "age", + "gender", + "sex", + "p_factor", + ], + cache_dir=self.data_dir, + ) + all_datasets.append(ds) + + combined = BaseConcatDataset(all_datasets) if len(all_datasets) > 1 else all_datasets[0] + + # Filter valid recordings + valid_datasets = [] + for ds in combined.datasets: + # Check: valid subject, enough samples, correct channels, valid p_factor + if ( + ds.description.subject not in EXCLUDED_SUBJECTS + and ds.raw.n_times >= int(window_size_s * self.sfreq) + and len(ds.raw.ch_names) == DEFAULT_N_CHANS + and not math.isnan(ds.description.get("p_factor", float("nan"))) + ): + valid_datasets.append(ds) + + if not valid_datasets: + self.dataset = None + self._window_datasets = [] + return + + filtered = BaseConcatDataset(valid_datasets) + + # Create fixed-length windows + windows_ds = create_fixed_length_windows( + filtered, + window_size_samples=int(window_size_s * self.sfreq), + window_stride_samples=int(window_stride_s * self.sfreq), + drop_last_window=True, + ) + + # Get unique subjects and split + all_subjects = list({ds.description.subject for ds in filtered.datasets}) + all_subjects = [s for s in all_subjects if s not in EXCLUDED_SUBJECTS] + + train_subj, valid_test_subj = train_test_split( + all_subjects, + test_size=(valid_frac + test_frac), + random_state=check_random_state(seed), + shuffle=True, + ) + valid_subj, test_subj = train_test_split( + valid_test_subj, + test_size=test_frac / (valid_frac + test_frac), + random_state=check_random_state(seed + 1), + shuffle=True, + ) + + split_map = {"train": train_subj, "val": valid_subj, "test": test_subj} + target_subjects = set(split_map[split]) + + # Filter windows by subject + self._window_datasets = [] + for wds in windows_ds.datasets: + if wds.description.subject in target_subjects: + self._window_datasets.append(wds) + + self.dataset = BaseConcatDataset(self._window_datasets) if self._window_datasets else None + + def __len__(self) -> int: + if not self._window_datasets: + return 0 + return sum(len(wds) for wds in self._window_datasets) + + def __getitem__(self, idx: int) -> tuple: + # Find which window dataset and local index + cumsum = 0 + for wds in self._window_datasets: + if idx < cumsum + len(wds): + local_idx = idx - cumsum + X, _, crop_inds = wds[local_idx] + + # Get target + target = float(wds.description.get(self.target_name, 0.0)) + + # Random crop + i_window, i_start, i_stop = crop_inds + window_len = i_stop - i_start + if window_len > self.crop_size_samples: + start_offset = self.rng.randint(0, window_len - self.crop_size_samples) + X = X[:, start_offset : start_offset + self.crop_size_samples] + + # Convert to float32 for MPS compatibility + X = X.astype(np.float32) + return X, np.float32(target) + cumsum += len(wds) + + raise IndexError(f"Index {idx} out of range") + + +def get_train_val_test_split( + subjects: list[str], + valid_frac: float = 0.1, + test_frac: float = 0.1, + seed: int = 2025, +) -> tuple[list[str], list[str], list[str]]: + """Split subjects into train/val/test sets. + + Args: + subjects: List of subject IDs. + valid_frac: Fraction for validation. + test_frac: Fraction for test. + seed: Random seed. + + Returns: + Tuple of (train_subjects, val_subjects, test_subjects). + + Raises: + ValueError: If valid_frac + test_frac is not in (0, 1). + """ + if valid_frac + test_frac <= 0: + raise ValueError("valid_frac + test_frac must be > 0") + if valid_frac + test_frac >= 1.0: + raise ValueError("valid_frac + test_frac must be < 1.0") + + # Remove excluded subjects + subjects = [s for s in subjects if s not in EXCLUDED_SUBJECTS] + + train_subj, valid_test_subj = train_test_split( + subjects, + test_size=(valid_frac + test_frac), + random_state=check_random_state(seed), + shuffle=True, + ) + valid_subj, test_subj = train_test_split( + valid_test_subj, + test_size=test_frac / (valid_frac + test_frac), + random_state=check_random_state(seed + 1), + shuffle=True, + ) + + return train_subj, valid_subj, test_subj diff --git a/projects/eeg/models/__init__.py b/projects/eeg/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/eeg/models/eeg_model.py b/projects/eeg/models/eeg_model.py new file mode 100644 index 00000000..1363e7c4 --- /dev/null +++ b/projects/eeg/models/eeg_model.py @@ -0,0 +1,178 @@ +"""EEG Regression Model for EEG 2025 Challenge. + +This module implements the LighterModule for both challenges: +- Challenge 1: Cross-Task Transfer Learning (predict response time) +- Challenge 2: Externalizing Factor Prediction (predict psychopathology scores) + +Both are regression tasks using EEG data from the HBN-EEG dataset. +""" + +from typing import Any + +from lighter import LighterModule + + +class EEGRegressionModel(LighterModule): + """EEG Regression Model for the EEG 2025 Challenge. + + Handles EEG input tensors of shape [B, C, T] where: + - B: batch size + - C: number of EEG channels (129 for HBN-EEG) + - T: number of time samples (200 for 2s at 100Hz) + + Supports both challenges: + - Challenge 1: Predict response time from CCD task + - Challenge 2: Predict externalizing factor scores + + Attributes: + network: The neural network backbone (EEGNeX, EEGNet, etc.) + criterion: Loss function (MSELoss or L1Loss for regression) + optimizer: Optimizer for training + scheduler: Learning rate scheduler (optional) + train_metrics: Training metrics (MAE, RMSE) + val_metrics: Validation metrics + test_metrics: Test metrics + """ + + def _shared_step(self, batch: tuple, metrics) -> dict[str, Any]: + """Shared logic for train/val/test steps. + + Args: + batch: Tuple of (EEG data, target values, crop_indices, metadata) + or (EEG data, target values) for simpler datasets + metrics: MetricCollection or None + + Returns: + Dictionary with loss, predictions, and targets + """ + # Handle different batch formats from EEGDash + if len(batch) == 4: + # Full format: (X, y, crop_inds, infos) + eeg, target = batch[0], batch[1] + elif len(batch) == 3: + # Format: (X, y, crop_inds) + eeg, target = batch[0], batch[1] + else: + # Simple format: (X, y) + eeg, target = batch + + # Ensure correct dtype + eeg = eeg.float() + target = target.float() + + # Ensure target has correct shape [B, 1] + if target.dim() == 1: + target = target.unsqueeze(1) + + # Forward pass + pred = self(eeg) + + # Compute loss + loss = self.criterion(pred, target) + + # Update metrics + if metrics is not None: + # Flatten for metrics + metrics(pred.squeeze(), target.squeeze()) + + return { + "loss": loss, + "pred": pred.squeeze(), + "target": target.squeeze(), + } + + def training_step(self, batch: tuple, batch_idx: int) -> dict[str, Any]: + """Training step with loss computation and metric updates.""" + return self._shared_step(batch, self.train_metrics) + + def validation_step(self, batch: tuple, batch_idx: int) -> dict[str, Any]: + """Validation step with loss and metrics.""" + return self._shared_step(batch, self.val_metrics) + + def test_step(self, batch: tuple, batch_idx: int) -> dict[str, Any]: + """Test step with metrics only (no loss required).""" + result = self._shared_step(batch, self.test_metrics) + # Optionally remove loss for test + return result + + def predict_step(self, batch: tuple, batch_idx: int) -> dict[str, Any]: + """Prediction step for inference and submission. + + Returns predictions along with metadata for aggregation. + """ + # Handle different batch formats + if len(batch) == 4: + eeg, target, crop_inds, infos = batch + elif len(batch) == 3: + eeg, target, crop_inds = batch + infos = None + else: + eeg, target = batch if isinstance(batch, (tuple, list)) else (batch, None) + infos = None + + eeg = eeg.float() + + # Forward pass + pred = self(eeg) + + result = { + "prediction": pred.squeeze().tolist() if pred.dim() > 0 else [pred.item()], + } + + if target is not None: + target = target.float() + if target.dim() == 1: + target = target.unsqueeze(1) + result["target"] = target.squeeze().tolist() + + if infos is not None: + result["subject"] = infos.get("subject", []) + + return result + + +class EEGWindowAggregator: + """Aggregates predictions from multiple windows per subject. + + For the EEG 2025 Challenge, multiple windows are extracted per recording. + This class aggregates window-level predictions to subject/recording level. + + Aggregation strategies: + - mean: Average of all window predictions (default) + - median: Median of all window predictions + + Args: + strategy: Aggregation strategy ('mean', 'median') + """ + + def __init__(self, strategy: str = "mean") -> None: + self.strategy = strategy + self.predictions: dict[str, list[float]] = {} + + def add(self, subject: str, prediction: float) -> None: + """Add a window prediction for a subject.""" + if subject not in self.predictions: + self.predictions[subject] = [] + self.predictions[subject].append(prediction) + + def aggregate(self) -> dict[str, float]: + """Aggregate all predictions by subject.""" + results = {} + for subject, preds in self.predictions.items(): + if self.strategy == "mean": + results[subject] = sum(preds) / len(preds) + elif self.strategy == "median": + sorted_preds = sorted(preds) + mid = len(sorted_preds) // 2 + if len(sorted_preds) % 2 == 0: + results[subject] = (sorted_preds[mid - 1] + sorted_preds[mid]) / 2 + else: + results[subject] = sorted_preds[mid] + else: + # Default to mean + results[subject] = sum(preds) / len(preds) + return results + + def reset(self) -> None: + """Clear all stored predictions.""" + self.predictions.clear() diff --git a/projects/eeg/submission.py b/projects/eeg/submission.py new file mode 100644 index 00000000..f1796f76 --- /dev/null +++ b/projects/eeg/submission.py @@ -0,0 +1,312 @@ +"""Competition submission template for EEG 2025 Challenge. + +This module provides the Submission class required by CodaBench. + +Submission format: +- submission.py (this file) +- weights_challenge_1.pt (trained weights for Challenge 1) +- weights_challenge_2.pt (trained weights for Challenge 2) + +All files should be in a flat zip archive for submission. + +Usage: + # After training, export weights and create submission: + python -m projects.eeg.submission --export + + # Test submission locally: + python -m projects.eeg.submission --test +""" + +import argparse +import zipfile +from pathlib import Path + +import torch +from braindecode.models import EEGNeX + +# Default parameters matching the challenge +DEFAULT_SFREQ = 100 +DEFAULT_N_CHANS = 129 +DEFAULT_N_TIMES = 200 # 2 seconds at 100 Hz + + +def resolve_path(name: str = "model_file_name") -> str: + """Resolve model weight file path across different environments. + + CodaBench may place files in different locations depending on + the execution environment. + + Args: + name: Name of the weight file. + + Returns: + Resolved path to the file. + + Raises: + FileNotFoundError: If file cannot be found. + """ + # Check various possible locations + candidates = [ + Path(f"/app/input/res/{name}"), + Path(f"/app/input/{name}"), + Path(name), + Path(__file__).parent / name, + ] + + for path in candidates: + if path.exists(): + return str(path) + + raise FileNotFoundError(f"Could not find {name} in expected locations: {[str(c) for c in candidates]}") + + +class Submission: + """Competition submission class for EEG 2025 Challenge. + + This class is instantiated by the CodaBench evaluation system + and used to load models for both challenges. + + Args: + SFREQ: Sampling frequency (100 Hz for HBN-EEG). + DEVICE: Device to load models on ('cuda' or 'cpu'). + """ + + def __init__(self, SFREQ: int, DEVICE: str) -> None: + self.sfreq = SFREQ + self.device = DEVICE + self.n_times = int(2 * SFREQ) # 2 second windows + + def get_model_challenge_1(self) -> torch.nn.Module: + """Load trained model for Challenge 1 (Response Time Prediction). + + Returns: + Loaded EEGNeX model ready for inference. + """ + model = EEGNeX( + n_chans=DEFAULT_N_CHANS, + n_outputs=1, + n_times=self.n_times, + sfreq=self.sfreq, + ).to(self.device) + + # Load trained weights + weights_path = resolve_path("weights_challenge_1.pt") + state_dict = torch.load(weights_path, map_location=self.device, weights_only=True) + + # Handle potential 'model.' prefix from wrapper + if any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()} + + model.load_state_dict(state_dict) + model.eval() + + return model + + def get_model_challenge_2(self) -> torch.nn.Module: + """Load trained model for Challenge 2 (Externalizing Factor). + + Returns: + Loaded EEGNeX model ready for inference. + """ + model = EEGNeX( + n_chans=DEFAULT_N_CHANS, + n_outputs=1, + n_times=self.n_times, + sfreq=self.sfreq, + ).to(self.device) + + # Load trained weights + weights_path = resolve_path("weights_challenge_2.pt") + state_dict = torch.load(weights_path, map_location=self.device, weights_only=True) + + # Handle potential 'model.' prefix from wrapper + if any(k.startswith("model.") for k in state_dict.keys()): + state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()} + + model.load_state_dict(state_dict) + model.eval() + + return model + + +def export_weights( + checkpoint_path_1: str | None = None, + checkpoint_path_2: str | None = None, + output_dir: str = "submission", +) -> None: + """Export trained model weights for submission. + + Args: + checkpoint_path_1: Path to Challenge 1 checkpoint (Lightning format). + checkpoint_path_2: Path to Challenge 2 checkpoint (Lightning format). + output_dir: Directory to save exported weights. + """ + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + for challenge, ckpt_path in [ + (1, checkpoint_path_1), + (2, checkpoint_path_2), + ]: + if ckpt_path is None: + print(f"Skipping Challenge {challenge} - no checkpoint provided") + continue + + ckpt_path = Path(ckpt_path) + if not ckpt_path.exists(): + print(f"Warning: Checkpoint not found: {ckpt_path}") + continue + + # Load Lightning checkpoint + checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False) + state_dict = checkpoint.get("state_dict", checkpoint) + + # Extract network weights (remove 'network.' prefix from LighterModule) + network_state = {} + for key, value in state_dict.items(): + if key.startswith("network."): + # Remove 'network.' prefix + new_key = key[8:] + # Also handle 'model.' prefix from wrapper + if new_key.startswith("model."): + new_key = new_key[6:] + network_state[new_key] = value + + # Save weights + output_file = output_path / f"weights_challenge_{challenge}.pt" + torch.save(network_state, output_file) + print(f"Exported Challenge {challenge} weights to: {output_file}") + + +def create_submission_zip( + output_dir: str = "submission", + zip_name: str = "submission.zip", +) -> None: + """Create submission zip file. + + Args: + output_dir: Directory containing submission files. + zip_name: Name of the output zip file. + """ + output_path = Path(output_dir) + zip_path = output_path / zip_name + + required_files = [ + "submission.py", + "weights_challenge_1.pt", + "weights_challenge_2.pt", + ] + + # Copy submission.py to output directory + import shutil + + shutil.copy(__file__, output_path / "submission.py") + + # Create zip + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: + for filename in required_files: + filepath = output_path / filename + if filepath.exists(): + zf.write(filepath, filename) + print(f"Added to zip: {filename}") + else: + print(f"Warning: Missing file: {filename}") + + print(f"\nCreated submission: {zip_path}") + + +def test_submission(output_dir: str = "submission") -> None: + """Test submission locally. + + Args: + output_dir: Directory containing submission files. + """ + import sys + + sys.path.insert(0, str(Path(output_dir))) + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Testing on device: {device}") + + try: + sub = Submission(SFREQ=DEFAULT_SFREQ, DEVICE=device) + + # Test Challenge 1 + print("\nTesting Challenge 1 model...") + model1 = sub.get_model_challenge_1() + x = torch.randn(2, DEFAULT_N_CHANS, DEFAULT_N_TIMES).to(device) + with torch.no_grad(): + y = model1(x) + print(f" Input shape: {x.shape}") + print(f" Output shape: {y.shape}") + print(f" Output: {y.squeeze().tolist()}") + + # Test Challenge 2 + print("\nTesting Challenge 2 model...") + model2 = sub.get_model_challenge_2() + with torch.no_grad(): + y = model2(x) + print(f" Input shape: {x.shape}") + print(f" Output shape: {y.shape}") + print(f" Output: {y.squeeze().tolist()}") + + print("\nSubmission test PASSED!") + + except Exception as e: + print(f"\nSubmission test FAILED: {e}") + raise + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="EEG 2025 Challenge Submission Tools") + parser.add_argument( + "--export", + action="store_true", + help="Export trained weights for submission", + ) + parser.add_argument( + "--test", + action="store_true", + help="Test submission locally", + ) + parser.add_argument( + "--zip", + action="store_true", + help="Create submission zip file", + ) + parser.add_argument( + "--ckpt1", + type=str, + default=None, + help="Path to Challenge 1 checkpoint", + ) + parser.add_argument( + "--ckpt2", + type=str, + default=None, + help="Path to Challenge 2 checkpoint", + ) + parser.add_argument( + "--output-dir", + type=str, + default="submission", + help="Output directory for submission files", + ) + + args = parser.parse_args() + + if args.export: + export_weights( + checkpoint_path_1=args.ckpt1, + checkpoint_path_2=args.ckpt2, + output_dir=args.output_dir, + ) + + if args.zip: + create_submission_zip(output_dir=args.output_dir) + + if args.test: + test_submission(output_dir=args.output_dir) + + if not any([args.export, args.test, args.zip]): + parser.print_help() diff --git a/projects/huggingface_llm/README.md b/projects/huggingface_llm/README.md new file mode 100644 index 00000000..be8bb538 --- /dev/null +++ b/projects/huggingface_llm/README.md @@ -0,0 +1,39 @@ +# HuggingFace Text Classification + +DistilBERT sentiment classification on IMDB reviews. + +## Dataset + +**IMDB** - 50,000 movie reviews, binary sentiment. Auto-downloaded via HuggingFace `datasets`. + +## Architecture + +**DistilBERT** - 66M parameters, 6 layers. Pre-trained on Wikipedia + BookCorpus. + +## Lighter Features Demonstrated + +- **HuggingFace integration** - `transformers` and `datasets` +- **Model-computed loss** - no criterion config needed +- **`vars::` section** - reusable config values +- **CsvWriter callback** - prediction logging + +## Usage + +```bash +pip install lighter transformers datasets +cd projects/huggingface_llm + +# Train +lighter fit configs/imdb.yaml + +# Quick test +lighter fit configs/imdb.yaml trainer::fast_dev_run=true + +# Different model +lighter fit configs/imdb.yaml vars::pretrained_model=bert-base-uncased +``` + +## References + +- [HuggingFace Transformers](https://huggingface.co/docs/transformers) +- [DistilBERT Paper](https://arxiv.org/abs/1910.01108) diff --git a/projects/huggingface_llm/__init__.py b/projects/huggingface_llm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/huggingface_llm/__lighter__.py b/projects/huggingface_llm/__lighter__.py new file mode 100644 index 00000000..f3ffbab8 --- /dev/null +++ b/projects/huggingface_llm/__lighter__.py @@ -0,0 +1 @@ +# Lighter project marker - enables `project.` imports in configs diff --git a/projects/huggingface_llm/configs/imdb.yaml b/projects/huggingface_llm/configs/imdb.yaml new file mode 100644 index 00000000..815f4215 --- /dev/null +++ b/projects/huggingface_llm/configs/imdb.yaml @@ -0,0 +1,77 @@ +# HuggingFace LLM Text Classification +# +# Dataset: IMDB sentiment classification +# Model: DistilBERT +# +# Requirements: +# cd projects/huggingface_llm +# uv sync +# +# Run: +# uv run lighter fit configs/imdb.yaml +# +# Or from repo root: +# lighter fit projects/huggingface_llm/configs/imdb.yaml + +vars: + num_classes: 2 + pretrained_model: distilbert-base-uncased + batch_size: 16 + lr: 2.0e-5 + +trainer: + _target_: pytorch_lightning.Trainer + max_epochs: 3 + accelerator: auto + + callbacks: + - _target_: lighter.callbacks.CsvWriter + path: predictions.csv + keys: ["prediction"] + +model: + _target_: project.models.model.TextClassificationModel + + network: + _target_: transformers.AutoModelForSequenceClassification.from_pretrained + pretrained_model_name_or_path: "%vars::pretrained_model" + num_labels: "%vars::num_classes" + + optimizer: + _target_: torch.optim.AdamW + params: "$@model::network.parameters()" + lr: "%vars::lr" + + train_metrics: + _target_: torchmetrics.MetricCollection + metrics: + - _target_: torchmetrics.Accuracy + task: multiclass + num_classes: "%vars::num_classes" + + val_metrics: "%model::train_metrics" + +data: + _target_: lighter.LighterDataModule + + train_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + shuffle: true + dataset: + _target_: project.dataset.TextClassificationDataset + dataset_name: imdb + split: train + tokenizer_name: "%vars::pretrained_model" + + val_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + shuffle: false + dataset: + _target_: project.dataset.TextClassificationDataset + dataset_name: imdb + split: test + tokenizer_name: "%vars::pretrained_model" + + predict_dataloader: "%data::val_dataloader" diff --git a/projects/huggingface_llm/dataset.py b/projects/huggingface_llm/dataset.py new file mode 100644 index 00000000..dfb0e652 --- /dev/null +++ b/projects/huggingface_llm/dataset.py @@ -0,0 +1,33 @@ +from datasets import load_dataset +from torch.utils.data import Dataset +from transformers import AutoTokenizer + + +class TextClassificationDataset(Dataset): + def __init__(self, dataset_name, split, tokenizer_name, max_length=128): + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) + + # Load the dataset + self.dataset = load_dataset(dataset_name, split=split) + + # Reduce dataset size for faster demo + if "train" in split: + self.dataset = self.dataset.shuffle(seed=42).select(range(2000)) + else: + self.dataset = self.dataset.shuffle(seed=42).select(range(500)) + + # Tokenize the dataset + self.dataset = self.dataset.map( + lambda e: self.tokenizer(e["text"], truncation=True, padding="max_length", max_length=max_length), batched=True + ) + + # Rename 'label' to 'labels' for Hugging Face model compatibility + self.dataset = self.dataset.rename_column("label", "labels") + + self.dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"]) + + def __len__(self): + return len(self.dataset) + + def __getitem__(self, idx): + return self.dataset[idx] diff --git a/projects/huggingface_llm/models/__init__.py b/projects/huggingface_llm/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/huggingface_llm/models/model.py b/projects/huggingface_llm/models/model.py new file mode 100644 index 00000000..aa25f02c --- /dev/null +++ b/projects/huggingface_llm/models/model.py @@ -0,0 +1,49 @@ +"""Text classification model using LighterModule.""" + +from lighter import LighterModule + + +class TextClassificationModel(LighterModule): + """ + HuggingFace text classification model wrapper. + + The HuggingFace model computes its own loss when labels are provided, + so we don't need a separate criterion. + """ + + def _shared_step(self, batch, metrics): + """Shared step logic for train/val/test.""" + # HuggingFace models expect named arguments + outputs = self.network( + input_ids=batch["input_ids"], + attention_mask=batch["attention_mask"], + labels=batch["labels"], + ) + + # Update metrics if available + if metrics is not None: + preds = outputs.logits.argmax(dim=-1) + metrics(preds, batch["labels"]) + + return outputs + + def training_step(self, batch, batch_idx): + outputs = self._shared_step(batch, self.train_metrics) + return {"loss": outputs.loss} + + def validation_step(self, batch, batch_idx): + outputs = self._shared_step(batch, self.val_metrics) + return {"loss": outputs.loss} + + def test_step(self, batch, batch_idx): + outputs = self._shared_step(batch, self.test_metrics) + return {"loss": outputs.loss} + + def predict_step(self, batch, batch_idx): + """Prediction step - return predictions for CsvWriter.""" + outputs = self.network( + input_ids=batch["input_ids"], + attention_mask=batch["attention_mask"], + ) + preds = outputs.logits.argmax(dim=-1) + return {"prediction": preds.tolist()} diff --git a/projects/lora/README.md b/projects/lora/README.md new file mode 100644 index 00000000..c3d85f7b --- /dev/null +++ b/projects/lora/README.md @@ -0,0 +1,123 @@ +# LoRA (Low-Rank Adaptation) + +Parameter-efficient fine-tuning using [HuggingFace PEFT](https://github.com/huggingface/peft). + +## Overview + +LoRA (Low-Rank Adaptation) enables efficient fine-tuning by adding small trainable low-rank matrices to frozen pretrained weights. This reduces trainable parameters by 10-100x while maintaining performance comparable to full fine-tuning. + +This project uses the industry-standard **PEFT library** from HuggingFace, which provides: +- Actively maintained implementation +- Support for quantization (QLoRA) +- Built-in adapter saving, loading, and merging +- Extensive model support + +## Dataset + +**CIFAR-100** - 100 classes, 60k images. Auto-downloaded. + +## Lighter Features Demonstrated + +- **PEFT integration** - Standard LoRA via HuggingFace +- **`$` expressions** - Filter trainable parameters: `$[p for p in @model::network.parameters() if p.requires_grad]` +- **CsvWriter** - Prediction logging +- **Adapter management** - Save/load/merge adapters + +## Usage + +```bash +pip install lighter peft +cd projects/lora + +# Train +lighter fit configs/lora.yaml + +# Quick test +lighter fit configs/lora.yaml trainer::fast_dev_run=true + +# Different rank (more expressive) +lighter fit configs/lora.yaml vars::lora_rank=16 vars::lora_alpha=32 + +# Higher dropout for regularization +lighter fit configs/lora.yaml vars::lora_dropout=0.2 +``` + +## LoRA Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `lora_rank` | 8 | Low-rank dimension. Higher = more expressive but more params. Typical: 4-64 | +| `lora_alpha` | 16 | Scaling factor. Typically 2x rank. Controls adaptation strength | +| `lora_dropout` | 0.1 | Dropout for regularization. Increase if overfitting | + +## How LoRA Works + +Standard fine-tuning updates all weights W: +``` +h = Wx +``` + +LoRA adds low-rank decomposition BA: +``` +h = Wx + (BA)x * (alpha/rank) +``` + +Where: +- **W**: Frozen pretrained weights +- **B**: Trainable down-projection (out_features × rank) +- **A**: Trainable up-projection (rank × in_features) +- **alpha/rank**: Scaling factor + +Benefits: +- 10-100x fewer trainable parameters +- Reduced memory footprint +- Faster training +- Can store multiple task-specific adapters +- Merge adapters into base model for inference + +## Targeting Specific Modules + +By default, PEFT applies LoRA based on model architecture. For explicit control: + +```yaml +network: + _target_: project.networks.lora.LoRAWrapper + target_modules: ["fc"] # ResNet: FC layers only + + # For Vision Transformers: + # target_modules: ["query", "value"] + + # For Transformers: + # target_modules: ["q_proj", "v_proj", "k_proj", "o_proj"] +``` + +## Saving and Loading Adapters + +The LoRAWrapper provides methods for adapter management: + +```python +# Save adapter (small file, just LoRA weights) +model.network.save_adapter("./adapters/my_task") + +# Load adapter +model.network.load_adapter("./adapters/my_task") + +# Merge for deployment (single model file) +merged_model = model.network.merge_and_unload() +``` + +## Advanced: QLoRA + +For even more memory efficiency, use quantized LoRA (QLoRA): + +```bash +pip install bitsandbytes +``` + +Then configure quantization in your base model. See [PEFT QLoRA docs](https://huggingface.co/docs/peft/main/en/developer_guides/quantization). + +## References + +- [LoRA Paper](https://arxiv.org/abs/2106.09685) - Hu et al., 2021 +- [PEFT Library](https://github.com/huggingface/peft) - HuggingFace +- [QLoRA Paper](https://arxiv.org/abs/2305.14314) - Quantized LoRA diff --git a/projects/lora/__init__.py b/projects/lora/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/lora/__lighter__.py b/projects/lora/__lighter__.py new file mode 100644 index 00000000..9f1a3f37 --- /dev/null +++ b/projects/lora/__lighter__.py @@ -0,0 +1,8 @@ +# This file marks the directory as a Lighter project. +# Efficient Fine-tuning (LoRA) example demonstrating: +# - Parameter-efficient fine-tuning (PEFT) +# - Low-Rank Adaptation (LoRA) implementation +# - Freezer callback for freezing base model +# - Minimal trainable parameters (~0.1-1% of full model) +# - Dynamic rank configuration with $ expressions +# - _disabled_ for conditional adapter enable/disable diff --git a/projects/lora/configs/lora.yaml b/projects/lora/configs/lora.yaml new file mode 100644 index 00000000..9c6fd8e9 --- /dev/null +++ b/projects/lora/configs/lora.yaml @@ -0,0 +1,200 @@ +# LoRA (Low-Rank Adaptation) +# +# Parameter-efficient fine-tuning using HuggingFace PEFT. +# +# LoRA adds small trainable low-rank matrices to frozen pretrained weights, +# reducing trainable parameters by 10-100x while maintaining performance. +# +# Benefits: +# - 10-100x fewer trainable parameters +# - Reduced memory footprint +# - Faster training +# - Store multiple task-specific adapters +# - Merge adapters into base model for inference +# +# Requirements: +# pip install peft +# +# Run: lighter fit projects/lora/configs/lora.yaml + +vars: + # Task configuration + num_classes: 100 # CIFAR100 + image_size: 224 + batch_size: 32 + + # LoRA configuration + lora_rank: 8 # Low rank = fewer params (typically 4-64) + lora_alpha: 16 # Scaling factor (typically 2*rank) + lora_dropout: 0.1 # Dropout for regularization + + num_workers: 4 + +trainer: + _target_: pytorch_lightning.Trainer + max_epochs: 20 + accelerator: auto + log_every_n_steps: 10 + + callbacks: + # Save predictions + - _target_: lighter.callbacks.CsvWriter + path: predictions.csv + keys: [prediction, confidence, ground_truth] + + # Early stopping + - _target_: pytorch_lightning.callbacks.EarlyStopping + monitor: val/metrics/MulticlassAccuracy/epoch + mode: max + patience: 5 + + # Best model checkpoint + - _target_: pytorch_lightning.callbacks.ModelCheckpoint + monitor: val/metrics/MulticlassAccuracy/epoch + mode: max + save_top_k: 1 + filename: lora-{epoch}-{val_accuracy:.4f} + + # For quick testing, uncomment: + # fast_dev_run: true + +model: + _target_: project.models.model.LoRAClassificationModel + + # LoRA-wrapped model using HuggingFace PEFT + network: + _target_: project.networks.lora.LoRAWrapper + lora_rank: "%vars::lora_rank" + lora_alpha: "%vars::lora_alpha" + lora_dropout: "%vars::lora_dropout" + + # Target specific modules for LoRA + # For ImageClassifier: target the Linear in classifier Sequential (index 2) + # For ViT: target attention projections ["query", "value"] + target_modules: ["classifier.2"] + + # Base model to adapt + base_model: + _target_: project.networks.network.ImageClassifier + backbone: resnet50 + num_classes: "%vars::num_classes" + pretrained: true + freeze_backbone: false # LoRA handles freezing + + criterion: + _target_: torch.nn.CrossEntropyLoss + label_smoothing: 0.1 + + # Only LoRA parameters will be trainable + optimizer: + _target_: torch.optim.AdamW + params: "$[p for p in @model::network.parameters() if p.requires_grad]" + lr: 0.001 # Can use higher LR for LoRA + weight_decay: 0.01 + + scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + optimizer: "@model::optimizer" + T_max: 20 + + # Metrics + train_metrics: + _target_: torchmetrics.MetricCollection + metrics: + - _target_: torchmetrics.classification.MulticlassAccuracy + num_classes: "%vars::num_classes" + average: macro + - _target_: torchmetrics.classification.MulticlassF1Score + num_classes: "%vars::num_classes" + average: macro + + val_metrics: "%model::train_metrics" + test_metrics: "%model::train_metrics" + +data: + _target_: lighter.LighterDataModule + + train_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: true + dataset: + _target_: project.dataset.CIFAR100Dataset + root: ./.datasets/ + train: true + download: true + transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: ["%vars::image_size", "%vars::image_size"] + - _target_: torchvision.transforms.RandomHorizontalFlip + - _target_: torchvision.transforms.RandomRotation + degrees: 15 + - _target_: torchvision.transforms.ToTensor + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + val_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: false + dataset: + _target_: project.dataset.CIFAR100Dataset + root: ./.datasets/ + train: false + download: true + transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: ["%vars::image_size", "%vars::image_size"] + - _target_: torchvision.transforms.ToTensor + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + test_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + shuffle: false + dataset: + _target_: project.dataset.CIFAR100Dataset + root: ./.datasets/ + train: false + download: true + transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: ["%vars::image_size", "%vars::image_size"] + - _target_: torchvision.transforms.ToTensor + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + predict_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: 16 + num_workers: "%vars::num_workers" + shuffle: false + dataset: + _target_: project.dataset.CIFAR100Dataset + root: ./.datasets/ + train: false + download: true + transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: ["%vars::image_size", "%vars::image_size"] + - _target_: torchvision.transforms.ToTensor + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] diff --git a/projects/lora/dataset.py b/projects/lora/dataset.py new file mode 100644 index 00000000..0931434a --- /dev/null +++ b/projects/lora/dataset.py @@ -0,0 +1,104 @@ +"""Datasets for efficient fine-tuning examples.""" + +from collections.abc import Callable + +import torch +from torch.utils.data import Dataset +from torchvision import datasets + + +class CIFAR100Dataset(Dataset): + """CIFAR100 dataset wrapper for fine-tuning. + + CIFAR100 is useful for demonstrating fine-tuning because: + - 100 classes provide a challenging classification task + - Images are small enough for quick experimentation + - Pretrained models on ImageNet transfer well + + Args: + root: Root directory for dataset. + train: Whether to use training set. + transform: Transform to apply to images. + download: Whether to download if not present. + """ + + def __init__( + self, + root: str = "./.datasets/", + train: bool = True, + transform: Callable | None = None, + download: bool = True, + ) -> None: + self.dataset = datasets.CIFAR100( + root=root, + train=train, + download=download, + ) + self.transform = transform + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]: + image, label = self.dataset[idx] + + if self.transform: + image = self.transform(image) + + return image, label + + +class FewShotDataset(Dataset): + """Few-shot learning dataset for fine-tuning evaluation. + + Creates a dataset with limited samples per class, which is + the typical scenario where LoRA shines. + + Args: + base_dataset: The full dataset to sample from. + samples_per_class: Number of samples per class (k-shot). + num_classes: Number of classes to include (n-way). + seed: Random seed for reproducibility. + """ + + def __init__( + self, + base_dataset: Dataset, + samples_per_class: int = 5, + num_classes: int = 10, + seed: int = 42, + ) -> None: + self.samples_per_class = samples_per_class + self.num_classes = num_classes + + # Group samples by class + class_indices: dict[int, list[int]] = {} + for idx in range(len(base_dataset)): + _, label = base_dataset[idx] + if label not in class_indices: + class_indices[label] = [] + class_indices[label].append(idx) + + # Sample k examples from n classes + generator = torch.Generator().manual_seed(seed) + selected_classes = list(class_indices.keys())[:num_classes] + + self.indices = [] + self.label_map = {old: new for new, old in enumerate(selected_classes)} + + for cls in selected_classes: + cls_indices = class_indices[cls] + perm = torch.randperm(len(cls_indices), generator=generator) + selected = perm[:samples_per_class].tolist() + self.indices.extend([cls_indices[i] for i in selected]) + + self.base_dataset = base_dataset + + def __len__(self) -> int: + return len(self.indices) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]: + real_idx = self.indices[idx] + image, old_label = self.base_dataset[real_idx] + new_label = self.label_map[old_label] + return image, new_label diff --git a/projects/lora/models/__init__.py b/projects/lora/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/lora/models/model.py b/projects/lora/models/model.py new file mode 100644 index 00000000..d25c429a --- /dev/null +++ b/projects/lora/models/model.py @@ -0,0 +1,80 @@ +"""Efficient fine-tuning model using LighterModule.""" + +import torch + +from lighter import LighterModule + + +class LoRAClassificationModel(LighterModule): + """Classification model with LoRA fine-tuning support. + + Showcases: + - Parameter-efficient fine-tuning + - Tracking trainable parameter percentage + - Standard classification workflow + """ + + def _shared_step(self, batch: tuple, metrics) -> dict: + """Shared logic for train/val/test steps.""" + image, label = batch + + # Forward pass + logits = self(image) + + # Compute loss + loss = self.criterion(logits, label) + + # Get predictions + pred = logits.argmax(dim=1) + + # Update metrics + if metrics is not None: + metrics(pred, label) + + return { + "loss": loss, + "pred": pred, + "label": label, + "logits": logits, + } + + def training_step(self, batch, batch_idx): + return self._shared_step(batch, self.train_metrics) + + def validation_step(self, batch, batch_idx): + return self._shared_step(batch, self.val_metrics) + + def test_step(self, batch, batch_idx): + result = self._shared_step(batch, self.test_metrics) + del result["loss"] + return result + + def predict_step(self, batch, batch_idx): + """Prediction step.""" + image, label = batch if isinstance(batch, (tuple, list)) else (batch, None) + + logits = self(image) + pred = logits.argmax(dim=1) + confidence = torch.softmax(logits, dim=1).max(dim=1).values + + result = { + "prediction": pred.tolist(), + "confidence": confidence.tolist(), + } + + if label is not None: + result["ground_truth"] = label.tolist() + + return result + + def on_train_start(self) -> None: + """Log trainable parameters at training start.""" + trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) + total = sum(p.numel() for p in self.parameters()) + pct = 100 * trainable / total if total > 0 else 0 + + self.log("trainable_params", float(trainable)) + self.log("total_params", float(total)) + self.log("trainable_pct", pct) + + print(f"\nTrainable parameters: {trainable:,} / {total:,} ({pct:.2f}%)") diff --git a/projects/lora/networks/__init__.py b/projects/lora/networks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/lora/networks/lora.py b/projects/lora/networks/lora.py new file mode 100644 index 00000000..4cffd859 --- /dev/null +++ b/projects/lora/networks/lora.py @@ -0,0 +1,165 @@ +"""LoRA (Low-Rank Adaptation) using HuggingFace PEFT. + +This module provides integration with the standard PEFT library for parameter-efficient +fine-tuning. PEFT supports LoRA, QLoRA, Prefix Tuning, Prompt Tuning, and more. + +LoRA enables efficient fine-tuning by adding trainable low-rank matrices to frozen +pretrained weights, reducing trainable parameters by 10-100x while maintaining +performance comparable to full fine-tuning. + +Reference: + Hu et al. (2021) "LoRA: Low-Rank Adaptation of Large Language Models" + https://arxiv.org/abs/2106.09685 + +Requirements: + pip install peft +""" + +from typing import Any + +import torch.nn as nn + + +class LoRAWrapper(nn.Module): + """Wrapper that applies LoRA to a base model using HuggingFace PEFT. + + This wrapper uses the industry-standard PEFT library for parameter-efficient + fine-tuning. Benefits include: + - Actively maintained by HuggingFace + - Supports quantization (QLoRA) for memory efficiency + - Built-in adapter saving, loading, and merging + - Extensive model and method support + + Args: + base_model: The pretrained model to adapt. + lora_rank: Rank of the low-rank decomposition. Higher = more expressive + but more parameters. Typical values: 4-64. Default: 8. + lora_alpha: Scaling factor. The adaptation is scaled by alpha/rank. + Typically 2x the rank. Default: 16. + lora_dropout: Dropout probability for LoRA layers. Default: 0.1. + target_modules: Which modules to apply LoRA to. If None, PEFT uses + sensible defaults based on the model architecture. + modules_to_save: Modules to save in addition to LoRA adapters (e.g., classifier head). + task_type: Task type for PEFT config. Options: 'SEQ_CLS', 'SEQ_2_SEQ_LM', + 'CAUSAL_LM', 'TOKEN_CLS', 'QUESTION_ANS', 'FEATURE_EXTRACTION'. + bias: Bias training strategy. Options: 'none', 'all', 'lora_only'. Default: 'none'. + **peft_kwargs: Additional arguments passed to LoraConfig. + + Example: + ```yaml + model: + network: + _target_: project.networks.lora.LoRAWrapper + lora_rank: 8 + lora_alpha: 16 + lora_dropout: 0.1 + target_modules: ["query", "value"] + base_model: + _target_: torchvision.models.resnet50 + weights: IMAGENET1K_V2 + ``` + + Note: + - For vision models, target_modules might need to be set explicitly + - Use modules_to_save for classifier heads that need full fine-tuning + - QLoRA requires additional setup (BitsAndBytesConfig) + """ + + def __init__( + self, + base_model: nn.Module, + lora_rank: int = 8, + lora_alpha: int = 16, + lora_dropout: float = 0.1, + target_modules: list[str] | None = None, + modules_to_save: list[str] | None = None, + task_type: str | None = None, + bias: str = "none", + **peft_kwargs: Any, + ) -> None: + super().__init__() + + try: + from peft import LoraConfig, TaskType, get_peft_model + except ImportError as e: + raise ImportError( + "PEFT library required for LoRA. Install with:\n" + " pip install peft\n\n" + "For QLoRA (quantized), also install:\n" + " pip install bitsandbytes" + ) from e + + # Map task_type string to TaskType enum if provided + peft_task_type = None + if task_type: + peft_task_type = getattr(TaskType, task_type.upper(), None) + if peft_task_type is None: + valid_types = [t.name for t in TaskType] + raise ValueError(f"Invalid task_type '{task_type}'. Valid options: {valid_types}") + + # Create LoRA config + config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + target_modules=target_modules, + modules_to_save=modules_to_save, + task_type=peft_task_type, + bias=bias, + **peft_kwargs, + ) + + # Apply PEFT + self.model = get_peft_model(base_model, config) + + # Print trainable parameters summary + self.model.print_trainable_parameters() + + def forward(self, *args, **kwargs): + """Forward through LoRA-adapted model.""" + return self.model(*args, **kwargs) + + def save_adapter(self, path: str) -> None: + """Save LoRA adapter weights to disk. + + Only saves the adapter weights (small), not the full model. + + Args: + path: Directory to save adapter weights. + """ + self.model.save_pretrained(path) + print(f"LoRA adapter saved to: {path}") + + def load_adapter(self, path: str, adapter_name: str = "default") -> None: + """Load a LoRA adapter from disk. + + Args: + path: Directory containing adapter weights. + adapter_name: Name for the adapter. Default: "default". + """ + self.model.load_adapter(path, adapter_name) + print(f"LoRA adapter loaded from: {path}") + + def merge_and_unload(self) -> nn.Module: + """Merge LoRA weights into base model for efficient inference. + + After merging, the model no longer has separate adapter weights. + This is useful for deployment where you want a single model file. + + Returns: + Base model with LoRA weights merged in. + """ + return self.model.merge_and_unload() + + def get_nb_trainable_parameters(self) -> tuple[int, int, float]: + """Get trainable parameter statistics. + + Returns: + Tuple of (trainable_params, all_params, percentage). + """ + return self.model.get_nb_trainable_parameters() + + @property + def base_model(self) -> nn.Module: + """Access the underlying base model.""" + return self.model.get_base_model() diff --git a/projects/lora/networks/network.py b/projects/lora/networks/network.py new file mode 100644 index 00000000..a02fd13e --- /dev/null +++ b/projects/lora/networks/network.py @@ -0,0 +1,95 @@ +"""Base networks for efficient fine-tuning examples.""" + +import torch +import torch.nn as nn +from torchvision import models + + +class ImageClassifier(nn.Module): + """Image classifier using pretrained backbone with custom head. + + This model is designed to be adapted with LoRA - the backbone is + pretrained and the classification head matches the target task. + + Args: + backbone: Name of torchvision model. + num_classes: Number of output classes. + pretrained: Use pretrained weights. + freeze_backbone: Whether to freeze backbone initially. + """ + + def __init__( + self, + backbone: str = "resnet50", + num_classes: int = 100, + pretrained: bool = True, + freeze_backbone: bool = False, + ) -> None: + super().__init__() + + # Load backbone + if backbone == "resnet50": + weights = models.ResNet50_Weights.IMAGENET1K_V1 if pretrained else None + base = models.resnet50(weights=weights) + self.backbone = nn.Sequential(*list(base.children())[:-1]) + backbone_dim = 2048 + elif backbone == "resnet18": + weights = models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None + base = models.resnet18(weights=weights) + self.backbone = nn.Sequential(*list(base.children())[:-1]) + backbone_dim = 512 + elif backbone == "vit_b_16": + weights = models.ViT_B_16_Weights.IMAGENET1K_V1 if pretrained else None + self.backbone = models.vit_b_16(weights=weights) + self.backbone.heads = nn.Identity() + backbone_dim = 768 + else: + raise ValueError(f"Unsupported backbone: {backbone}") + + # Freeze backbone if requested + if freeze_backbone: + for param in self.backbone.parameters(): + param.requires_grad = False + + # Classification head + self.classifier = nn.Sequential( + nn.Flatten(), + nn.Dropout(0.2), + nn.Linear(backbone_dim, num_classes), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + features = self.backbone(x) + return self.classifier(features) + + +class VisionTransformerForClassification(nn.Module): + """Vision Transformer optimized for LoRA fine-tuning. + + ViT is particularly well-suited for LoRA because: + - Attention layers have many linear projections + - MLP blocks are linear layers + - LoRA can target q_proj, k_proj, v_proj, out_proj + + Args: + num_classes: Number of output classes. + pretrained: Use pretrained weights. + img_size: Input image size. + """ + + def __init__( + self, + num_classes: int = 100, + pretrained: bool = True, + img_size: int = 224, + ) -> None: + super().__init__() + + weights = models.ViT_B_16_Weights.IMAGENET1K_V1 if pretrained else None + self.vit = models.vit_b_16(weights=weights, image_size=img_size) + + # Replace classification head + self.vit.heads = nn.Linear(768, num_classes) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.vit(x) diff --git a/projects/medical_segmentation/README.md b/projects/medical_segmentation/README.md new file mode 100644 index 00000000..aa8f9423 --- /dev/null +++ b/projects/medical_segmentation/README.md @@ -0,0 +1,42 @@ +# Medical Image Segmentation with MONAI + +3D CT spleen segmentation using MONAI. + +## Dataset + +**Medical Segmentation Decathlon - Task09_Spleen** +- 41 training + 20 validation CT volumes +- Auto-downloaded (~1.5 GB) + +## Architecture + +**MONAI UNet** - 3D U-Net with residual connections and instance normalization. + +## Lighter Features Demonstrated + +- **MONAI integration** - DecathlonDataset, transforms, UNet +- **Sliding window inference** - process large 3D volumes in patches +- **Mixed precision** - 16-bit training +- **FileWriter callback** - saves segmentation masks as .seg.nrrd +- **`_mode_: callable`** - for custom writer function + +## Usage + +```bash +pip install lighter monai itk +cd projects/medical_segmentation + +# Train +lighter fit configs/spleen.yaml + +# Quick test +lighter fit configs/spleen.yaml trainer::fast_dev_run=true + +# Multi-GPU (recommended for 3D) +lighter fit configs/spleen.yaml trainer::devices=2 trainer::strategy=ddp +``` + +## References + +- [MONAI Documentation](https://docs.monai.io/) +- [Medical Segmentation Decathlon](http://medicaldecathlon.com/) diff --git a/projects/medical_segmentation/__init__.py b/projects/medical_segmentation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/medical_segmentation/__lighter__.py b/projects/medical_segmentation/__lighter__.py new file mode 100644 index 00000000..ad7cdae9 --- /dev/null +++ b/projects/medical_segmentation/__lighter__.py @@ -0,0 +1,6 @@ +# This file marks the directory as a Lighter project. +# Medical Image Segmentation example demonstrating: +# - FileWriter callback for saving segmentation masks +# - Freezer callback for transfer learning +# - Differential learning rates (frozen encoder, trainable decoder) +# - MetricCollection with Dice and IoU metrics diff --git a/projects/medical_segmentation/configs/spleen.yaml b/projects/medical_segmentation/configs/spleen.yaml new file mode 100644 index 00000000..2301f7e4 --- /dev/null +++ b/projects/medical_segmentation/configs/spleen.yaml @@ -0,0 +1,227 @@ +# Medical Image Segmentation with MONAI +# +# 3D Spleen CT Segmentation using the Medical Segmentation Decathlon dataset. +# Uses MONAI (https://github.com/Project-MONAI/MONAI) for medical imaging. +# +# Dataset: Task09_Spleen from Medical Segmentation Decathlon +# - 41 training CT scans +# - 20 validation CT scans +# - Auto-downloaded on first run (~1.5 GB) +# +# Features demonstrated: +# - FileWriter callback for saving segmentation masks as .seg.nrrd +# - Custom writer function for medical imaging formats +# - MONAI integration with sliding window inference +# +# Install: +# pip install lighter monai itk +# +# Run: +# lighter fit projects/medical_segmentation/configs/spleen.yaml +# lighter predict projects/medical_segmentation/configs/spleen.yaml # saves .seg.nrrd files + +vars: + num_classes: 2 # Background + Spleen + roi_size: [96, 96, 96] + batch_size: 2 + sw_batch_size: 4 + num_workers: 4 + +trainer: + _target_: pytorch_lightning.Trainer + max_epochs: 100 + accelerator: auto + log_every_n_steps: 5 + precision: 16-mixed # Use mixed precision for 3D volumes + + callbacks: + - _target_: pytorch_lightning.callbacks.EarlyStopping + monitor: val/MulticlassJaccardIndex_epoch + mode: max + patience: 20 + + - _target_: pytorch_lightning.callbacks.ModelCheckpoint + monitor: val/MulticlassJaccardIndex_epoch + mode: max + save_top_k: 1 + filename: spleen-{epoch}-{val/MulticlassJaccardIndex_epoch:.4f} + + # FileWriter saves segmentation masks as .seg.nrrd files + - _target_: lighter.callbacks.FileWriter + directory: ./predictions + value_key: pred + writer_fn: + _target_: project.writers.write_nrrd + _mode_: callable + + # For quick testing, uncomment: + # fast_dev_run: true + +model: + _target_: project.models.model.SegmentationModel + sw_batch_size: "%vars::sw_batch_size" + roi_size: "%vars::roi_size" + + network: + _target_: monai.networks.nets.UNet + spatial_dims: 3 + in_channels: 1 + out_channels: "%vars::num_classes" + channels: [16, 32, 64, 128, 256] + strides: [2, 2, 2, 2] + num_res_units: 2 + dropout: 0.2 + + optimizer: + _target_: torch.optim.AdamW + params: $@model::network.parameters() + lr: 0.0001 + weight_decay: 0.00001 + + scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + optimizer: "@model::optimizer" + T_max: 100 + + train_metrics: + _target_: torchmetrics.MetricCollection + metrics: + - _target_: torchmetrics.classification.MulticlassJaccardIndex + num_classes: "%vars::num_classes" + average: macro + + val_metrics: "%model::train_metrics" + test_metrics: "%model::train_metrics" + +data: + _target_: lighter.LighterDataModule + + train_dataloader: + _target_: monai.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: true + dataset: + _target_: monai.apps.DecathlonDataset + root_dir: ./.datasets/ + task: Task09_Spleen + section: training + download: true + transform: + _target_: monai.transforms.Compose + transforms: + - _target_: monai.transforms.LoadImaged + keys: [image, label] + - _target_: monai.transforms.EnsureChannelFirstd + keys: [image, label] + - _target_: monai.transforms.Orientationd + keys: [image, label] + axcodes: RAS + - _target_: monai.transforms.Spacingd + keys: [image, label] + pixdim: [1.5, 1.5, 2.0] + mode: [bilinear, nearest] + - _target_: monai.transforms.ScaleIntensityRanged + keys: [image] + a_min: -57 + a_max: 164 + b_min: 0.0 + b_max: 1.0 + clip: true + - _target_: monai.transforms.RandCropByPosNegLabeld + keys: [image, label] + label_key: label + spatial_size: "%vars::roi_size" + pos: 1 + neg: 1 + num_samples: 4 + - _target_: monai.transforms.RandFlipd + keys: [image, label] + prob: 0.5 + spatial_axis: 0 + - _target_: monai.transforms.RandFlipd + keys: [image, label] + prob: 0.5 + spatial_axis: 1 + - _target_: monai.transforms.RandFlipd + keys: [image, label] + prob: 0.5 + spatial_axis: 2 + - _target_: monai.transforms.RandRotate90d + keys: [image, label] + prob: 0.5 + max_k: 3 + - _target_: monai.transforms.EnsureTyped + keys: [image, label] + + val_dataloader: + _target_: monai.data.DataLoader + batch_size: 1 # Full volume for validation + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: false + dataset: + _target_: monai.apps.DecathlonDataset + root_dir: ./.datasets/ + task: Task09_Spleen + section: validation + download: true + transform: + _target_: monai.transforms.Compose + transforms: + - _target_: monai.transforms.LoadImaged + keys: [image, label] + - _target_: monai.transforms.EnsureChannelFirstd + keys: [image, label] + - _target_: monai.transforms.Orientationd + keys: [image, label] + axcodes: RAS + - _target_: monai.transforms.Spacingd + keys: [image, label] + pixdim: [1.5, 1.5, 2.0] + mode: [bilinear, nearest] + - _target_: monai.transforms.ScaleIntensityRanged + keys: [image] + a_min: -57 + a_max: 164 + b_min: 0.0 + b_max: 1.0 + clip: true + - _target_: monai.transforms.EnsureTyped + keys: [image, label] + + test_dataloader: + _target_: monai.data.DataLoader + batch_size: 1 + num_workers: "%vars::num_workers" + shuffle: false + dataset: + _target_: monai.apps.DecathlonDataset + root_dir: ./.datasets/ + task: Task09_Spleen + section: validation + download: true + transform: + _target_: monai.transforms.Compose + transforms: + - _target_: monai.transforms.LoadImaged + keys: [image, label] + - _target_: monai.transforms.EnsureChannelFirstd + keys: [image, label] + - _target_: monai.transforms.Orientationd + keys: [image, label] + axcodes: RAS + - _target_: monai.transforms.Spacingd + keys: [image, label] + pixdim: [1.5, 1.5, 2.0] + mode: [bilinear, nearest] + - _target_: monai.transforms.ScaleIntensityRanged + keys: [image] + a_min: -57 + a_max: 164 + b_min: 0.0 + b_max: 1.0 + clip: true + - _target_: monai.transforms.EnsureTyped + keys: [image, label] diff --git a/projects/medical_segmentation/models/__init__.py b/projects/medical_segmentation/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/medical_segmentation/models/model.py b/projects/medical_segmentation/models/model.py new file mode 100644 index 00000000..0ada6893 --- /dev/null +++ b/projects/medical_segmentation/models/model.py @@ -0,0 +1,120 @@ +"""Medical segmentation model using MONAI and LighterModule.""" + +import torch + +from lighter import LighterModule + + +class SegmentationModel(LighterModule): + """3D Medical image segmentation model using MONAI. + + Uses MONAI's DiceLoss and sliding window inference for volumetric data. + """ + + def __init__( + self, + network: torch.nn.Module, + optimizer: torch.optim.Optimizer | None = None, + scheduler=None, + sw_batch_size: int = 4, + roi_size: tuple[int, int, int] = (96, 96, 96), + **kwargs, + ) -> None: + super().__init__( + network=network, + optimizer=optimizer, + scheduler=scheduler, + **kwargs, + ) + self.sw_batch_size = sw_batch_size + self.roi_size = roi_size + + # Use MONAI's DiceLoss + from monai.losses import DiceLoss + + self.dice_loss = DiceLoss(to_onehot_y=True, softmax=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass.""" + return self.network(x) + + def training_step(self, batch, batch_idx): + """Training step with Dice loss.""" + images = batch["image"] + labels = batch["label"] + + # Forward pass + outputs = self(images) + + # Compute Dice loss + loss = self.dice_loss(outputs, labels) + + # Get predictions for metrics + preds = outputs.argmax(dim=1) + + if self.train_metrics is not None: + self.train_metrics(preds, labels.squeeze(1).long()) + + self.log("train/loss", loss, prog_bar=True) + return {"loss": loss} + + def validation_step(self, batch, batch_idx): + """Validation step with sliding window inference.""" + from monai.inferers import sliding_window_inference + + images = batch["image"] + labels = batch["label"] + + # Use sliding window inference for full volumes + outputs = sliding_window_inference( + images, + roi_size=self.roi_size, + sw_batch_size=self.sw_batch_size, + predictor=self.network, + ) + + loss = self.dice_loss(outputs, labels) + preds = outputs.argmax(dim=1) + + if self.val_metrics is not None: + self.val_metrics(preds, labels.squeeze(1).long()) + + self.log("val/loss", loss, prog_bar=True) + return {"loss": loss} + + def test_step(self, batch, batch_idx): + """Test step with sliding window inference.""" + from monai.inferers import sliding_window_inference + + images = batch["image"] + labels = batch["label"] + + outputs = sliding_window_inference( + images, + roi_size=self.roi_size, + sw_batch_size=self.sw_batch_size, + predictor=self.network, + ) + + preds = outputs.argmax(dim=1) + + if self.test_metrics is not None: + self.test_metrics(preds, labels.squeeze(1).long()) + + return {"pred": preds, "label": labels} + + def predict_step(self, batch, batch_idx): + """Prediction step returning segmentation masks.""" + from monai.inferers import sliding_window_inference + + images = batch["image"] + + outputs = sliding_window_inference( + images, + roi_size=self.roi_size, + sw_batch_size=self.sw_batch_size, + predictor=self.network, + ) + + preds = outputs.argmax(dim=1) + return {"pred": preds, "output": outputs} diff --git a/projects/medical_segmentation/writers.py b/projects/medical_segmentation/writers.py new file mode 100644 index 00000000..2879724b --- /dev/null +++ b/projects/medical_segmentation/writers.py @@ -0,0 +1,60 @@ +"""Custom writer functions for medical imaging outputs.""" + +from collections.abc import Sequence +from pathlib import Path + +import numpy as np +import torch + + +def write_nrrd( + path: Path, + tensor: torch.Tensor, + *, + suffix: str = ".seg.nrrd", + affine: np.ndarray | None = None, + spatial_shape: Sequence[int] | None = None, +) -> None: + """Write a 3D segmentation tensor to NRRD format. + + Uses MONAI's ITKWriter for robust I/O with proper orientation handling. + NRRD (Nearly Raw Raster Data) is commonly used for medical segmentation masks. + + Args: + path: Output path (suffix will be replaced). + tensor: 3D tensor (D, H, W) or (C, D, H, W) containing segmentation labels. + suffix: File suffix (default: .seg.nrrd for segmentation convention). + affine: 4x4 affine matrix for spatial orientation. If None, uses identity. + spatial_shape: Original spatial shape for resampling. If None, uses tensor shape. + """ + from monai.data.image_writer import ITKWriter + + path = path.with_suffix(suffix) + + # Convert tensor to numpy + if isinstance(tensor, torch.Tensor): + data = tensor.detach().cpu().numpy() + else: + data = np.asarray(tensor) + + # Remove batch dim if present: (B, C, D, H, W) -> (C, D, H, W) or (D, H, W) + if data.ndim == 5: + if data.shape[0] != 1: + raise ValueError(f"Expected batch size 1, got {data.shape[0]}. Use batch_size=1 for prediction.") + data = data[0] + + # Validate final shape is 3D or 4D + if data.ndim not in (3, 4): + raise ValueError(f"Expected 3D (D,H,W) or 4D (C,D,H,W) data after removing batch, got shape {data.shape}") + + # Determine channel_dim: if 4D assume first dim is channel, else no channel + channel_dim = 0 if data.ndim == 4 else None + + # Use identity affine if not provided + if affine is None: + affine = np.eye(4) + + writer = ITKWriter(output_dtype=np.int16) + writer.set_data_array(data, channel_dim=channel_dim) + writer.set_metadata({"affine": affine, "spatial_shape": spatial_shape or data.shape[-3:]}) + writer.write(str(path)) diff --git a/projects/self_supervised/README.md b/projects/self_supervised/README.md new file mode 100644 index 00000000..0fca46ff --- /dev/null +++ b/projects/self_supervised/README.md @@ -0,0 +1,111 @@ +# Self-Supervised Learning with SimCLR + +Contrastive learning using the [lightly](https://github.com/lightly-ai/lightly) library. + +## Overview + +Self-supervised learning (SSL) learns representations from unlabeled data. This project implements **SimCLR** (Simple Framework for Contrastive Learning), which learns by maximizing agreement between differently augmented views of the same image. + +## How SimCLR Works + +1. **Two views**: Each image is augmented twice to create two different views +2. **Encoding**: Both views pass through the same encoder network +3. **Projection**: Features are projected to a lower-dimensional space +4. **Contrastive loss**: NT-Xent loss pulls together views of the same image while pushing apart views of different images + +## Dataset + +**CIFAR-10** - 60,000 32x32 images. Auto-downloaded. + +The lightly library applies SimCLR-specific augmentations: +- Random resized crop +- Color jitter (brightness, contrast, saturation, hue) +- Random grayscale +- Gaussian blur +- Random horizontal flip + +## Architecture + +- **Backbone**: ResNet-18 (configurable to ResNet-50) +- **Projection head**: MLP with hidden layer (2048) → output (128) +- **Loss**: NT-Xent (Normalized Temperature-scaled Cross Entropy) + +## Lighter Features Demonstrated + +- **lightly integration** - SSL transforms, losses, and dataset wrappers +- **Model-computed loss** - No `criterion` config needed (loss computed internally) +- **Custom `__init__`** - Extending LighterModule with additional parameters +- **Temperature parameter** - Configurable via `vars::` section +- **`drop_last: true`** - Required for contrastive learning (batch size consistency) + +## Usage + +```bash +pip install lighter lightly +cd projects/self_supervised + +# Train (100 epochs) +lighter fit configs/simclr.yaml + +# Quick test +lighter fit configs/simclr.yaml trainer::fast_dev_run=true + +# Different backbone +lighter fit configs/simclr.yaml model::network::backbone=resnet50 + +# Adjust temperature +lighter fit configs/simclr.yaml vars::temperature=0.1 + +# Larger batch size (important for SSL performance) +lighter fit configs/simclr.yaml vars::batch_size=512 +``` + +## Configuration Highlights + +```yaml +vars: + batch_size: 256 # Larger is better for contrastive learning + temperature: 0.5 # Lower = harder negatives + projection_dim: 128 # Output dimension of projection head + +model: + _target_: project.models.model.SimCLRModel + temperature: "%vars::temperature" + # No criterion - loss is computed internally + +data: + train_dataloader: + drop_last: true # Required for contrastive learning +``` + +## Downstream Evaluation + +After pretraining, extract features for downstream tasks: + +```bash +lighter predict configs/simclr.yaml --ckpt_path path/to/checkpoint.ckpt +``` + +The `predict_step` returns features (not projections) suitable for: +- **Linear probing**: Train a linear classifier on frozen features +- **Fine-tuning**: Use pretrained backbone with task-specific head + +## Key Hyperparameters + +| Parameter | Default | Notes | +|-----------|---------|-------| +| `batch_size` | 256 | Larger batches improve performance significantly | +| `temperature` | 0.5 | Lower values focus on hard negatives | +| `projection_dim` | 128 | Output dimension of projection head | +| `hidden_dim` | 2048 | Hidden layer size in projection head | +| Learning rate | 0.06 | Scales with batch size (0.3 * batch_size / 256) | + +## Alternative Methods + +The `networks/encoder.py` also includes a BYOL (Bootstrap Your Own Latent) implementation which doesn't require negative pairs. To use BYOL, you would need to create a corresponding model class. + +## References + +- [SimCLR Paper](https://arxiv.org/abs/2002.05709) - Chen et al., 2020 +- [lightly Documentation](https://docs.lightly.ai/) +- [SimCLR v2 Paper](https://arxiv.org/abs/2006.10029) - Improved version with larger models diff --git a/projects/self_supervised/__init__.py b/projects/self_supervised/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/self_supervised/__lighter__.py b/projects/self_supervised/__lighter__.py new file mode 100644 index 00000000..1183950a --- /dev/null +++ b/projects/self_supervised/__lighter__.py @@ -0,0 +1,7 @@ +# This file marks the directory as a Lighter project. +# Self-Supervised Learning (SimCLR) example demonstrating: +# - Contrastive learning with InfoNCE loss +# - _mode_: callable for lazy augmentation pipelines +# - Custom loss computed in step (no criterion) +# - Temperature parameter scheduling +# - Dual-view data augmentation diff --git a/projects/self_supervised/configs/simclr.yaml b/projects/self_supervised/configs/simclr.yaml new file mode 100644 index 00000000..2850f725 --- /dev/null +++ b/projects/self_supervised/configs/simclr.yaml @@ -0,0 +1,99 @@ +# Self-Supervised Learning with SimCLR using Lightly +# +# This example uses the lightly library (https://github.com/lightly-ai/lightly) +# for state-of-the-art self-supervised learning methods. +# +# Dataset: CIFAR10 (50k training images, 10k test images) +# Method: SimCLR - Contrastive learning with augmented views +# +# Install: +# pip install lighter lightly +# +# Run: +# lighter fit projects/self_supervised/configs/simclr.yaml + +vars: + batch_size: 256 + temperature: 0.5 + input_size: 32 + projection_dim: 128 + hidden_dim: 2048 + num_workers: 4 + +trainer: + _target_: pytorch_lightning.Trainer + max_epochs: 100 + accelerator: auto + log_every_n_steps: 10 + + callbacks: + - _target_: pytorch_lightning.callbacks.ModelCheckpoint + monitor: val/loss + mode: min + save_top_k: 1 + filename: simclr-{epoch}-{val/loss:.4f} + + # For quick testing, uncomment: + # fast_dev_run: true + +model: + _target_: project.models.model.SimCLRModel + temperature: "%vars::temperature" + + network: + _target_: project.networks.encoder.create_simclr_model + backbone: resnet18 + projection_dim: "%vars::projection_dim" + hidden_dim: "%vars::hidden_dim" + pretrained: false + + optimizer: + _target_: torch.optim.SGD + params: $@model::network.parameters() + lr: 0.06 + momentum: 0.9 + weight_decay: 0.0005 + + scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + optimizer: "@model::optimizer" + T_max: 100 + +data: + _target_: lighter.LighterDataModule + + train_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: true + drop_last: true + dataset: + _target_: lightly.data.LightlyDataset.from_torch_dataset + dataset: + _target_: torchvision.datasets.CIFAR10 + root: ./.datasets/ + train: true + download: true + transform: + _target_: lightly.transforms.SimCLRTransform + input_size: "%vars::input_size" + + val_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: false + drop_last: true + dataset: + _target_: lightly.data.LightlyDataset.from_torch_dataset + dataset: + _target_: torchvision.datasets.CIFAR10 + root: ./.datasets/ + train: false + download: true + transform: + _target_: lightly.transforms.SimCLRTransform + input_size: "%vars::input_size" diff --git a/projects/self_supervised/models/__init__.py b/projects/self_supervised/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/self_supervised/models/model.py b/projects/self_supervised/models/model.py new file mode 100644 index 00000000..eeedbb2b --- /dev/null +++ b/projects/self_supervised/models/model.py @@ -0,0 +1,176 @@ +"""SimCLR model using lightly library with LighterModule. + +This module implements the SimCLR (Simple Framework for Contrastive Learning of +Visual Representations) self-supervised learning method using the lightly library. + +SimCLR learns representations by maximizing agreement between differently augmented +views of the same image via a contrastive loss (NT-Xent). + +Reference: + Chen et al., "A Simple Framework for Contrastive Learning of Visual Representations" + https://arxiv.org/abs/2002.05709 + +Requirements: + pip install lightly +""" + +from typing import Any + +import torch + +from lighter import LighterModule + + +class SimCLRModel(LighterModule): + """SimCLR self-supervised learning model using lightly. + + This model implements contrastive learning where two augmented views of + the same image are pulled together in the embedding space while views + from different images are pushed apart. + + The model uses NT-Xent (Normalized Temperature-scaled Cross Entropy) loss + from the lightly library, which is an efficient implementation of the + InfoNCE loss used in SimCLR. + + Args: + network: Encoder network that returns (features, projections). + The network should output a tuple where: + - features: Representations for downstream tasks + - projections: Embeddings used for contrastive loss + optimizer: Optimizer for training. Required for training. + scheduler: Learning rate scheduler. Optional. + temperature: Temperature parameter for NT-Xent loss. Lower values + make the model focus more on hard negatives. Default: 0.5. + **kwargs: Additional arguments passed to LighterModule. + + Example: + ```yaml + model: + _target_: project.models.SimCLRModel + temperature: 0.5 + network: + _target_: project.networks.encoder.create_simclr_model + backbone: resnet18 + optimizer: + _target_: torch.optim.SGD + params: $@model::network.parameters() + lr: 0.06 + ``` + + Note: + - Use `drop_last: true` in dataloaders for contrastive learning + - Batch size significantly impacts performance (larger is better) + - The network must return (features, projections) tuple + """ + + def __init__( + self, + network: torch.nn.Module, + optimizer: torch.optim.Optimizer | None = None, + scheduler: Any = None, + temperature: float = 0.5, + **kwargs: Any, + ) -> None: + super().__init__( + network=network, + optimizer=optimizer, + scheduler=scheduler, + **kwargs, + ) + # Import here to allow module to load even if lightly is not installed, + # providing a clear error message when the model is actually instantiated. + from lightly.loss import NTXentLoss + + self.ssl_loss = NTXentLoss(temperature=temperature) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Forward pass through encoder. + + Args: + x: Input images of shape [B, C, H, W]. + + Returns: + Tuple of (features, projections) where: + - features: Shape [B, feature_dim] for downstream tasks + - projections: Shape [B, projection_dim] for contrastive loss + """ + return self.network(x) + + def training_step(self, batch: tuple, batch_idx: int) -> dict[str, torch.Tensor]: + """Training step with contrastive loss. + + Computes NT-Xent loss between projections of two augmented views. + The loss encourages the model to produce similar representations + for different augmentations of the same image. + + Args: + batch: Tuple from LightlyDataset containing: + - (view0, view1): Two augmented views of each image + - targets: Original labels (unused in SSL) + - filenames: Image filenames + batch_idx: Index of the current batch. + + Returns: + Dict with 'loss' key for automatic logging. + """ + (view0, view1), targets, filenames = batch + + # Get projections for both views + _, z0 = self(view0) + _, z1 = self(view1) + + # Compute contrastive loss + loss = self.ssl_loss(z0, z1) + + self.log("train/loss", loss, prog_bar=True) + return {"loss": loss} + + def validation_step(self, batch: tuple, batch_idx: int) -> dict[str, torch.Tensor]: + """Validation step with contrastive loss. + + Same as training_step but without gradient computation. + + Args: + batch: Tuple from LightlyDataset (see training_step). + batch_idx: Index of the current batch. + + Returns: + Dict with 'loss' key for automatic logging. + """ + (view0, view1), targets, filenames = batch + + _, z0 = self(view0) + _, z1 = self(view1) + + loss = self.ssl_loss(z0, z1) + + self.log("val/loss", loss, prog_bar=True) + return {"loss": loss} + + def predict_step(self, batch: tuple, batch_idx: int) -> dict[str, Any]: + """Extract features for downstream tasks. + + Returns features (not projections) which are typically used for + downstream classification or other tasks via linear probing or + fine-tuning. + + Args: + batch: Tuple from LightlyDataset (see training_step). + batch_idx: Index of the current batch. + + Returns: + Dict containing: + - features: Extracted feature representations + - targets: Original labels for evaluation + - filenames: Image filenames for identification + """ + (view0, view1), targets, filenames = batch + + # Get features (not projections) for downstream tasks + features, _ = self(view0) + + return { + "features": features, + "targets": targets, + "filenames": filenames, + } diff --git a/projects/self_supervised/networks/__init__.py b/projects/self_supervised/networks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/self_supervised/networks/encoder.py b/projects/self_supervised/networks/encoder.py new file mode 100644 index 00000000..cf5c4fb5 --- /dev/null +++ b/projects/self_supervised/networks/encoder.py @@ -0,0 +1,137 @@ +"""Self-supervised learning models using lightly library. + +Lightly provides state-of-the-art SSL methods including SimCLR, BYOL, DINO, etc. +This module wraps lightly components for use with Lighter. +""" + +import torch +import torch.nn as nn +from torchvision import models + + +def create_simclr_model( + backbone: str = "resnet18", + projection_dim: int = 128, + hidden_dim: int = 2048, + pretrained: bool = False, +) -> nn.Module: + """Create a SimCLR model using lightly components. + + Args: + backbone: Backbone architecture name. + projection_dim: Output dimension of projection head. + hidden_dim: Hidden dimension of projection head. + pretrained: Use pretrained backbone weights. + + Returns: + SimCLR model with backbone and projection head. + """ + from lightly.models.modules import SimCLRProjectionHead + + # Create backbone + if backbone == "resnet18": + weights = models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None + base = models.resnet18(weights=weights) + backbone_dim = 512 + elif backbone == "resnet50": + weights = models.ResNet50_Weights.IMAGENET1K_V1 if pretrained else None + base = models.resnet50(weights=weights) + backbone_dim = 2048 + else: + raise ValueError(f"Unsupported backbone: {backbone}") + + # Remove classification head + backbone_model = nn.Sequential(*list(base.children())[:-1], nn.Flatten()) + + # Create projection head + projection_head = SimCLRProjectionHead(backbone_dim, hidden_dim, projection_dim) + + return SimCLRNetwork(backbone_model, projection_head) + + +class SimCLRNetwork(nn.Module): + """SimCLR network combining backbone and projection head.""" + + def __init__(self, backbone: nn.Module, projection_head: nn.Module) -> None: + super().__init__() + self.backbone = backbone + self.projection_head = projection_head + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Forward pass returning both features and projections. + + Args: + x: Input images [B, C, H, W]. + + Returns: + Tuple of (features, projections). + """ + features = self.backbone(x) + projections = self.projection_head(features) + return features, projections + + +def create_byol_model( + backbone: str = "resnet18", + projection_dim: int = 256, + hidden_dim: int = 4096, + pretrained: bool = False, +) -> nn.Module: + """Create a BYOL model using lightly components. + + BYOL (Bootstrap Your Own Latent) doesn't require negative pairs. + + Args: + backbone: Backbone architecture name. + projection_dim: Output dimension of projection head. + hidden_dim: Hidden dimension. + pretrained: Use pretrained backbone weights. + + Returns: + BYOL model. + """ + from lightly.models.modules import BYOLPredictionHead, BYOLProjectionHead + + # Create backbone + if backbone == "resnet18": + weights = models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None + base = models.resnet18(weights=weights) + backbone_dim = 512 + elif backbone == "resnet50": + weights = models.ResNet50_Weights.IMAGENET1K_V1 if pretrained else None + base = models.resnet50(weights=weights) + backbone_dim = 2048 + else: + raise ValueError(f"Unsupported backbone: {backbone}") + + backbone_model = nn.Sequential(*list(base.children())[:-1], nn.Flatten()) + projection_head = BYOLProjectionHead(backbone_dim, hidden_dim, projection_dim) + prediction_head = BYOLPredictionHead(projection_dim, hidden_dim, projection_dim) + + return BYOLNetwork(backbone_model, projection_head, prediction_head) + + +class BYOLNetwork(nn.Module): + """BYOL network with backbone, projection, and prediction heads.""" + + def __init__( + self, + backbone: nn.Module, + projection_head: nn.Module, + prediction_head: nn.Module, + ) -> None: + super().__init__() + self.backbone = backbone + self.projection_head = projection_head + self.prediction_head = prediction_head + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass. + + Returns: + Tuple of (features, projections, predictions). + """ + features = self.backbone(x) + projections = self.projection_head(features) + predictions = self.prediction_head(projections) + return features, projections, predictions diff --git a/projects/video_recognition/.gitignore b/projects/video_recognition/.gitignore new file mode 100644 index 00000000..b95f2c5f --- /dev/null +++ b/projects/video_recognition/.gitignore @@ -0,0 +1 @@ +.datasets/ diff --git a/projects/video_recognition/README.md b/projects/video_recognition/README.md new file mode 100644 index 00000000..0f78ac50 --- /dev/null +++ b/projects/video_recognition/README.md @@ -0,0 +1,52 @@ +# Video Action Recognition + +3D CNN and Video Transformer for action recognition. + +## Dataset + +**UCF101** - 101 action classes, 13,320 videos. Manual download required: + +```bash +cd projects/video_recognition && mkdir -p .datasets +# Download from https://www.crcv.ucf.edu/data/UCF101.php: +# - UCF101.rar +# - UCF101TrainTestSplits-RecognitionTask.zip +unar -o .datasets/ UCF101.rar +unzip UCF101TrainTestSplits-RecognitionTask.zip -d .datasets/ +``` + +## Models + +- **R3D** - 3D ResNet with spatiotemporal convolutions +- **ViViT** - Video Vision Transformer with tubelet embedding + +## Lighter Features Demonstrated + +- **Config composition** - `base.yaml` + model config merged via CLI +- **FileWriter + CsvWriter** - save video clips and classification results +- **`_mode_: callable`** - custom collate and writer functions +- **`%` raw references** - shared dataloader settings + +## Usage + +```bash +pip install lighter pytorchvideo av +cd projects/video_recognition + +# macOS Apple Silicon: enable MPS fallback +export PYTORCH_ENABLE_MPS_FALLBACK=1 + +# R3D (3D CNN) +lighter fit configs/base.yaml configs/r3d.yaml + +# ViViT (Transformer) +lighter fit configs/base.yaml configs/transformer.yaml + +# Quick test +lighter fit configs/base.yaml configs/r3d.yaml trainer::fast_dev_run=true +``` + +## References + +- [UCF101 Dataset](https://www.crcv.ucf.edu/data/UCF101.php) +- [pytorchvideo](https://pytorchvideo.org/) diff --git a/projects/video_recognition/__init__.py b/projects/video_recognition/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/video_recognition/__lighter__.py b/projects/video_recognition/__lighter__.py new file mode 100644 index 00000000..9b1eb488 --- /dev/null +++ b/projects/video_recognition/__lighter__.py @@ -0,0 +1,7 @@ +# This file marks the directory as a Lighter project. +# Video Action Recognition example demonstrating: +# - 3D CNN and Video Transformer architectures +# - CsvWriter for per-clip prediction logging +# - Config inheritance (base + dataset-specific configs) +# - Custom video dataloaders with temporal sampling +# - Multi-GPU training support for memory-intensive video models diff --git a/projects/video_recognition/configs/base.yaml b/projects/video_recognition/configs/base.yaml new file mode 100644 index 00000000..1a6ddee6 --- /dev/null +++ b/projects/video_recognition/configs/base.yaml @@ -0,0 +1,61 @@ +# Base configuration for Video Action Recognition +# +# Shared settings across all video models. Merge with model-specific config: +# cd projects/video_recognition && uv sync +# lighter fit configs/base.yaml,configs/r3d.yaml +# lighter fit configs/base.yaml,configs/transformer.yaml +# lighter predict configs/base.yaml,configs/r3d.yaml # saves .mp4 files +# +# Features demonstrated: +# - Config composition (base + model configs) +# - FileWriter callback for saving video clips as .mp4 +# - CsvWriter for classification results +# - Video handling (5D tensors: B, C, T, H, W) + +vars: + num_classes: 101 # UCF101 + frames_per_clip: 16 + frame_rate: 15 + num_workers: 4 + +trainer: + _target_: pytorch_lightning.Trainer + accelerator: auto + log_every_n_steps: 10 + # For quick testing, uncomment: + # fast_dev_run: true + + callbacks: + # Save classification results to CSV + - _target_: lighter.callbacks.CsvWriter + path: predictions.csv + keys: [prediction, confidence, ground_truth] + + # Save video clips as MP4 files + - _target_: lighter.callbacks.FileWriter + directory: ./video_outputs + value_key: video + writer_fn: + _target_: project.writers.write_video_mp4 + _mode_: callable + +data: + _target_: lighter.LighterDataModule + + train_dataloader: + _target_: torch.utils.data.DataLoader + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: true + collate_fn: + _target_: project.dataset.video_collate_fn + _mode_: callable + + val_dataloader: + _target_: torch.utils.data.DataLoader + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: false + collate_fn: + _target_: project.dataset.video_collate_fn + _mode_: callable diff --git a/projects/video_recognition/configs/r3d.yaml b/projects/video_recognition/configs/r3d.yaml new file mode 100644 index 00000000..422c6fd7 --- /dev/null +++ b/projects/video_recognition/configs/r3d.yaml @@ -0,0 +1,92 @@ +# R3D (3D ResNet) for Video Action Recognition +# +# 3D CNN using spatiotemporal convolutions for video classification. +# +# Run: lighter fit configs/base.yaml configs/r3d.yaml + +vars: + height: 112 + width: 112 + batch_size: 4 + +trainer: + max_epochs: 50 + accumulate_grad_batches: 4 # Effective batch size = 16 + + callbacks: + - _target_: pytorch_lightning.callbacks.EarlyStopping + monitor: val/top1_accuracy_epoch + mode: max + patience: 10 + + - _target_: pytorch_lightning.callbacks.ModelCheckpoint + monitor: val/top1_accuracy_epoch + mode: max + save_top_k: 1 + filename: r3d-{epoch}-{val/top1_accuracy_epoch:.4f} + +model: + _target_: project.models.model.VideoClassificationModel + + network: + _target_: project.networks.video_models.R3D + num_classes: "%vars::num_classes" + in_channels: 3 + num_frames: "%vars::frames_per_clip" + + criterion: + _target_: torch.nn.CrossEntropyLoss + label_smoothing: 0.1 + + optimizer: + _target_: torch.optim.AdamW + params: $@model::network.parameters() + lr: 0.0001 + weight_decay: 0.05 + + scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + optimizer: "@model::optimizer" + T_max: 50 + + train_metrics: + _target_: torchmetrics.MetricCollection + metrics: + top1_accuracy: + _target_: torchmetrics.classification.MulticlassAccuracy + num_classes: "%vars::num_classes" + top_k: 1 + top5_accuracy: + _target_: torchmetrics.classification.MulticlassAccuracy + num_classes: "%vars::num_classes" + top_k: 5 + + val_metrics: "%model::train_metrics" + test_metrics: "%model::train_metrics" + +data: + train_dataloader: + batch_size: "%vars::batch_size" + dataset: + _target_: project.dataset.VideoClipDataset + size: ["%vars::height", "%vars::width"] + dataset: + _target_: torchvision.datasets.UCF101 + root: ./.datasets/UCF-101 + annotation_path: ./.datasets/ucfTrainTestlist + frames_per_clip: "%vars::frames_per_clip" + frame_rate: "%vars::frame_rate" + train: true + + val_dataloader: + batch_size: "%vars::batch_size" + dataset: + _target_: project.dataset.VideoClipDataset + size: ["%vars::height", "%vars::width"] + dataset: + _target_: torchvision.datasets.UCF101 + root: ./.datasets/UCF-101 + annotation_path: ./.datasets/ucfTrainTestlist + frames_per_clip: "%vars::frames_per_clip" + frame_rate: "%vars::frame_rate" + train: false diff --git a/projects/video_recognition/configs/transformer.yaml b/projects/video_recognition/configs/transformer.yaml new file mode 100644 index 00000000..7786d8ad --- /dev/null +++ b/projects/video_recognition/configs/transformer.yaml @@ -0,0 +1,94 @@ +# Video Transformer (ViViT) for Video Action Recognition +# +# Tubelet embedding + Transformer encoder for video classification. +# Requires more memory than R3D - uses smaller batch size and gradient accumulation. +# +# Run: lighter fit configs/base.yaml configs/transformer.yaml + +vars: + height: 224 # Larger spatial size for ViT + width: 224 + batch_size: 2 # Transformers need more memory + +trainer: + max_epochs: 100 + accumulate_grad_batches: 8 # Effective batch size = 16 + + callbacks: + - _target_: pytorch_lightning.callbacks.EarlyStopping + monitor: val/MulticlassAccuracy_epoch + mode: max + patience: 10 + + - _target_: pytorch_lightning.callbacks.ModelCheckpoint + monitor: val/MulticlassAccuracy_epoch + mode: max + save_top_k: 1 + filename: transformer-{epoch}-{val/MulticlassAccuracy_epoch:.4f} + +model: + _target_: project.models.model.VideoClassificationModel + + network: + _target_: project.networks.video_models.VideoTransformer + num_classes: "%vars::num_classes" + img_size: "%vars::height" + num_frames: "%vars::frames_per_clip" + patch_size: [2, 16, 16] + embed_dim: 384 + num_heads: 6 + num_layers: 6 + in_channels: 3 + + criterion: + _target_: torch.nn.CrossEntropyLoss + label_smoothing: 0.1 + + optimizer: + _target_: torch.optim.AdamW + params: $@model::network.parameters() + lr: 0.00005 # Lower LR for transformers + weight_decay: 0.1 + + scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingWarmRestarts + optimizer: "@model::optimizer" + T_0: 10 + T_mult: 2 + + train_metrics: + _target_: torchmetrics.MetricCollection + metrics: + - _target_: torchmetrics.classification.MulticlassAccuracy + num_classes: "%vars::num_classes" + top_k: 1 + + val_metrics: "%model::train_metrics" + test_metrics: "%model::train_metrics" + +data: + train_dataloader: + batch_size: "%vars::batch_size" + dataset: + _target_: project.dataset.VideoClipDataset + size: ["%vars::height", "%vars::width"] + dataset: + _target_: torchvision.datasets.UCF101 + root: ./.datasets/UCF-101 + annotation_path: ./.datasets/ucfTrainTestlist + frames_per_clip: "%vars::frames_per_clip" + frame_rate: "%vars::frame_rate" + train: true + + val_dataloader: + batch_size: "%vars::batch_size" + dataset: + _target_: project.dataset.VideoClipDataset + size: ["%vars::height", "%vars::width"] + dataset: + _target_: torchvision.datasets.UCF101 + root: ./.datasets/UCF-101 + annotation_path: ./.datasets/ucfTrainTestlist + frames_per_clip: "%vars::frames_per_clip" + frame_rate: "%vars::frame_rate" + train: false diff --git a/projects/video_recognition/dataset.py b/projects/video_recognition/dataset.py new file mode 100644 index 00000000..9b8a9d58 --- /dev/null +++ b/projects/video_recognition/dataset.py @@ -0,0 +1,74 @@ +"""Video datasets for action recognition.""" + +import torch +from torch.utils.data import Dataset + + +class VideoClipDataset(Dataset): + """Wrapper for video datasets to ensure consistent output format. + + Handles video data format conversion and resizing for different backends. + + Args: + dataset: Base video dataset (UCF101, Kinetics, etc.). + size: Target size (height, width) for resizing frames. If None, no resizing. + """ + + def __init__( + self, + dataset: Dataset, + size: tuple[int, int] | None = None, + ) -> None: + self.dataset = dataset + self.size = size + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]: + # UCF101/Kinetics return (video, audio, label) + video, audio, label = self.dataset[idx] + + # video shape from UCF101: [T, H, W, C] + # Convert to float and normalize to [0, 1] + video = video.float() + if video.max() > 1.0: + video = video / 255.0 + + # Permute to [T, C, H, W] for per-frame processing + video = video.permute(0, 3, 1, 2) # [T, H, W, C] -> [T, C, H, W] + + # Resize if size is specified + if self.size is not None: + video = torch.nn.functional.interpolate( + video, + size=self.size, + mode="bilinear", + align_corners=False, + ) + + # Permute to [C, T, H, W] for 3D CNN + video = video.permute(1, 0, 2, 3) # [T, C, H, W] -> [C, T, H, W] + + return video, label + + +def video_collate_fn(batch: list) -> tuple[torch.Tensor, torch.Tensor]: + """Collate function for video batches. + + Handles variable-length videos by padding or truncating. + """ + videos, labels = zip(*batch, strict=True) + + # Find max temporal dimension + max_t = max(v.shape[1] for v in videos) + + # Pad videos shorter than max_t (no truncation needed since max_t is the max of all lengths) + padded_videos = [] + for v in videos: + if v.shape[1] < max_t: + pad_size = max_t - v.shape[1] + v = torch.nn.functional.pad(v, (0, 0, 0, 0, 0, pad_size)) + padded_videos.append(v) + + return torch.stack(padded_videos), torch.tensor(labels) diff --git a/projects/video_recognition/models/__init__.py b/projects/video_recognition/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/video_recognition/models/model.py b/projects/video_recognition/models/model.py new file mode 100644 index 00000000..b6eeb84b --- /dev/null +++ b/projects/video_recognition/models/model.py @@ -0,0 +1,80 @@ +"""Video action recognition model using LighterModule.""" + +import torch + +from lighter import LighterModule + + +class VideoClassificationModel(LighterModule): + """Video action recognition model. + + Showcases: + - Handling video inputs (5D tensors: B, C, T, H, W) + - Top-k accuracy metrics + - Predictions for CsvWriter callback + """ + + def _shared_step(self, batch: tuple, metrics) -> dict: + """Shared logic for train/val/test steps.""" + video, label = batch + + # Forward pass - network handles [B, C, T, H, W] input + logits = self(video) + + # Compute loss + loss = self.criterion(logits, label) + + # Get predictions + pred = logits.argmax(dim=1) + + # Update metrics + if metrics is not None: + metrics(logits, label) + + return { + "loss": loss, + "pred": pred, + "label": label, + "logits": logits, + } + + def training_step(self, batch, batch_idx): + return self._shared_step(batch, self.train_metrics) + + def validation_step(self, batch, batch_idx): + return self._shared_step(batch, self.val_metrics) + + def test_step(self, batch, batch_idx): + result = self._shared_step(batch, self.test_metrics) + del result["loss"] + return result + + def predict_step(self, batch, batch_idx): + """Prediction step returning class labels, probabilities, and video tensors. + + Returns dict compatible with CsvWriter and FileWriter callbacks. + - CsvWriter: uses prediction, confidence, ground_truth + - FileWriter: uses video tensor for mp4 output + """ + video, label = batch if isinstance(batch, (tuple, list)) else (batch, None) + + logits = self(video) + probs = torch.softmax(logits, dim=1) + pred = logits.argmax(dim=1) + + # Get top-k predictions (up to 5, or fewer if model has fewer classes) + k = min(5, probs.shape[1]) + topk_probs, topk_indices = probs.topk(k, dim=1) + + result = { + "prediction": pred.tolist(), + "confidence": probs.max(dim=1).values.tolist(), + "top5_classes": topk_indices.tolist(), + "top5_probs": topk_probs.tolist(), + "video": video, # Include video tensor for FileWriter + } + + if label is not None: + result["ground_truth"] = label.tolist() + + return result diff --git a/projects/video_recognition/networks/__init__.py b/projects/video_recognition/networks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/video_recognition/networks/video_models.py b/projects/video_recognition/networks/video_models.py new file mode 100644 index 00000000..151f4454 --- /dev/null +++ b/projects/video_recognition/networks/video_models.py @@ -0,0 +1,227 @@ +"""Video models for action recognition. + +Includes 3D CNN (R3D) and Video Transformer architectures. +""" + +import torch +import torch.nn as nn + + +class Conv3DBlock(nn.Module): + """3D convolution block with batch norm and ReLU.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: tuple[int, int, int] = (3, 3, 3), + stride: tuple[int, int, int] = (1, 1, 1), + padding: tuple[int, int, int] = (1, 1, 1), + ) -> None: + super().__init__() + self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, bias=False) + self.bn = nn.BatchNorm3d(out_channels) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.relu(self.bn(self.conv(x))) + + +class R3DBlock(nn.Module): + """Residual 3D block.""" + + def __init__(self, in_channels: int, out_channels: int, stride: int = 1) -> None: + super().__init__() + self.conv1 = Conv3DBlock(in_channels, out_channels, stride=(stride, stride, stride)) + self.conv2 = nn.Sequential( + nn.Conv3d(out_channels, out_channels, 3, 1, 1, bias=False), + nn.BatchNorm3d(out_channels), + ) + self.relu = nn.ReLU(inplace=True) + + self.downsample = None + if stride != 1 or in_channels != out_channels: + self.downsample = nn.Sequential( + nn.Conv3d(in_channels, out_channels, 1, (stride, stride, stride), bias=False), + nn.BatchNorm3d(out_channels), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + identity = x + out = self.conv1(x) + out = self.conv2(out) + if self.downsample: + identity = self.downsample(x) + out += identity + return self.relu(out) + + +class R3D(nn.Module): + """R3D (3D ResNet) for video classification. + + A simplified R3D-18 style architecture. + + Args: + num_classes: Number of action classes. + in_channels: Number of input channels (3 for RGB). + num_frames: Expected number of frames (for documentation). + """ + + def __init__( + self, + num_classes: int = 101, + in_channels: int = 3, + num_frames: int = 16, + ) -> None: + super().__init__() + self.num_frames = num_frames + + # Stem + self.stem = nn.Sequential( + nn.Conv3d(in_channels, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2), padding=(1, 3, 3), bias=False), + nn.BatchNorm3d(64), + nn.ReLU(inplace=True), + nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)), + ) + + # Residual layers + self.layer1 = self._make_layer(64, 64, 2) + self.layer2 = self._make_layer(64, 128, 2, stride=2) + self.layer3 = self._make_layer(128, 256, 2, stride=2) + self.layer4 = self._make_layer(256, 512, 2, stride=2) + + # Classification head + self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1)) + self.fc = nn.Linear(512, num_classes) + + def _make_layer(self, in_channels: int, out_channels: int, blocks: int, stride: int = 1) -> nn.Sequential: + layers = [R3DBlock(in_channels, out_channels, stride)] + for _ in range(1, blocks): + layers.append(R3DBlock(out_channels, out_channels)) + return nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass. + + Args: + x: Video tensor of shape [B, C, T, H, W]. + + Returns: + Class logits [B, num_classes]. + """ + x = self.stem(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + x = self.avgpool(x) + x = x.flatten(1) + return self.fc(x) + + +class PatchEmbed3D(nn.Module): + """3D patch embedding for video.""" + + def __init__( + self, + img_size: int = 224, + num_frames: int = 16, + patch_size: tuple[int, int, int] = (2, 16, 16), + in_channels: int = 3, + embed_dim: int = 768, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.num_patches = (num_frames // patch_size[0]) * (img_size // patch_size[1]) * (img_size // patch_size[2]) + self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: [B, C, T, H, W] + x = self.proj(x) # [B, embed_dim, T', H', W'] + x = x.flatten(2).transpose(1, 2) # [B, num_patches, embed_dim] + return x + + +class VideoTransformer(nn.Module): + """Simple Video Transformer (ViViT-style) for action recognition. + + Uses tubelet embedding and standard transformer encoder. + + Args: + num_classes: Number of action classes. + img_size: Spatial size of input frames. + num_frames: Number of input frames. + patch_size: Spatiotemporal patch size (T, H, W). + embed_dim: Transformer embedding dimension. + num_heads: Number of attention heads. + num_layers: Number of transformer layers. + mlp_ratio: MLP hidden dimension ratio. + """ + + def __init__( + self, + num_classes: int = 101, + img_size: int = 224, + num_frames: int = 16, + patch_size: tuple[int, int, int] = (2, 16, 16), + embed_dim: int = 768, + num_heads: int = 12, + num_layers: int = 12, + mlp_ratio: float = 4.0, + in_channels: int = 3, + ) -> None: + super().__init__() + + # Patch embedding + self.patch_embed = PatchEmbed3D(img_size, num_frames, patch_size, in_channels, embed_dim) + num_patches = self.patch_embed.num_patches + + # Positional embedding + CLS token + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) + + # Transformer encoder + encoder_layer = nn.TransformerEncoderLayer( + d_model=embed_dim, + nhead=num_heads, + dim_feedforward=int(embed_dim * mlp_ratio), + batch_first=True, + norm_first=True, + ) + self.transformer = nn.TransformerEncoder(encoder_layer, num_layers) + + # Classification head + self.norm = nn.LayerNorm(embed_dim) + self.head = nn.Linear(embed_dim, num_classes) + + # Initialize + nn.init.trunc_normal_(self.pos_embed, std=0.02) + nn.init.trunc_normal_(self.cls_token, std=0.02) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass. + + Args: + x: Video tensor [B, C, T, H, W]. + + Returns: + Class logits [B, num_classes]. + """ + B = x.shape[0] + + # Patch embedding + x = self.patch_embed(x) # [B, N, D] + + # Add CLS token + cls_tokens = self.cls_token.expand(B, -1, -1) + x = torch.cat([cls_tokens, x], dim=1) + + # Add positional embedding + x = x + self.pos_embed + + # Transformer + x = self.transformer(x) + + # Classification from CLS token + x = self.norm(x[:, 0]) + return self.head(x) diff --git a/projects/video_recognition/writers.py b/projects/video_recognition/writers.py new file mode 100644 index 00000000..74c592e8 --- /dev/null +++ b/projects/video_recognition/writers.py @@ -0,0 +1,45 @@ +"""Custom writer functions for video recognition outputs.""" + +from pathlib import Path + +import torch + + +def write_video_mp4(path: Path, tensor: torch.Tensor, *, suffix: str = ".mp4", fps: int = 15) -> None: + """Write a video tensor to MP4 format. + + Args: + path: Output path (suffix will be replaced). + tensor: Video tensor of shape (T, H, W, C) with values in [0, 255] as uint8, + or (C, T, H, W) / (T, C, H, W) with float values in [0, 1]. + suffix: File suffix (default: .mp4). + fps: Frames per second for output video. + """ + import torchvision.io + + path = path.with_suffix(suffix) + + # Handle different tensor layouts + if isinstance(tensor, torch.Tensor): + tensor = tensor.detach().cpu() + + # Determine layout and convert to (T, H, W, C) for torchvision + if tensor.ndim == 5: + # (B, C, T, H, W) -> take first batch + tensor = tensor[0] + + if tensor.ndim == 4: + # Could be (C, T, H, W) or (T, H, W, C) or (T, C, H, W) + if tensor.shape[0] == 3: + # (C, T, H, W) -> (T, H, W, C) + tensor = tensor.permute(1, 2, 3, 0) + elif tensor.shape[1] == 3: + # (T, C, H, W) -> (T, H, W, C) + tensor = tensor.permute(0, 2, 3, 1) + # else assume already (T, H, W, C) + + # Convert to uint8 if float + if tensor.dtype in (torch.float32, torch.float64, torch.float16): + tensor = (tensor.clamp(0, 1) * 255).to(torch.uint8) + + torchvision.io.write_video(str(path), tensor, fps=fps) diff --git a/projects/vision_language/README.md b/projects/vision_language/README.md new file mode 100644 index 00000000..82959d82 --- /dev/null +++ b/projects/vision_language/README.md @@ -0,0 +1,47 @@ +# Vision-Language Learning (CLIP-style) + +Dual-encoder for learning aligned image-text representations. + +## Task + +Learn joint embeddings for images and text for retrieval and zero-shot classification. + +## Dataset + +**Flickr8k** (default) - 8,000 images with 5 captions each +- Download from Kaggle: https://www.kaggle.com/datasets/adityajn105/flickr8k +- Extract to `.datasets/flickr8k/` + +**Flickr30k** (larger) - 31,000 images with 5 captions each +- Request access: https://shannon.cs.illinois.edu/DenotationGraph/ + +## Architecture + +- **Image Encoder**: ResNet-50 with projection head +- **Text Encoder**: Transformer with learned embeddings +- **Loss**: Symmetric contrastive loss with learnable temperature + +## Lighter Features Demonstrated + +- **Freezer callback** - freeze image encoder backbone during warmup +- **Differential learning rates** - lower LR for pretrained, higher for new layers +- **`_mode_: callable`** for custom collate function +- **`$` expressions** for parameter group filtering + +## Usage + +```bash +pip install lighter transformers +cd projects/vision_language + +# Train +lighter fit configs/clip.yaml + +# Quick test +lighter fit configs/clip.yaml trainer::fast_dev_run=true +``` + +## References + +- [CLIP Paper](https://arxiv.org/abs/2103.00020) +- [Flickr8k on Kaggle](https://www.kaggle.com/datasets/adityajn105/flickr8k) diff --git a/projects/vision_language/__init__.py b/projects/vision_language/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/vision_language/__lighter__.py b/projects/vision_language/__lighter__.py new file mode 100644 index 00000000..7f4f435e --- /dev/null +++ b/projects/vision_language/__lighter__.py @@ -0,0 +1,8 @@ +# This file marks the directory as a Lighter project. +# Vision-Language (CLIP-style) example demonstrating: +# - Dual-encoder architecture (image + text) +# - Contrastive learning between modalities +# - Freezer callback for frozen text encoder +# - HuggingFace transformers integration +# - Zero-shot classification evaluation +# - Temperature-scaled similarity diff --git a/projects/vision_language/configs/clip.yaml b/projects/vision_language/configs/clip.yaml new file mode 100644 index 00000000..57077df9 --- /dev/null +++ b/projects/vision_language/configs/clip.yaml @@ -0,0 +1,182 @@ +# Vision-Language (CLIP-style) Training with Lighter +# +# This example showcases: +# - Dual-encoder architecture (image + text encoders) +# - Contrastive learning between vision and language +# - Freezer callback for frozen text encoder (transfer learning) +# - Learnable temperature parameter +# - Custom collate function for multi-modal batches +# - Image-to-text and text-to-image retrieval metrics +# +# Applications: +# - Zero-shot image classification +# - Image-text retrieval +# - Visual search +# - Multi-modal embeddings +# +# Run: lighter fit projects/vision_language/configs/clip.yaml + +vars: + embed_dim: 512 + batch_size: 64 + image_size: 224 + max_text_length: 77 + vocab_size: 30522 # BERT tokenizer vocab size + temperature: 0.07 + num_workers: 4 + # Dataset path - download Flickr8k from https://www.kaggle.com/datasets/adityajn105/flickr8k + # For Flickr30k, use: .datasets/flickr30k and change dataset class below + data_root: .datasets/flickr8k + +trainer: + _target_: pytorch_lightning.Trainer + max_epochs: 50 + accelerator: auto + log_every_n_steps: 10 + gradient_clip_val: 1.0 # Important for CLIP training stability + + callbacks: + # Freeze image encoder backbone for first 5 epochs + # This helps when using pretrained vision encoder + - _target_: lighter.callbacks.Freezer + names: + - image_encoder.backbone + until_epoch: 5 + + # Early stopping on retrieval accuracy + - _target_: pytorch_lightning.callbacks.EarlyStopping + monitor: val/i2t_accuracy + mode: max + patience: 10 + + # Save best model + - _target_: pytorch_lightning.callbacks.ModelCheckpoint + monitor: val/i2t_accuracy + mode: max + save_top_k: 1 + filename: clip-{epoch}-{val/i2t_accuracy:.4f} + + # For quick testing, uncomment: + # fast_dev_run: true + +model: + _target_: project.models.model.CLIPLighterModel + + network: + _target_: project.networks.clip_model.CLIPModel + embed_dim: "%vars::embed_dim" + temperature: "%vars::temperature" + learnable_temperature: true + + image_encoder: + _target_: project.networks.clip_model.ImageEncoder + backbone: resnet50 + embed_dim: "%vars::embed_dim" + pretrained: true + + text_encoder: + _target_: project.networks.clip_model.TextEncoder + vocab_size: "%vars::vocab_size" + embed_dim: "%vars::embed_dim" + num_heads: 8 + num_layers: 4 + max_length: "%vars::max_text_length" + + # No criterion - loss computed in step + + # Different learning rates for different components + optimizer: + _target_: torch.optim.AdamW + params: + # Image encoder (pretrained) - lower learning rate + - params: "$[p for n, p in @model::network.named_parameters() if 'image_encoder' in n]" + lr: 0.00001 + # Text encoder - higher learning rate + - params: "$[p for n, p in @model::network.named_parameters() if 'text_encoder' in n]" + lr: 0.0001 + # Temperature parameter + - params: "$[p for n, p in @model::network.named_parameters() if 'log_temperature' in n]" + lr: 0.001 + weight_decay: 0.01 + + scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingWarmRestarts + optimizer: "@model::optimizer" + T_0: 10 + T_mult: 2 + +data: + _target_: lighter.LighterDataModule + + train_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: true + drop_last: true # Important for contrastive learning + collate_fn: + _target_: project.dataset.collate_fn + _mode_: callable + dataset: + # Use Flickr8kDataset (easier to obtain) or Flickr30kDataset + _target_: project.dataset.Flickr8kDataset + root: "%vars::data_root" + split: train + max_length: "%vars::max_text_length" + image_transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.RandomResizedCrop + size: "%vars::image_size" + scale: [0.8, 1.0] + - _target_: torchvision.transforms.RandomHorizontalFlip + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + val_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: "%vars::batch_size" + num_workers: "%vars::num_workers" + pin_memory: true + shuffle: false + drop_last: false + collate_fn: + _target_: project.dataset.collate_fn + _mode_: callable + dataset: + _target_: project.dataset.Flickr8kDataset + root: "%vars::data_root" + split: val + max_length: "%vars::max_text_length" + image_transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: ["%vars::image_size", "%vars::image_size"] + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + predict_dataloader: + _target_: torch.utils.data.DataLoader + batch_size: 16 + num_workers: "%vars::num_workers" + shuffle: false + collate_fn: + _target_: project.dataset.collate_fn + _mode_: callable + dataset: + _target_: project.dataset.Flickr8kDataset + root: "%vars::data_root" + split: test + max_length: "%vars::max_text_length" + image_transform: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: ["%vars::image_size", "%vars::image_size"] + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] diff --git a/projects/vision_language/dataset.py b/projects/vision_language/dataset.py new file mode 100644 index 00000000..49ec4a02 --- /dev/null +++ b/projects/vision_language/dataset.py @@ -0,0 +1,386 @@ +"""Vision-Language datasets for CLIP-style training. + +Real datasets for vision-language learning: +- Flickr30k: 31k images with 5 captions each +- COCO Captions: 330k images with captions +- Flickr8k: 8k images with 5 captions each (smaller, easier to obtain) +""" + +from collections.abc import Callable +from pathlib import Path + +import torch +from torch.utils.data import Dataset + + +class Flickr30kDataset(Dataset): + """Flickr30k dataset for vision-language training. + + Flickr30k contains 31,000 images with 5 captions each. + Note: Split filtering requires split manifest files (not included). + Currently loads all data regardless of split parameter. + + Download: + 1. Request access: https://shannon.cs.illinois.edu/DenotationGraph/ + 2. Extract images to root/flickr30k-images/ + 3. Download captions file to root/results_20130124.token + + Args: + root: Root directory containing images and captions. + tokenizer: HuggingFace tokenizer for text encoding. + max_length: Maximum text sequence length. + image_transform: Transform for images. + """ + + def __init__( + self, + root: str, + tokenizer=None, + max_length: int = 77, + image_transform: Callable | None = None, + ) -> None: + self.root = Path(root) + self.tokenizer = tokenizer + self.max_length = max_length + self.image_transform = image_transform + + # Load captions + self.data = self._load_captions() + + if len(self.data) == 0: + raise FileNotFoundError( + f"No data found for Flickr30k dataset at '{self.root}'.\n\n" + f"Please download the dataset:\n" + f"1. Request access: https://shannon.cs.illinois.edu/DenotationGraph/\n" + f"2. Extract images to: {self.root}/flickr30k-images/\n" + f"3. Download captions to: {self.root}/results_20130124.token\n\n" + f"Or use Flickr8kDataset for a smaller, easier-to-obtain alternative." + ) + + def _load_captions(self) -> list[dict]: + """Load image-caption pairs.""" + caption_file = self.root / "results_20130124.token" + image_dir = self.root / "flickr30k-images" + + data = [] + if caption_file.exists(): + with open(caption_file) as f: + for line in f: + parts = line.strip().split("\t") + if len(parts) >= 2: + image_id = parts[0].split("#")[0] + caption = parts[1] + image_path = image_dir / image_id + if image_path.exists(): + data.append({"image_path": image_path, "caption": caption}) + return data + + def __len__(self) -> int: + return len(self.data) + + def __getitem__(self, idx: int) -> dict: + item = self.data[idx] + + # Load image + image = self._load_image(item["image_path"]) + if self.image_transform: + image = self.image_transform(image) + + # Tokenize text + if self.tokenizer: + encoding = self.tokenizer( + item["caption"], + max_length=self.max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + input_ids = encoding["input_ids"].squeeze(0) + attention_mask = encoding["attention_mask"].squeeze(0) + else: + # Dummy tokenization for testing + input_ids = torch.zeros(self.max_length, dtype=torch.long) + attention_mask = torch.ones(self.max_length, dtype=torch.long) + + return { + "image": image, + "input_ids": input_ids, + "attention_mask": attention_mask, + "caption": item["caption"], + } + + def _load_image(self, path: Path) -> torch.Tensor: + """Load image from file.""" + try: + from torchvision.io import read_image + + return read_image(str(path)).float() / 255.0 + except Exception: + return torch.zeros(3, 224, 224) + + +class Flickr8kDataset(Dataset): + """Flickr8k dataset for vision-language training. + + Flickr8k contains 8,000 images with 5 captions each. + Smaller and easier to obtain than Flickr30k. + + Download from Kaggle: + https://www.kaggle.com/datasets/adityajn105/flickr8k + + Args: + root: Root directory containing images and captions. + split: Dataset split ('train', 'val', 'test'). If split files exist, + only images from that split are loaded. Otherwise loads all images. + tokenizer: HuggingFace tokenizer for text encoding. + max_length: Maximum text sequence length. + image_transform: Transform for images. + """ + + # Standard Flickr8k splits + SPLITS = { + "train": "Flickr_8k.trainImages.txt", + "val": "Flickr_8k.devImages.txt", + "test": "Flickr_8k.testImages.txt", + } + + def __init__( + self, + root: str, + split: str = "train", + tokenizer=None, + max_length: int = 77, + image_transform: Callable | None = None, + ) -> None: + self.root = Path(root) + self.split = split + self.tokenizer = tokenizer + self.max_length = max_length + self.image_transform = image_transform + + self.data = self._load_data() + + if len(self.data) == 0: + raise FileNotFoundError( + f"No data found for Flickr8k dataset at '{self.root}'.\n\n" + "Please download the dataset from:\n" + "https://www.kaggle.com/datasets/adityajn105/flickr8k\n\n" + "Expected structure:\n" + f" {self.root}/Images/ (image files)\n" + f" {self.root}/captions.txt (caption file)" + ) + + def _load_split_images(self) -> set[str] | None: + """Load image names for the current split, if split file exists.""" + if self.split not in self.SPLITS: + return None + split_file = self.root / self.SPLITS[self.split] + if not split_file.exists(): + return None + with open(split_file) as f: + return {line.strip() for line in f if line.strip()} + + def _load_data(self) -> list[dict]: + """Load image-caption pairs.""" + caption_file = self.root / "captions.txt" + image_dir = self.root / "Images" + + # Try alternative paths + if not caption_file.exists(): + caption_file = self.root / "Flickr8k.token.txt" + if not image_dir.exists(): + image_dir = self.root / "Flickr8k_Dataset" + + # Load split filter if available + valid_images = self._load_split_images() + + data = [] + if caption_file.exists(): + with open(caption_file) as f: + # Skip header if present + first_line = f.readline() + if not first_line.startswith("image"): + f.seek(0) + + for line in f: + line = line.strip() + if not line: + continue + + # Handle both CSV format and token format + if "," in line: + # CSV format: image,caption + parts = line.split(",", 1) + if len(parts) >= 2: + image_name = parts[0].strip() + caption = parts[1].strip() + else: + # Token format: image#idx\tcaption + parts = line.split("\t") + if len(parts) >= 2: + image_name = parts[0].split("#")[0] + caption = parts[1] + else: + continue + + # Filter by split if split file exists + if valid_images is not None and image_name not in valid_images: + continue + + image_path = image_dir / image_name + if image_path.exists(): + data.append({"image_path": image_path, "caption": caption}) + + return data + + def __len__(self) -> int: + return len(self.data) + + def __getitem__(self, idx: int) -> dict: + item = self.data[idx] + + image = self._load_image(item["image_path"]) + if self.image_transform: + image = self.image_transform(image) + + if self.tokenizer: + encoding = self.tokenizer( + item["caption"], + max_length=self.max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + input_ids = encoding["input_ids"].squeeze(0) + attention_mask = encoding["attention_mask"].squeeze(0) + else: + input_ids = torch.zeros(self.max_length, dtype=torch.long) + attention_mask = torch.ones(self.max_length, dtype=torch.long) + + return { + "image": image, + "input_ids": input_ids, + "attention_mask": attention_mask, + "caption": item["caption"], + } + + def _load_image(self, path: Path) -> torch.Tensor: + """Load image from file.""" + from torchvision.io import read_image + + return read_image(str(path)).float() / 255.0 + + +class COCOCaptionsDataset(Dataset): + """COCO Captions dataset for vision-language training. + + COCO Captions: 330k images with 5 captions each. + + Download: + 1. Images: https://cocodataset.org/#download + 2. Annotations: Download 2017 Train/Val annotations + + Args: + root: Root directory containing images. + ann_file: Path to annotations JSON file. + tokenizer: HuggingFace tokenizer. + max_length: Maximum text sequence length. + image_transform: Transform for images. + """ + + def __init__( + self, + root: str, + ann_file: str, + tokenizer=None, + max_length: int = 77, + image_transform: Callable | None = None, + ) -> None: + self.root = Path(root) + self.ann_file = Path(ann_file) + self.tokenizer = tokenizer + self.max_length = max_length + self.image_transform = image_transform + + self.data = self._load_annotations() + + if len(self.data) == 0: + raise FileNotFoundError( + f"No data found for COCO Captions dataset.\n\n" + f"Images root: {self.root}\n" + f"Annotations: {self.ann_file}\n\n" + "Please download the dataset from:\n" + "https://cocodataset.org/#download\n\n" + "Expected files:\n" + " - train2017/ or val2017/ (image directories)\n" + " - annotations/captions_train2017.json" + ) + + def _load_annotations(self) -> list[dict]: + """Load COCO annotations.""" + import json + + data = [] + if self.ann_file.exists(): + with open(self.ann_file) as f: + coco = json.load(f) + + # Build image id to filename mapping + id_to_file = {img["id"]: img["file_name"] for img in coco["images"]} + + # Load caption annotations + for ann in coco["annotations"]: + image_id = ann["image_id"] + if image_id in id_to_file: + data.append( + { + "image_path": self.root / id_to_file[image_id], + "caption": ann["caption"], + } + ) + return data + + def __len__(self) -> int: + return len(self.data) + + def __getitem__(self, idx: int) -> dict: + item = self.data[idx] + + image = self._load_image(item["image_path"]) + if self.image_transform: + image = self.image_transform(image) + + if self.tokenizer: + encoding = self.tokenizer( + item["caption"], + max_length=self.max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + input_ids = encoding["input_ids"].squeeze(0) + attention_mask = encoding["attention_mask"].squeeze(0) + else: + input_ids = torch.zeros(self.max_length, dtype=torch.long) + attention_mask = torch.ones(self.max_length, dtype=torch.long) + + return { + "image": image, + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + def _load_image(self, path: Path) -> torch.Tensor: + """Load image from file.""" + from torchvision.io import read_image + + return read_image(str(path)).float() / 255.0 + + +def collate_fn(batch: list[dict]) -> dict: + """Collate function for vision-language batches.""" + return { + "image": torch.stack([item["image"] for item in batch]), + "input_ids": torch.stack([item["input_ids"] for item in batch]), + "attention_mask": torch.stack([item["attention_mask"] for item in batch]), + } diff --git a/projects/vision_language/models/__init__.py b/projects/vision_language/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/vision_language/models/model.py b/projects/vision_language/models/model.py new file mode 100644 index 00000000..8a88048b --- /dev/null +++ b/projects/vision_language/models/model.py @@ -0,0 +1,136 @@ +"""Vision-Language CLIP model using LighterModule.""" + +import torch +import torch.nn.functional as F + +from lighter import LighterModule + + +def clip_loss( + image_features: torch.Tensor, + text_features: torch.Tensor, + temperature: torch.Tensor, +) -> torch.Tensor: + """Compute symmetric CLIP contrastive loss. + + Args: + image_features: L2-normalized image embeddings [B, D]. + text_features: L2-normalized text embeddings [B, D]. + temperature: Temperature parameter. + + Returns: + Scalar loss value. + """ + # Compute similarity matrix + logits = image_features @ text_features.t() / temperature + + # Labels: diagonal elements are positive pairs + batch_size = image_features.shape[0] + labels = torch.arange(batch_size, device=image_features.device) + + # Symmetric loss: image->text and text->image + loss_i2t = F.cross_entropy(logits, labels) + loss_t2i = F.cross_entropy(logits.t(), labels) + + return (loss_i2t + loss_t2i) / 2 + + +class CLIPLighterModel(LighterModule): + """CLIP-style vision-language model. + + Showcases: + - Multi-modal (image + text) inputs + - Custom contrastive loss + - Temperature as learnable parameter + - Retrieval metrics (image-to-text, text-to-image) + """ + + def forward( + self, + image: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass through both encoders.""" + return self.network(image, input_ids, attention_mask) + + def _shared_step(self, batch: dict, metrics) -> dict: + """Shared logic for train/val/test steps.""" + image = batch["image"] + input_ids = batch["input_ids"] + attention_mask = batch.get("attention_mask") + + # Forward pass + image_features, text_features, temperature = self(image, input_ids, attention_mask) + + # Compute CLIP loss + loss = clip_loss(image_features, text_features, temperature) + + # Compute retrieval accuracy + with torch.no_grad(): + logits = image_features @ text_features.t() + batch_size = image_features.shape[0] + labels = torch.arange(batch_size, device=image_features.device) + + # Image-to-text retrieval accuracy + i2t_acc = (logits.argmax(dim=1) == labels).float().mean() + # Text-to-image retrieval accuracy + t2i_acc = (logits.argmax(dim=0) == labels).float().mean() + + # Update metrics if available + if metrics is not None: + # Note: CLIP typically uses retrieval metrics, not classification metrics + pass + + return { + "loss": loss, + "i2t_accuracy": i2t_acc, + "t2i_accuracy": t2i_acc, + "temperature": temperature.detach(), + "image_features": image_features, + "text_features": text_features, + } + + def training_step(self, batch, batch_idx): + result = self._shared_step(batch, self.train_metrics) + + # Log retrieval accuracy + self.log("train/i2t_accuracy", result["i2t_accuracy"], prog_bar=True) + self.log("train/t2i_accuracy", result["t2i_accuracy"]) + self.log("train/temperature", result["temperature"]) + + return result + + def validation_step(self, batch, batch_idx): + result = self._shared_step(batch, self.val_metrics) + + self.log("val/i2t_accuracy", result["i2t_accuracy"], prog_bar=True) + self.log("val/t2i_accuracy", result["t2i_accuracy"]) + + return result + + def test_step(self, batch, batch_idx): + result = self._shared_step(batch, self.test_metrics) + del result["loss"] + + self.log("test/i2t_accuracy", result["i2t_accuracy"]) + self.log("test/t2i_accuracy", result["t2i_accuracy"]) + + return result + + def predict_step(self, batch, batch_idx): + """Prediction step returning embeddings for retrieval.""" + image = batch["image"] + input_ids = batch["input_ids"] + attention_mask = batch.get("attention_mask") + + image_features, text_features, temperature = self(image, input_ids, attention_mask) + + # Compute similarity + similarity = (image_features @ text_features.t()).diag() + + return { + "image_features": image_features.tolist(), + "text_features": text_features.tolist(), + "similarity": similarity.tolist(), + } diff --git a/projects/vision_language/networks/__init__.py b/projects/vision_language/networks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/projects/vision_language/networks/clip_model.py b/projects/vision_language/networks/clip_model.py new file mode 100644 index 00000000..d32d8128 --- /dev/null +++ b/projects/vision_language/networks/clip_model.py @@ -0,0 +1,221 @@ +"""CLIP-style Vision-Language model. + +Dual-encoder architecture that learns aligned image-text representations +through contrastive learning. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision import models + + +class ImageEncoder(nn.Module): + """Vision encoder using ResNet or ViT backbone. + + Projects images to a shared embedding space. + + Args: + backbone: Name of torchvision model ('resnet50', 'resnet18', 'vit_b_16'). + embed_dim: Output embedding dimension. + pretrained: Use pretrained weights. + """ + + def __init__( + self, + backbone: str = "resnet50", + embed_dim: int = 512, + pretrained: bool = True, + ) -> None: + super().__init__() + + if backbone == "resnet50": + weights = models.ResNet50_Weights.IMAGENET1K_V1 if pretrained else None + base = models.resnet50(weights=weights) + self.backbone = nn.Sequential(*list(base.children())[:-1]) # Remove FC + backbone_dim = 2048 + elif backbone == "resnet18": + weights = models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None + base = models.resnet18(weights=weights) + self.backbone = nn.Sequential(*list(base.children())[:-1]) + backbone_dim = 512 + elif backbone == "vit_b_16": + weights = models.ViT_B_16_Weights.IMAGENET1K_V1 if pretrained else None + self.backbone = models.vit_b_16(weights=weights) + self.backbone.heads = nn.Identity() # Remove classification head + backbone_dim = 768 + else: + raise ValueError(f"Unsupported backbone: {backbone}") + + # Project to shared embedding space + self.projection = nn.Sequential( + nn.Linear(backbone_dim, embed_dim), + nn.LayerNorm(embed_dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Encode images to embedding space. + + Args: + x: Images [B, C, H, W]. + + Returns: + Image embeddings [B, embed_dim]. + """ + features = self.backbone(x) + features = features.flatten(1) + return self.projection(features) + + +class TextEncoder(nn.Module): + """Text encoder using Transformer. + + Simple transformer encoder for text. In production, use + HuggingFace transformers (BERT, RoBERTa, etc.) for better results. + + Args: + vocab_size: Size of vocabulary. + embed_dim: Embedding dimension. + num_heads: Number of attention heads. + num_layers: Number of transformer layers. + max_length: Maximum sequence length. + """ + + def __init__( + self, + vocab_size: int = 30522, # BERT vocab size + embed_dim: int = 512, + num_heads: int = 8, + num_layers: int = 4, + max_length: int = 77, + ) -> None: + super().__init__() + + self.embedding = nn.Embedding(vocab_size, embed_dim) + self.pos_embedding = nn.Embedding(max_length, embed_dim) + + encoder_layer = nn.TransformerEncoderLayer( + d_model=embed_dim, + nhead=num_heads, + dim_feedforward=embed_dim * 4, + batch_first=True, + norm_first=True, + ) + self.transformer = nn.TransformerEncoder(encoder_layer, num_layers) + + self.norm = nn.LayerNorm(embed_dim) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Encode text to embedding space. + + Args: + input_ids: Token IDs [B, L]. + attention_mask: Attention mask [B, L]. + + Returns: + Text embeddings [B, embed_dim]. + """ + B, L = input_ids.shape + + # Token + positional embeddings + positions = torch.arange(L, device=input_ids.device).unsqueeze(0).expand(B, -1) + x = self.embedding(input_ids) + self.pos_embedding(positions) + + # Create attention mask for transformer + if attention_mask is not None: + # Convert to additive mask (0 = attend, -inf = ignore) + src_key_padding_mask = attention_mask == 0 + else: + src_key_padding_mask = None + + # Transformer encoding + x = self.transformer(x, src_key_padding_mask=src_key_padding_mask) + + # Pool: use [CLS] token (first token) or mean pooling + x = x[:, 0] # CLS token pooling + + return self.norm(x) + + +class CLIPModel(nn.Module): + """CLIP-style dual encoder for vision-language learning. + + Learns to align image and text representations through contrastive learning. + + Args: + image_encoder: Image encoder module. + text_encoder: Text encoder module. + embed_dim: Shared embedding dimension. + temperature: Initial temperature for contrastive loss. + learnable_temperature: Whether temperature is learnable. + """ + + def __init__( + self, + image_encoder: nn.Module, + text_encoder: nn.Module, + embed_dim: int = 512, + temperature: float = 0.07, + learnable_temperature: bool = True, + ) -> None: + super().__init__() + + self.image_encoder = image_encoder + self.text_encoder = text_encoder + + if learnable_temperature: + # Learnable log temperature (more stable optimization) + self.log_temperature = nn.Parameter(torch.log(torch.tensor(temperature))) + else: + self.register_buffer("log_temperature", torch.log(torch.tensor(temperature))) + + @property + def temperature(self) -> torch.Tensor: + return self.log_temperature.exp() + + def encode_image(self, image: torch.Tensor) -> torch.Tensor: + """Encode images and L2-normalize.""" + features = self.image_encoder(image) + return F.normalize(features, dim=-1) + + def encode_text( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Encode text and L2-normalize.""" + features = self.text_encoder(input_ids, attention_mask) + return F.normalize(features, dim=-1) + + def forward( + self, + image: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass computing image and text embeddings. + + Args: + image: Images [B, C, H, W]. + input_ids: Token IDs [B, L]. + attention_mask: Attention mask [B, L]. + + Returns: + Tuple of (image_features, text_features, temperature). + """ + image_features = self.encode_image(image) + text_features = self.encode_text(input_ids, attention_mask) + + return image_features, text_features, self.temperature + + def compute_similarity( + self, + image_features: torch.Tensor, + text_features: torch.Tensor, + ) -> torch.Tensor: + """Compute cosine similarity matrix scaled by temperature.""" + return image_features @ text_features.t() / self.temperature diff --git a/pyproject.toml b/pyproject.toml index 9e1e3c97..17ef168c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ dependencies = [ "torchmetrics>=1.2.0", "tensorboard>=2.11.2", "requests>=2.31.0", - "sparkwheel>=0.0.9", + "sparkwheel>=0.0.10", "rich>=13.7.0", "torchvision>=0.20.0", "cloudpickle>=3.0.0", diff --git a/src/lighter/utils/dynamic_imports.py b/src/lighter/utils/dynamic_imports.py index 2c393e26..ca1bfa1d 100644 --- a/src/lighter/utils/dynamic_imports.py +++ b/src/lighter/utils/dynamic_imports.py @@ -13,8 +13,6 @@ dispatch table. """ -from __future__ import annotations - import importlib.abc import importlib.machinery import importlib.util diff --git a/uv.lock b/uv.lock index bc6651de..5ff09409 100644 --- a/uv.lock +++ b/uv.lock @@ -628,7 +628,7 @@ requires-dist = [ { name = "pytorch-lightning", specifier = ">=2.1.3" }, { name = "requests", specifier = ">=2.31.0" }, { name = "rich", specifier = ">=13.7.0" }, - { name = "sparkwheel", specifier = ">=0.0.9" }, + { name = "sparkwheel", specifier = ">=0.0.10" }, { name = "tensorboard", specifier = ">=2.11.2" }, { name = "torch", specifier = ">=2.1.2" }, { name = "torchmetrics", specifier = ">=1.2.0" }, @@ -1851,14 +1851,14 @@ wheels = [ [[package]] name = "sparkwheel" -version = "0.0.9" +version = "0.0.10" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyyaml" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e5/3e/eea7646716b39fe523c42e7672df5c0a4d3087351c809ba6cf59e3ce09e8/sparkwheel-0.0.9.tar.gz", hash = "sha256:604cded3ecc6c8dceb5b769e9eb273e15e0dc206598b549ad18595f842ab80bc", size = 49543, upload-time = "2025-11-29T03:30:25.302Z" } +sdist = { url = "https://files.pythonhosted.org/packages/91/93/5c03db0a3b063a8cd11c2a873443b4f4e56f109057665454af9fe73a634a/sparkwheel-0.0.10.tar.gz", hash = "sha256:8982aa17ae7ef949803a320c5f7e8a2ff8c6e12e7232656ff890e9e25d95eec2", size = 49545, upload-time = "2025-12-09T03:06:47.76Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b6/a5/d03dce442b1226905837a07668399fa9232cc73a2684041df3b0cedc0229/sparkwheel-0.0.9-py3-none-any.whl", hash = "sha256:d18f6eef804ead414c1dad3fab055ee4e701e67d6f91461b8ce08f71c9cd7b28", size = 59158, upload-time = "2025-11-29T03:30:26.465Z" }, + { url = "https://files.pythonhosted.org/packages/e6/95/5528a9e0d9510ace944a1f4cad443e0796df6134e0bdad79c8cd30ace9f2/sparkwheel-0.0.10-py3-none-any.whl", hash = "sha256:14a54308f7eb691eb8089f376ee74694c9ee1bc417dbd71c01eda8b653fda7a8", size = 59167, upload-time = "2025-12-09T03:06:46.882Z" }, ] [[package]]