PyTorch Lightning Integration Example

This example demonstrates a complete EvoAug2 training workflow using PyTorch Lightning, implementing the two-stage training approach with comprehensive checkpoint management and performance evaluation.

Overview

The Lightning module example (example_lightning_module.py) showcases:

  • Two-stage training approach (augmentation + fine-tuning)

  • PyTorch Lightning integration with custom DataModules

  • Checkpoint management and resumption capabilities

  • Performance comparison between different training strategies

  • Comprehensive visualization of results

  • DeepSTARR model training on genomic regulatory data

Key Features

  1. RobustLoader Integration: Uses EvoAug2’s RobustLoader for efficient augmentation management

  2. Checkpoint Resumption: Automatically detects and resumes from existing checkpoints

  3. Performance Metrics: Comprehensive evaluation with Pearson and Spearman correlations

  4. Visualization: Automatic generation of comparison plots and performance summaries

  5. Flexible Configuration: Easy parameter modification for different experiments

File Structure

example_lightning_module.py
├── Configuration and setup
├── Checkpoint management functions
├── Plotting and visualization functions
├── Main training function
│   ├── Stage 1: Training with augmentations
│   ├── Stage 2: Fine-tuning on original data
│   └── Control: Standard training comparison
└── Results analysis and visualization

Usage

Basic Execution:

python example_lightning_module.py

Prerequisites:

# Install dependencies
pip install evoaug2[full]

# Download DeepSTARR data
wget https://zenodo.org/record/7265991/files/DeepSTARR_data.h5

Configuration:

The script can be customized by modifying these parameters:

# Experiment configuration
expt_name = 'DeepSTARR'
data_path = '.'
filepath = '/path/to/deepstarr-data.h5'
output_dir = '/path/to/output/'
batch_size = 128

# Augmentation parameters (DeepSTARR optimal settings)
augment_list = [
    RandomTranslocation(shift_min=0, shift_max=20),
    RandomRC(rc_prob=0.0),
    RandomMutation(mut_frac=0.05),
    RandomNoise(noise_mean=0, noise_std=0.3),
]

# Training parameters
max_augs_per_seq = 2  # Maximum augmentations per sequence
hard_aug = True        # Always apply exactly 2 augmentations
max_epochs = 100       # Stage 1 training epochs
finetune_epochs = 5    # Stage 2 fine-tuning epochs

Training Stages

Stage 1: Augmentation Training

# Create augmented data module
data_module = AugmentedDataModule(
    base_dataset,
    augment_list,
    max_augs_per_seq=2,
    hard_aug=True
)

# Train with augmentations
trainer.fit(model, datamodule=data_module)

Stage 2: Fine-tuning

# Load best augmented model
model_finetune = DeepSTARRModel.load_from_checkpoint(
    best_model_path,
    model=deepstarr
)

# Fine-tune on original data
model_finetune.learning_rate = 0.0001
trainer_finetune.fit(model_finetune, datamodule=data_module_finetune)

Control Training

# Train control model without augmentations
model_control = DeepSTARRModel(deepstarr_control)
trainer_control.fit(model_control, datamodule=data_module_control)

DataModule Implementation

AugmentedDataModule:

class AugmentedDataModule(pl.LightningDataModule):
    def __init__(self, base_dataset, augment_list, max_augs_per_seq, hard_aug):
        super().__init__()
        self.base_dataset = base_dataset
        self.augment_list = augment_list
        self.max_augs_per_seq = max_augs_per_seq
        self.hard_aug = hard_aug

    def train_dataloader(self):
        # Use RobustLoader with training dataset
        train_dataset = self.base_dataset.get_train_dataset()
        return RobustLoader(
            base_dataset=train_dataset,
            augment_list=self.augment_list,
            max_augs_per_seq=self.max_augs_per_seq,
            hard_aug=self.hard_aug,
            batch_size=self.base_dataset.batch_size,
            shuffle=True
        )

    def val_dataloader(self):
        # Validation with augmentations disabled
        val_dataset = self.base_dataset.get_val_dataset()
        loader = RobustLoader(...)
        loader.disable_augmentations()
        return loader

FineTuneDataModule:

class FineTuneDataModule(pl.LightningDataModule):
    def __init__(self, base_dataset):
        super().__init__()
        self.base_dataset = base_dataset

    def train_dataloader(self):
        return self.base_dataset.train_dataloader()

Checkpoint Management

Automatic Detection:

def check_existing_checkpoints(output_dir, expt_name):
    """Check for existing checkpoints and return their status."""
    checkpoint_status = {
        'augmented': {'exists': False, 'path': None, 'epochs': None},
        'finetuned': {'exists': False, 'path': None, 'epochs': None},
        'control': {'exists': False, 'path': None, 'epochs': None}
    }

    # Check each checkpoint type
    aug_path = os.path.join(output_dir, f"{expt_name}_aug.ckpt")
    if os.path.exists(aug_path):
        checkpoint_status['augmented']['exists'] = True
        # ... extract metadata

Resumption Logic:

if checkpoint_status['augmented']['exists']:
    print(f"✓ Found existing augmented model checkpoint")
    print("Skipping Stage 1 training - using existing model.")
    best_model_path = checkpoint_status['augmented']['path']
else:
    # Train new model
    trainer.fit(model, datamodule=data_module)

Performance Evaluation

Metrics Calculation:

# Get predictions
pred = utils.get_predictions(model, base_dataset.x_test, batch_size=batch_size)

# Calculate correlations
pearson_r = []
for class_index in range(y_true.shape[-1]):
    r = stats.pearsonr(y_true[:,class_index], y_score[:,class_index])[0]
    pearson_r.append(r)

spearman_r = []
for class_index in range(y_true.shape[-1]):
    r = stats.spearmanr(y_true[:,class_index], y_score[:,class_index])[0]
    spearman_r.append(r)

Results Storage:

metrics_data['augmented'] = {
    'pearson_r': pearson_aug,
    'spearman_r': spearman_r
}

Visualization

Comparison Plots:

def plot_metrics_comparison(metrics_data, plots_dir, expt_name):
    """Create comprehensive plots comparing metrics across model types."""

    # 1. Correlation Metrics Comparison (Bar Plot)
    # 2. Detailed Metrics by Class (Heatmap)
    # 3. Performance Improvement Analysis
    # 4. Individual Model Performance
    # 5. Summary Statistics Table

Generated Files:

  • {expt_name}_metrics_comparison.png - Overall performance comparison

  • {expt_name}_individual_performance.png - Individual model analysis

  • {expt_name}_performance_summary.png - Statistical summary table

Output Structure

Model Checkpoints:

output_dir/
├── DeepSTARR_aug.ckpt          # Stage 1: Augmented model
├── DeepSTARR_finetune.ckpt     # Stage 2: Fine-tuned model
└── DeepSTARR_standard.ckpt     # Control: Standard model

Plots and Results:

plots/
├── DeepSTARR_metrics_comparison.png
├── DeepSTARR_individual_performance.png
└── DeepSTARR_performance_summary.png

Training Logs:

  • PyTorch Lightning logs in lightning_logs/

  • Console output with training progress

  • Checkpoint status information

Customization

Modify Augmentation Strategy:

# Change augmentation types
augment_list = [
    RandomDeletion(delete_min=0, delete_max=30),
    RandomInsertion(insert_min=0, insert_max=20),
    RandomMutation(mut_frac=0.1),
]

# Adjust augmentation frequency
max_augs_per_seq = 3  # Apply up to 3 augmentations
hard_aug = False      # Stochastic application

Modify Training Parameters:

# Change learning rates
learning_rate = 0.0005      # Stage 1
finetune_lr = 0.00005      # Stage 2

# Adjust epochs
max_epochs = 50             # Stage 1
finetune_epochs = 10        # Stage 2

# Modify batch size
batch_size = 64             # Smaller for memory constraints

Custom Models:

# Use different model architecture
from evoaug_utils.model_zoo import YourCustomModel
model = YourCustomModel(num_classes=2)

Troubleshooting

Common Issues:

  1. Memory Errors: - Reduce batch_size - Use gradient accumulation - Enable mixed precision training

  2. Checkpoint Issues: - Verify output directory permissions - Check available disk space - Ensure consistent model architecture

  3. Data Loading Problems: - Verify data file path - Check data format compatibility - Ensure sufficient memory for dataset

Performance Tips:

  • Use GPU acceleration when available

  • Enable mixed precision training

  • Monitor memory usage during training

  • Use appropriate batch sizes for your hardware

Next Steps

After running this example:

  1. Analyze Results: Review generated plots and metrics

  2. Experiment: Modify parameters and compare results

  3. Scale Up: Apply to larger datasets or different models

  4. Customize: Adapt for your specific use case

Further Learning:

  • Read the user_guide/training for detailed training explanations

  • Explore the api/evoaug for API reference

  • Check the examples/vanilla_pytorch for PyTorch-only approach

  • Review the user_guide/augmentations for augmentation details

This example provides a production-ready template for implementing EvoAug2 in PyTorch Lightning workflows and can serve as a foundation for your own genomic sequence analysis projects.