---
name: zyplib-train
description: Use when an agent needs to set up a training loop, configure metrics/callbacks, customize training steps using `zyplib.train`. Trigger when the user asks to train a model, configure checkpointing/logging, add custom metrics, write custom training steps, or integrate a new training paradigm.
---

# zyplib-train

Set up and run training from the agent's perspective: the user gives a training requirement, and you must assemble the correct `zyplib.train` API combination.

Assume the caller is using a released/importable `zyplib` package. Prefer the public API exposed from `zyplib.train`, and only dive into internals when you need customization details.

The source code of zyplib.train can be viewed in this way:
- First locate the installed/importable `zyplib` package root.

  ```bash
  python -c "import zyplib; from pathlib import Path; print(Path(zyplib.__file__).resolve().parent)"
  ```
- Then use Read/View tools on `<zyplib>/train/` in this order.
  - Start with `__init__.py` for the public API surface exported as `zyplib.train`.
  - Then read `common/__init__.py` for `default_fabric` and common exports.
  - Focus on `common/procedure.py` for `train_loop`, `train_epoch`, `training_step`, `validation_step`.
  - Focus on `common/_base.py` for `FabricTrainingContext` dataclass.
  - Focus on `common/trainer.py` for `BasicFabricTrainer` (alternative OOP style).
  - Focus on `common/callbacks.py` for all built-in callbacks and `FabricCallbackProtocol`.
  - Focus on `metric/torch_metric.py` for `TorchMetricRecorder` and `ScalarMetric`.
  - Focus on `metric/__init__.py` for `clf_metrics_builder` and `reg_metrics_builder`.

## Understand the architecture

`zyplib.train` is organized in 4 layers, assembled bottom-up:

```
Layer 1: Fabric Setup
  default_fabric(...)  ->  Fabric with loggers + callbacks attached

Layer 2: Training Context
  FabricTrainingContext(model, optimizer, fabric, fn_loss, train_metric_recorder, val_metric_recorder)
  Auto-setup: fabric.setup(model), fabric.setup_optimizers(optimizer) on __post_init__

Layer 3: Training Loop
  procedure.train_loop(context, train_loader, max_epochs, val_loader, ...)
  Internally calls: train_epoch -> training_step (per batch) -> metric_recorder.update

Layer 4: Composable Functions (customization points)
  fn_train_step, fn_val_step       -- replace the step logic
  fn_batch2target                  -- extract y from non-standard batches
  fn_transform_batch               -- data augmentation (Mixup, CutMix, TTA)
  fn_on_trainepoch_end             -- custom epoch-end logic
```

Think in this order:
1. **What infrastructure?** -> `default_fabric(...)` to get a Fabric with logging + checkpointing.
2. **What context?** -> `FabricTrainingContext` bundles model, optimizer, loss, metric recorders.
3. **What loop?** -> `procedure.train_loop(...)` runs the epoch loop with callbacks.
4. **What customization?** -> Inject custom step functions, transforms, or paradigm procedures.

## Use the framework (procedure-style)

The recommended primary style is procedure-based. Minimal complete skeleton:

```python
from zyplib.train import default_fabric, FabricTrainingContext, procedure, TorchMetricRecorder, clf_metrics_builder

# 1. Fabric setup
fabric = default_fabric(dir_name='lightning_logs', keep_best_k=3)

# 2. Metric recorders (MANDATORY — train_loop crashes without them)
metrics_fn = lambda: clf_metrics_builder(binary=True)
train_recorder = TorchMetricRecorder(metrics_fn(), before_metric='sigmoid')
val_recorder   = TorchMetricRecorder(metrics_fn(), before_metric='sigmoid')

# 3. Training context (auto-setups model + optimizer via fabric)
ctx = FabricTrainingContext(
    model=model,
    optimizer=optimizer,
    fabric=fabric,
    fn_loss=nn.BCEWithLogitsLoss(),
    train_metric_recorder=train_recorder,
    val_metric_recorder=val_recorder,
)
ctx.launch()  # IMPORTANT: use ctx.launch(), not fabric.launch()

# 4. Train
procedure.train_loop(ctx, train_loader, max_epochs=100, val_loader=val_loader)

# 5. After training
ctx.load_best_ckpt('replace')
scores = procedure.test(ctx, test_loader)
```

**Alternative: Trainer style** — `BasicFabricTrainer` wraps the same procedure functions with OOP overrides. Use when you prefer `trainer.fit()` / subclass overrides. See `references/examples.md` Scenario C.

## Configure metrics

`TorchMetricRecorder` combines two metric types:
- **Computed metrics** (torchmetrics): Accuracy, F1, etc. — built from `clf_metrics_builder()` or `reg_metrics_builder()`
- **Scalar metrics**: values returned by step functions (e.g. `loss`) — recorded via `scalar_metrics` parameter

Key configuration:

```python
recorder = TorchMetricRecorder(
    metrics=clf_metrics_builder(binary=True, use_auc=True),  # torchmetrics
    before_metric='sigmoid',   # activation before metric computation
    scalar_metrics='loss',     # record 'loss' from step output dict
    # scalar_metrics='auto',   # auto-record ALL keys from step output
    # scalar_metrics=['loss', 'g_loss', 'd_loss'],  # explicit list
)
```

`before_metric` rules:
- Binary clf + `BCEWithLogitsLoss` -> `'sigmoid'`
- Multiclass + `CrossEntropyLoss` -> `'softmax'`
- Custom -> pass a callable `(Tensor) -> Tensor`
- Per-metric -> pass a dict with glob patterns: `{'accuracy': torch.sigmoid, 'U*': None}`

## Configure side effects

`default_fabric(...)` creates a Fabric with built-in loggers and callbacks. Defaults:

| Side effect | Default | Key |
|-------------|---------|-----|
| TensorBoard logger | ON | `tensorboard_logger` |
| CSV logger | OFF | `csv_logger` |
| Print training info | ON | `print_info_cb` |
| Save checkpoints | ON | `save_checkpoints_cb` |
| Save hparams | ON | `save_hparams_cb` |
| Save train info to file | ON | `save_train_info` |

Toggle via `enabled` dict:
```python
# Disable everything, then enable only print
fabric = default_fabric(enabled={'*': False, 'print_info_cb': True})

# Disable only checkpoints
fabric = default_fabric(enabled={'save_checkpoints_cb': False})
```

Checkpoint config: `save_best=True`, `keep_best_k=3`, `monitor='loss'`, `monitor_on='val'`, `less_is_better=True`.

Add custom callbacks/loggers via `fn_custom_callback` / `fn_custom_logger`:
```python
fabric = default_fabric(
    fn_custom_callback=lambda version: [EarlyStopping(patience=15), LRSchedulerCallback(scheduler)],
)
```

## Customize training

### Step function contract

The core composability interface. Custom step functions must follow this signature:

```python
def my_train_step(
    ctx: FabricTrainingContext,
    batch: tuple[Tensor, ...],
    *,
    batch_idx: int,
    epoch_idx: int,
    ) -> tuple[Tensor, dict[str, float]]:
    """
    Returns:
      [0] y_pred: Tensor — for standard supervised tasks this is passed to
          TorchMetricRecorder.update(y_pred, y, ...) so computed metrics can
          compare prediction vs target. If the recorder only tracks scalar
          metrics and has no torchmetrics, this can be a placeholder such as None.
      [1] metrics: dict   — keys must match scalar_metrics names (default: 'loss')
    """
    x, y = batch
    ctx.optimizer.zero_grad()
    y_pred = ctx.model(x)
    loss = my_custom_loss(y_pred, y)
    ctx.fabric.backward(loss)
    ctx.optimizer.step()
    return y_pred, {'loss': loss.item()}
```

Wire it in:
```python
procedure.train_loop(ctx, train_loader, max_epochs=100, fn_train_step=my_train_step)
```

### Other composable hooks

- `fn_batch2target`: `(batch) -> Tensor` — extracts y from batch. Default: `batch[1]`. In ordinary supervised tasks this target is paired with `y_pred` inside `TorchMetricRecorder.update(y_pred, y, **raw_metrics)`. Override for GAN or multi-input batches; in scalar-only setups the extracted target may be effectively unused by the recorder. Accepted by both `train_loop` and `train_epoch`/`val_epoch`.
- `fn_transform_batch`: `(batch, batch_idx, epoch_idx) -> batch` — data augmentation hook (Mixup, CutMix). Accepted by `train_epoch`/`val_epoch` only (NOT `train_loop`). Also available as `ctx.transform_on_train_batch_start` override, which is the recommended way when using `train_loop`.
- `fn_on_trainepoch_end`: `(ctx, epoch, train_metrics, val_metrics) -> None` — custom epoch-end logic without writing a full callback. Accepted by `train_loop`.

### Paradigm extensions

For non-standard training (GAN, EDL, TML), the pattern is:
1. Write a Procedure class with `xxx_train_step` / `xxx_val_step` matching the step function contract
2. Wire via `fn_train_step=` / `fn_val_step=` parameters
3. Use `scalar_metrics='auto'` or explicit list for custom loss names

```python
from zyplib.paradigm.edl import BasicEDLFabricTrainProcedure
edl_proc = BasicEDLFabricTrainProcedure(config)
procedure.train_loop(ctx, train_loader, fn_train_step=edl_proc.evidential_train_step, fn_val_step=edl_proc.evidential_val_step)
```

See `references/paradigm-extensions.md` for GAN, EDL, and TML examples.

### Manual Fabric loop (escape hatch)

If `procedure.train_loop` doesn't fit your needs (e.g. multi-phase training, complex scheduling between epochs), you can drop down to calling `train_epoch` / `val_epoch` directly, or write a fully manual loop using only `fabric.setup()` + your own for-loop. In this case, read `common/procedure.py` source to understand what `train_loop` does internally (callback dispatch, early stopping checks, metric logging) and replicate the parts you need.

## After training

```python
# Load best checkpoint (replaces current model in-place)
ctx.load_best_ckpt('replace')

# Or get a new model copy
best_model = ctx.load_best_ckpt('return')

# Run test evaluation
scores = procedure.test(ctx, test_loader)
```

`CurveEvaluation` is available for post-hoc comparison of training runs — see `zyplib.train.CurveEvaluator`.

## Read references only when needed

Each reference file covers a specific concern. Read only the one you need:

- `references/examples.md` — **6 complete code skeletons**:
  - Scenario A: Standard binary classification with procedure (recommended pattern)
  - Scenario B: Multiclass classification with procedure
  - Scenario C: Trainer-style (BasicFabricTrainer)
  - Scenario D: Custom step function with multi-loss
  - Scenario E: K-fold training loop
  - Scenario F: After-training evaluation and best model loading

- `references/step-function-contract.md` — **Composable function interface spec**:
  - `fn_train_step` / `fn_val_step` exact signatures and return requirements
  - `fn_batch2target` signature and when to override
  - `fn_transform_batch` signature and augmentation patterns
  - `fn_on_trainepoch_end` signature
  - How step output flows into TorchMetricRecorder
  - How custom scalar_metrics names must match step output keys

- `references/metrics-and-recording.md` — **TorchMetricRecorder deep-dive**:
  - `before_metric` modes: string presets, callable, dict-glob per-metric
  - `scalar_metrics` modes: string, list, 'auto', ScalarMetric instance
  - `clf_metrics_builder` full parameters (task, threshold, num_classes, optional metrics)
  - `reg_metrics_builder` parameters
  - How to add custom torchmetrics.Metric instances

- `references/fabric-and-side-effects.md` — **default_fabric configuration**:
  - Full parameter list for `default_fabric`
  - `EnabledSideEffect` dict keys and defaults
  - `enabled` wildcard (`'*'`) behavior
  - Version naming and directory layout (`root_dir/dir_name/version/`)
  - `fn_custom_logger` / `fn_custom_callback` signatures

- `references/callbacks.md` — **All 7 built-in callbacks**:
  - `SaveCheckpoints`: params, Top-K strategy, monitor config, `best_checkpoint()` method
  - `PrintTrainInfo`: output format, file saving
  - `SaveHparams`: hparams.yaml, TensorBoard integration
  - `EarlyStopping`: patience, mode, monitor_on, callbacks
  - `LRSchedulerCallback`: interval (epoch/step/auto), ReduceLROnPlateau handling
  - `MonitorLearningRate`: LR logging
  - `FabricLogMetrics`: automatic metric logging to all loggers
  - `FabricCallbackProtocol`: full hook signature reference

- `references/paradigm-extensions.md` — **Extension pattern with examples**:
  - The general pattern: Procedure class + step function contract
  - GAN: `BasicGANFabricProcedure`, multi-optimizer, `fn_batch2target` override
  - EDL: `BasicEDLFabricTrainProcedure`, epoch-dependent loss, custom scalar_metrics
  - TML: `TMLFabricTrainingProcedure`, multi-view input, custom batch format

- `references/pitfalls.md` — **10 common mistakes and best-practice table**:
  - #1 Missing metric recorders (crash on train_epoch)
  - #2 Using fabric.launch() instead of ctx.launch() (device mismatch)
  - #3 Wrong before_metric for task type (garbage metrics)
  - #4 Double model/optimizer setup (calling fabric.setup manually)
  - #5 fn_loss resolution order (ctx.fn_loss vs parameter)
  - #6 scalar_metrics name mismatch with step output keys
  - #7 Forgetting val_metric_recorder for validation
  - #8 LR scheduler not registered as callback
  - #9 Checkpoint monitor metric name mismatch
  - #10 Not calling load_best_ckpt before test

## Finish with these checks

1. Are both `train_metric_recorder` and `val_metric_recorder` provided in FabricTrainingContext?
2. Is `ctx.launch()` called (not `fabric.launch()`) before training?
3. Does `before_metric` match the loss function? (sigmoid for BCE, softmax for CE)
4. Is `fn_loss` set on `ctx.fn_loss` or passed to `train_loop`?
5. Do custom step function return values match `scalar_metrics` names?
6. Is `seed_everything()` called for reproducibility?
7. Is `load_best_ckpt()` called before `procedure.test()`?
8. Are checkpoint `monitor` and `monitor_on` configured correctly for the tracked metric?
