select-trainer

# Configure Training for a PINA Solver > [!IMPORTANT] > Read [RULES.md](../RULES.md) before using this skill — it applies to all skills. Use this skill to set up and run training once a solver has been chosen. PINA's `Trainer` wraps Lightning's `Trainer` with PINA-specific defaults: data splitting across conditions, batching strategies, device placement for inverse parameters, and gradient tracking for physics-informed solvers. ## Step 1 — Understand what the trainer needs The trainer needs three things. If any are missing, ask the user: 1. **Solver** (required) — already chosen and instantiated. 2. **Domain discretisation** — all `DomainEquationCondition` domains must be sampled before the trainer can create dataloaders. 3. **Training config** — epochs, batch size, accelerator, device. ### Domain discretisation Physics-informed problems use `domain=...` in their conditions. The trainer requires every such domain to have been sampled: ```python problem.discretise_domain(n=1000, mode="random") ``` If you forget, the trainer raises a clear error listing which domains are missing. ## Step 2 — Configure the trainer The constructor has two groups of parameters: PINA-specific and Lightning `**kwargs` passed through to the parent class. ### PINA-specific parameters | Parameter | Default | What it controls | |---|---|---| | `solver` | — (required) | The solver instance to train | | `batch_size` | `None` | `None` = full batch; int = mini-batches | | `train_size` / `val_size` / `test_size` | `1.0 / 0.0 / 0.0` | Fraction split per condition | | `batching_mode` | `"common_batch_size"` | How batches are built across conditions (see below) | | `automatic_batching` | `False` | `True` = Lightning's default collation; `False` = direct subset retrieval | | `num_workers` | `0` | Dataloader workers | | `pin_memory` | `False` | Pin memory for faster GPU transfer | | `shuffle` | `True` | Shuffle before splitting | ### Batching modes The mode controls how condition data is assembled into batches: | Mode | Behaviour | When to use | |---|---|---| | `"common_batch_size"` | Each condition supplies `batch_size` points per batch | Default — works for most cases | | `"proportional"` | Batch sizes are scaled by condition dataset sizes | Unbalanced datasets (e.g., many interior points but few boundary points) | | `"separate_conditions"` | Iterates through each condition separately | Each condition's data is heterogeneous (e.g., one is pointwise, another is a graph) | ### Common Lightning kwargs These are passed as `**kwargs` and fully documented by PyTorch Lightning. Key ones for PINA users: | Kwarg | What it does | |---|---| | `max_epochs` | Number of training epochs | | `accelerator` | `"cpu"`, `"gpu"`, `"mps"` (Apple Silicon) | | `devices` | Device index or count (e.g., `1`, `[0]`, `"auto"`) | | `precision` | `"16-mixed"`, `"32"`, `"64"`, `"bf16-mixed"` | | `enable_progress_bar` | `True` / `False` | | `gradient_clip_val` | Gradient clipping threshold | | `callbacks` | List of Lightning callbacks (e.g., `ModelCheckpoint`, `EarlyStopping`) | ## Step 3 — Train and test Usage is uniform regardless of solver type: ```python from pina import Trainer trainer = Trainer( solver=solver, max_epochs=1000, batch_size=32, accelerator="cpu", train_size=0.8, val_size=0.1, test_size=0.1, ) trainer.train() # Lightning fit() trainer.test() # Lightning test() — optional ``` ## Templates ### Full-batch physics-informed ```python problem.discretise_domain(n=5000, mode="random") trainer = Trainer( solver=solver, max_epochs=10000, batch_size=None, # full batch — no mini-batching accelerator="cpu", ) trainer.train() ``` ### Mini-batch supervised ```python trainer = Trainer( solver=solver, max_epochs=500, batch_size=64, train_size=0.8, val_size=0.1, test_size=0.1, accelerator="gpu", devices=1, num_workers=4, shuffle=True, ) trainer.train() trainer.test() ``` ### Unbalanced conditions (proportional batching) ```python trainer = Trainer( solver=solver, batch_size=512, batching_mode="proportional", max_epochs=1000, ) ``` ## Checklist - [ ] Solver already chosen and instantiated (see **select-solver** skill) - [ ] All domains discretised: `problem.discretise_domain(n=..., mode="...")` - [ ] Batch size decided: `None` for full-batch, an int for mini-batches - [ ] Batching mode selected based on condition balance - [ ] Train/val/test split configured - [ ] Accelerator and device chosen - [ ] Training started: `trainer.train()`