ma-agents 3.6.2 → 3.8.0

package/lib/bmad-extension-plugin/skills/ml-experiment/assets/quick_trainer_setup.py

@@ -0,0 +1,233 @@
+"""
+quick_trainer_setup.py — BMAD DL Lifecycle
+Ready-to-run Lightning Trainer configuration for standard DL training runs.
+
+Covers: callbacks (early stopping, checkpointing, LR monitor),
+        loggers (CSV + optional TensorBoard/W&B), and hardware-aware device selection.
+
+Usage:
+    python3 assets/quick_trainer_setup.py          # prints recommended config
+    python3 assets/quick_trainer_setup.py --run    # launches a training run (demo)
+
+Or import and call build_trainer() in your training script:
+
+    from assets.quick_trainer_setup import build_trainer
+    from src.models.your_model import YourModel
+    from src.data.your_datamodule import YourDataModule
+
+    trainer = build_trainer(max_epochs=50, experiment_name="run_001")
+    model = YourModel(num_classes=2)
+    dm = YourDataModule(data_dir="data/")
+    trainer.fit(model, dm)
+    trainer.test(model, dm)
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+
+# ── Lightning import ──────────────────────────────────────────────────────────
+
+try:
+    import lightning as L
+    from lightning.pytorch.callbacks import (
+        EarlyStopping, ModelCheckpoint, LearningRateMonitor, RichProgressBar,
+    )
+    from lightning.pytorch.loggers import CSVLogger
+    LIGHTNING_PKG = "lightning"
+except ImportError:
+    try:
+        import pytorch_lightning as L
+        from pytorch_lightning.callbacks import (
+            EarlyStopping, ModelCheckpoint, LearningRateMonitor,
+        )
+        from pytorch_lightning.loggers import CSVLogger
+        RichProgressBar = None
+        LIGHTNING_PKG = "pytorch_lightning"
+    except ImportError:
+        L = None  # type: ignore
+        LIGHTNING_PKG = None
+
+
+def _detect_accelerator() -> tuple[str, int]:
+    """Return (accelerator, devices) based on available hardware, with explicit status output."""
+    try:
+        import torch
+        if torch.cuda.is_available():
+            device_name = torch.cuda.get_device_name(0)
+            vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+            print(f"GPU: {device_name} ({vram_gb:.1f} GB VRAM) — using CUDA")
+            return "gpu", torch.cuda.device_count()
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            print("GPU: Apple MPS — using Metal Performance Shaders")
+            return "mps", 1
+    except ImportError:
+        pass
+    print("WARNING: No GPU detected — training will run on CPU and be significantly slower.")
+    print("         If you expected a GPU, check your CUDA installation and driver.")
+    return "cpu", 1
+
+
+def build_trainer(
+    max_epochs: int = 50,
+    experiment_name: str = "experiment",
+    version: str | None = None,
+    log_dir: str | Path = "logs/",
+    monitor_metric: str = "val/loss",
+    monitor_mode: str = "min",
+    early_stopping_patience: int = 10,
+    gradient_clip_val: float = 1.0,
+    accumulate_grad_batches: int = 1,
+    precision: str = "16-mixed",
+) -> "L.Trainer":
+    """
+    Build a Lightning Trainer with standard callbacks and logging.
+
+    Args:
+        max_epochs: Maximum training epochs.
+        experiment_name: Name used for checkpoint dir and log subdir.
+        version: Run identifier appended to log path (e.g. "fold_0", "run_001").
+                 Prevents different runs from overwriting each other's TensorBoard logs.
+        log_dir: Root directory for logs and checkpoints.
+        monitor_metric: Metric to monitor for early stopping and checkpointing.
+        monitor_mode: "min" (for loss) or "max" (for accuracy/F1).
+        early_stopping_patience: Stop after N epochs without improvement.
+        gradient_clip_val: Max gradient norm (0.0 to disable clipping).
+        accumulate_grad_batches: Simulate larger batch size via gradient accumulation.
+        precision: Training precision ("32", "16-mixed", "bf16-mixed").
+
+    Returns:
+        Configured Lightning Trainer.
+    """
+    if L is None:
+        raise ImportError(
+            "PyTorch Lightning not installed.\n"
+            "  pip install lightning   (recommended)\n"
+            "  or: pip install pytorch-lightning"
+        )
+
+    log_dir = Path(log_dir)
+    ckpt_dir = log_dir / "checkpoints" / experiment_name
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+    accelerator, devices = _detect_accelerator()
+
+    # ── Callbacks ─────────────────────────────────────────────────────────────
+    callbacks = [
+        ModelCheckpoint(
+            dirpath=ckpt_dir,
+            filename=f"{experiment_name}-{{epoch:02d}}-{{{monitor_metric}:.4f}}",
+            monitor=monitor_metric,
+            mode=monitor_mode,
+            save_top_k=3,
+            save_last=True,
+            verbose=True,
+        ),
+        EarlyStopping(
+            monitor=monitor_metric,
+            mode=monitor_mode,
+            patience=early_stopping_patience,
+            verbose=True,
+        ),
+        LearningRateMonitor(logging_interval="epoch"),
+    ]
+    if RichProgressBar is not None:
+        callbacks.append(RichProgressBar())
+
+    # ── Loggers ───────────────────────────────────────────────────────────────
+    # TensorBoard is required — install with: uv add tensorboard
+    # version= keeps each fold/run in its own subdir so they never overwrite each other
+    if LIGHTNING_PKG == "lightning":
+        from lightning.pytorch.loggers import TensorBoardLogger
+    else:
+        from pytorch_lightning.loggers import TensorBoardLogger
+
+    tb_logger = TensorBoardLogger(save_dir=str(log_dir), name=experiment_name, version=version)
+    csv_logger = CSVLogger(save_dir=str(log_dir), name=experiment_name, version=version)
+    loggers = [tb_logger, csv_logger]
+
+    log_path = Path(log_dir) / experiment_name / (version or f"version_{tb_logger.version}")
+    print(f"Logs → {log_path}/   run: tensorboard --logdir={log_dir}")
+
+    # ── Precision ─────────────────────────────────────────────────────────────
+    # Fall back to 32-bit on CPU (mixed precision not supported)
+    if accelerator == "cpu" and precision != "32":
+        precision = "32"
+
+    trainer = L.Trainer(
+        max_epochs=max_epochs,
+        accelerator=accelerator,
+        devices=devices,
+        precision=precision,
+        gradient_clip_val=gradient_clip_val if gradient_clip_val > 0 else None,
+        accumulate_grad_batches=accumulate_grad_batches,
+        callbacks=callbacks,
+        logger=loggers,
+        log_every_n_steps=10,
+        deterministic=False,  # set True for full reproducibility (slower)
+    )
+
+    return trainer
+
+
+def print_config(max_epochs: int, experiment_name: str, log_dir: str) -> None:
+    accelerator, devices = _detect_accelerator()
+    print(f"""
+┌─────────────────────────────────────────────────────┐
+│  BMAD DL — Quick Trainer Configuration              │
+├─────────────────────────────────────────────────────┤
+│  Lightning package : {LIGHTNING_PKG or 'NOT INSTALLED':<30} │
+│  Hardware          : {accelerator.upper()} ({devices} device(s)){'':<19} │
+│  Max epochs        : {max_epochs:<30} │
+│  Experiment name   : {experiment_name:<30} │
+│  Log directory     : {log_dir:<30} │
+├─────────────────────────────────────────────────────┤
+│  Callbacks active:                                  │
+│    ✓ ModelCheckpoint  (top-3 + last)                │
+│    ✓ EarlyStopping    (patience=10)                 │
+│    ✓ LearningRateMonitor                            │
+│    ✓ RichProgressBar  (if available)                │
+│  Loggers active:                                    │
+│    ✓ CSVLogger                                      │
+│    ✓ TensorBoardLogger (if tensorboard installed)   │
+└─────────────────────────────────────────────────────┘
+
+  Quick start in your training script:
+
+    from assets.quick_trainer_setup import build_trainer
+    trainer = build_trainer(max_epochs=50, experiment_name="run_001")
+    trainer.fit(model, datamodule)
+    trainer.test(model, datamodule)
+
+  After training, parse results with:
+
+    python3 scripts/parse_training_logs.py \\
+        logs/{experiment_name}/version_0/metrics.csv \\
+        docs/prd/01_PRD.md
+""")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Quick Trainer Setup — BMAD DL Lifecycle")
+    parser.add_argument("--run", action="store_true", help="Launch a demo training run")
+    parser.add_argument("--max-epochs", type=int, default=50)
+    parser.add_argument("--experiment-name", type=str, default="run_001")
+    parser.add_argument("--log-dir", type=str, default="logs/")
+    args = parser.parse_args()
+
+    if args.run:
+        if L is None:
+            print("Error: PyTorch Lightning not installed.", file=sys.stderr)
+            return 2
+        print("Demo run requires a model and datamodule. See module docstring.")
+        return 1
+
+    print_config(args.max_epochs, args.experiment_name, args.log_dir)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

package/lib/bmad-extension-plugin/skills/ml-experiment/assets/template_datamodule.py

@@ -0,0 +1,219 @@
+"""
+template_datamodule.py — BMAD DL Lifecycle
+PyTorch Lightning LightningDataModule template.
+
+Handles train/val/test dataset loading, transforms, and DataLoader creation
+in a clean, reproducible, and Lightning-compatible way.
+
+Usage:
+    Copy to src/data/your_datamodule.py and implement the TODO sections.
+    Then pass it directly to the Trainer — no manual DataLoaders needed.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import torch
+from torch.utils.data import DataLoader, Dataset, random_split
+
+try:
+    import lightning as L
+    LightningDataModule = L.LightningDataModule
+except ImportError:
+    try:
+        import pytorch_lightning as pl
+        LightningDataModule = pl.LightningDataModule
+    except ImportError:
+        raise ImportError("Install PyTorch Lightning: pip install lightning")
+
+try:
+    from torchvision import transforms
+    HAS_TORCHVISION = True
+except ImportError:
+    HAS_TORCHVISION = False
+
+
+# ── TODO: Define or import your Dataset ───────────────────────────────────────
+# Replace this stub with your actual Dataset class.
+
+class YourDataset(Dataset):
+    """
+    Stub dataset — replace with your implementation.
+
+    Expected output per __getitem__: (input_tensor, label_tensor)
+    """
+
+    def __init__(self, data_dir: Path, split: str = "train", transform=None):
+        self.data_dir = data_dir
+        self.split = split
+        self.transform = transform
+
+        # TODO: load file list, annotations, CSV rows, etc.
+        self.samples: list = []   # list of (path_or_data, label)
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int):
+        sample, label = self.samples[idx]
+        # TODO: load image/array/features from `sample`
+        # x = Image.open(sample).convert("RGB")
+        # if self.transform:
+        #     x = self.transform(x)
+        # return x, label
+        raise NotImplementedError("Implement __getitem__ in your Dataset")
+
+# ── END TODO ──────────────────────────────────────────────────────────────────
+
+
+class YourDataModule(LightningDataModule):
+    """
+    Template LightningDataModule.
+
+    Replace 'YourDataModule' with a descriptive name (e.g. DefectDataModule).
+
+    Args:
+        data_dir: Root directory of your dataset.
+        batch_size: Batch size for all DataLoaders.
+        num_workers: Number of worker processes for data loading.
+        val_split: Fraction of training data to use for validation
+                   (only used when no explicit val/ directory exists).
+        seed: Random seed for reproducibility.
+        image_size: (H, W) for image resizing — set None to skip.
+    """
+
+    def __init__(
+        self,
+        data_dir: str | Path = "data/",
+        batch_size: int = 32,
+        num_workers: int = 4,
+        val_split: float = 0.15,
+        seed: int = 42,
+        image_size: tuple[int, int] | None = (224, 224),
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        self.data_dir = Path(data_dir)
+
+        # Built in setup()
+        self.train_dataset: Optional[Dataset] = None
+        self.val_dataset: Optional[Dataset] = None
+        self.test_dataset: Optional[Dataset] = None
+
+    # ── Transforms ────────────────────────────────────────────────────────────
+
+    def _train_transform(self):
+        """
+        TODO: Define augmentation pipeline for training.
+        """
+        if not HAS_TORCHVISION:
+            return None
+        steps = []
+        if self.hparams.image_size:
+            steps.append(transforms.Resize(self.hparams.image_size))
+        steps += [
+            transforms.RandomHorizontalFlip(),
+            transforms.RandomRotation(10),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225]),
+        ]
+        return transforms.Compose(steps)
+
+    def _eval_transform(self):
+        """
+        TODO: Define deterministic transform for val/test (no augmentation).
+        """
+        if not HAS_TORCHVISION:
+            return None
+        steps = []
+        if self.hparams.image_size:
+            steps.append(transforms.Resize(self.hparams.image_size))
+        steps += [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225]),
+        ]
+        return transforms.Compose(steps)
+
+    # ── Setup ─────────────────────────────────────────────────────────────────
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        """
+        Called by Lightning before fit/test. Initializes dataset splits.
+
+        stage: "fit" (train+val), "test", "predict", or None (all).
+        """
+        # TODO: Adjust split detection logic for your directory layout.
+        # Option A — explicit split directories: data/train/, data/val/, data/test/
+        has_split_dirs = (
+            (self.data_dir / "train").exists() and
+            (self.data_dir / "val").exists()
+        )
+
+        if stage in (None, "fit"):
+            if has_split_dirs:
+                self.train_dataset = YourDataset(
+                    self.data_dir / "train", split="train",
+                    transform=self._train_transform(),
+                )
+                self.val_dataset = YourDataset(
+                    self.data_dir / "val", split="val",
+                    transform=self._eval_transform(),
+                )
+            else:
+                # Option B — random split from single dataset directory
+                full_dataset = YourDataset(
+                    self.data_dir, split="train",
+                    transform=self._train_transform(),
+                )
+                val_size = int(len(full_dataset) * self.hparams.val_split)
+                train_size = len(full_dataset) - val_size
+                self.train_dataset, self.val_dataset = random_split(
+                    full_dataset,
+                    [train_size, val_size],
+                    generator=torch.Generator().manual_seed(self.hparams.seed),
+                )
+
+        if stage in (None, "test"):
+            test_dir = self.data_dir / "test"
+            if test_dir.exists():
+                self.test_dataset = YourDataset(
+                    test_dir, split="test",
+                    transform=self._eval_transform(),
+                )
+
+    # ── DataLoaders ───────────────────────────────────────────────────────────
+
+    def train_dataloader(self) -> DataLoader:
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=True,
+            num_workers=self.hparams.num_workers,
+            pin_memory=True,
+            drop_last=True,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=False,
+            num_workers=self.hparams.num_workers,
+            pin_memory=True,
+        )
+
+    def test_dataloader(self) -> DataLoader:
+        if self.test_dataset is None:
+            raise RuntimeError("No test dataset found. Check data_dir/test/ exists.")
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=False,
+            num_workers=self.hparams.num_workers,
+            pin_memory=True,
+        )