Add more scientific skills

scientific-packages/pytorch-lightning/scripts/quick_trainer_setup.py

@@ -0,0 +1,262 @@
+"""
+Helper script to quickly set up a PyTorch Lightning Trainer with common configurations.
+
+This script provides preset configurations for different training scenarios
+and makes it easy to create a Trainer with best practices.
+"""
+
+import lightning as L
+from lightning.pytorch.callbacks import (
+    ModelCheckpoint,
+    EarlyStopping,
+    LearningRateMonitor,
+    RichProgressBar,
+    ModelSummary,
+)
+from lightning.pytorch.loggers import TensorBoardLogger, CSVLogger
+
+
+def create_trainer(
+    preset: str = "default",
+    max_epochs: int = 100,
+    accelerator: str = "auto",
+    devices: int = 1,
+    log_dir: str = "./logs",
+    experiment_name: str = "lightning_experiment",
+    enable_checkpointing: bool = True,
+    enable_early_stopping: bool = True,
+    **kwargs
+):
+    """
+    Create a Lightning Trainer with preset configurations.
+
+    Args:
+        preset: Configuration preset - "default", "fast_dev", "production", "distributed"
+        max_epochs: Maximum number of training epochs
+        accelerator: Device to use ("auto", "gpu", "cpu", "tpu")
+        devices: Number of devices to use
+        log_dir: Directory for logs and checkpoints
+        experiment_name: Name for the experiment
+        enable_checkpointing: Whether to enable model checkpointing
+        enable_early_stopping: Whether to enable early stopping
+        **kwargs: Additional arguments to pass to Trainer
+
+    Returns:
+        Configured Lightning Trainer instance
+    """
+
+    callbacks = []
+    logger_list = []
+
+    # Configure based on preset
+    if preset == "fast_dev":
+        # Fast development run - minimal epochs, quick debugging
+        config = {
+            "fast_dev_run": False,
+            "max_epochs": 3,
+            "limit_train_batches": 100,
+            "limit_val_batches": 50,
+            "log_every_n_steps": 10,
+            "enable_progress_bar": True,
+            "enable_model_summary": True,
+        }
+
+    elif preset == "production":
+        # Production-ready configuration with all bells and whistles
+        config = {
+            "max_epochs": max_epochs,
+            "precision": "16-mixed",
+            "gradient_clip_val": 1.0,
+            "log_every_n_steps": 50,
+            "enable_progress_bar": True,
+            "enable_model_summary": True,
+            "deterministic": True,
+            "benchmark": True,
+        }
+
+        # Add model checkpointing
+        if enable_checkpointing:
+            callbacks.append(
+                ModelCheckpoint(
+                    dirpath=f"{log_dir}/{experiment_name}/checkpoints",
+                    filename="{epoch}-{val_loss:.2f}",
+                    monitor="val_loss",
+                    mode="min",
+                    save_top_k=3,
+                    save_last=True,
+                    verbose=True,
+                )
+            )
+
+        # Add early stopping
+        if enable_early_stopping:
+            callbacks.append(
+                EarlyStopping(
+                    monitor="val_loss",
+                    patience=10,
+                    mode="min",
+                    verbose=True,
+                )
+            )
+
+        # Add learning rate monitor
+        callbacks.append(LearningRateMonitor(logging_interval="epoch"))
+
+        # Add TensorBoard logger
+        logger_list.append(
+            TensorBoardLogger(
+                save_dir=log_dir,
+                name=experiment_name,
+                version=None,
+            )
+        )
+
+    elif preset == "distributed":
+        # Distributed training configuration
+        config = {
+            "max_epochs": max_epochs,
+            "strategy": "ddp",
+            "precision": "16-mixed",
+            "sync_batchnorm": True,
+            "use_distributed_sampler": True,
+            "log_every_n_steps": 50,
+            "enable_progress_bar": True,
+        }
+
+        # Add model checkpointing
+        if enable_checkpointing:
+            callbacks.append(
+                ModelCheckpoint(
+                    dirpath=f"{log_dir}/{experiment_name}/checkpoints",
+                    filename="{epoch}-{val_loss:.2f}",
+                    monitor="val_loss",
+                    mode="min",
+                    save_top_k=3,
+                    save_last=True,
+                )
+            )
+
+    else:  # default
+        # Default configuration - balanced for most use cases
+        config = {
+            "max_epochs": max_epochs,
+            "log_every_n_steps": 50,
+            "enable_progress_bar": True,
+            "enable_model_summary": True,
+        }
+
+        # Add basic checkpointing
+        if enable_checkpointing:
+            callbacks.append(
+                ModelCheckpoint(
+                    dirpath=f"{log_dir}/{experiment_name}/checkpoints",
+                    filename="{epoch}-{val_loss:.2f}",
+                    monitor="val_loss",
+                    save_last=True,
+                )
+            )
+
+        # Add CSV logger
+        logger_list.append(
+            CSVLogger(
+                save_dir=log_dir,
+                name=experiment_name,
+            )
+        )
+
+    # Add progress bar
+    if config.get("enable_progress_bar", True):
+        callbacks.append(RichProgressBar())
+
+    # Merge with provided kwargs
+    final_config = {
+        **config,
+        "accelerator": accelerator,
+        "devices": devices,
+        "callbacks": callbacks,
+        "logger": logger_list if logger_list else True,
+        **kwargs,
+    }
+
+    # Create and return trainer
+    return L.Trainer(**final_config)
+
+
+def create_debugging_trainer():
+    """Create a trainer optimized for debugging."""
+    return create_trainer(
+        preset="fast_dev",
+        max_epochs=1,
+        limit_train_batches=10,
+        limit_val_batches=5,
+        num_sanity_val_steps=2,
+    )
+
+
+def create_gpu_trainer(num_gpus: int = 1, precision: str = "16-mixed"):
+    """Create a trainer optimized for GPU training."""
+    return create_trainer(
+        preset="production",
+        accelerator="gpu",
+        devices=num_gpus,
+        precision=precision,
+    )
+
+
+def create_distributed_trainer(num_gpus: int = 2, num_nodes: int = 1):
+    """Create a trainer for distributed training across multiple GPUs."""
+    return create_trainer(
+        preset="distributed",
+        accelerator="gpu",
+        devices=num_gpus,
+        num_nodes=num_nodes,
+        strategy="ddp",
+    )
+
+
+# Example usage
+if __name__ == "__main__":
+    print("Creating different trainer configurations...\n")
+
+    # 1. Default trainer
+    print("1. Default trainer:")
+    trainer_default = create_trainer(preset="default", max_epochs=50)
+    print(f"   Max epochs: {trainer_default.max_epochs}")
+    print(f"   Accelerator: {trainer_default.accelerator}")
+    print(f"   Callbacks: {len(trainer_default.callbacks)}")
+    print()
+
+    # 2. Fast development trainer
+    print("2. Fast development trainer:")
+    trainer_dev = create_trainer(preset="fast_dev")
+    print(f"   Max epochs: {trainer_dev.max_epochs}")
+    print(f"   Train batches limit: {trainer_dev.limit_train_batches}")
+    print()
+
+    # 3. Production trainer
+    print("3. Production trainer:")
+    trainer_prod = create_trainer(
+        preset="production",
+        max_epochs=100,
+        experiment_name="my_experiment"
+    )
+    print(f"   Max epochs: {trainer_prod.max_epochs}")
+    print(f"   Precision: {trainer_prod.precision}")
+    print(f"   Callbacks: {len(trainer_prod.callbacks)}")
+    print()
+
+    # 4. Debugging trainer
+    print("4. Debugging trainer:")
+    trainer_debug = create_debugging_trainer()
+    print(f"   Max epochs: {trainer_debug.max_epochs}")
+    print(f"   Train batches: {trainer_debug.limit_train_batches}")
+    print()
+
+    # 5. GPU trainer
+    print("5. GPU trainer:")
+    trainer_gpu = create_gpu_trainer(num_gpus=1)
+    print(f"   Accelerator: {trainer_gpu.accelerator}")
+    print(f"   Precision: {trainer_gpu.precision}")
+    print()
+
+    print("All trainer configurations created successfully!")