Trainer

Trainer#

class braintools.trainer.Trainer(max_epochs=1000, min_epochs=1, max_steps=-1, val_check_interval=1.0, check_val_every_n_epoch=1, callbacks=None, logger=True, enable_progress_bar=True, enable_checkpointing=True, default_root_dir=None, gradient_clip_val=None, gradient_clip_algorithm='norm', accumulate_grad_batches=1, devices='auto', strategy='auto', precision='32', deterministic=False, benchmark=False, seed=None)#

Orchestrates the training process.

The Trainer handles the training loop, validation, testing, and prediction, integrating callbacks, logging, checkpointing, and distributed training.

Parameters:

  • max_epochs (int) – Maximum number of training epochs.

  • min_epochs (int) – Minimum number of training epochs.

  • max_steps (int) – Maximum number of training steps. -1 means no limit. A “step” is counted per processed batch; when accumulate_grad_batches > 1 this counts micro-batches rather than optimizer updates.

  • val_check_interval (int | float) – How often to run validation within a training epoch. Integer = every N batches, float = fraction of epoch.

  • check_val_every_n_epoch (int) – Run validation every N epochs.

  • callbacks (List[Callback] | None) – List of callbacks to use.

  • logger (Logger | List[Logger] | bool) – Logger(s) to use. True = CSVLogger, False = no logging.

  • enable_progress_bar (bool) – Whether to show progress bars.

  • enable_checkpointing (bool) – Whether to enable automatic checkpointing.

  • default_root_dir (str | None) – Default root directory for logs and checkpoints.

  • gradient_clip_val (float | None) – Value for gradient clipping.

  • gradient_clip_algorithm (str) – Gradient clipping algorithm (‘norm’ or ‘value’).

  • accumulate_grad_batches (int) – Number of batches to accumulate gradients over before each optimizer step. Must be >= 1.

  • devices (int | List[int] | str) – Devices to use for training. 'auto' uses all visible JAX devices; an int selects the first N; a list selects by index.

  • strategy (str | Strategy) – Distributed training strategy. The active training loop drives a single optimizer on the selected device(s); multi-host strategies are provided but not exercised by fit().

  • precision (str) – Requested compute precision ('32', '16', 'bf16', optionally with a '-mixed'/'-true' suffix). Recorded for API compatibility; the loop currently computes in float32 and emits a warning for any non-32 value. Cast your model/data explicitly for reduced precision.

  • deterministic (bool) – Recorded for API compatibility. Determinism in JAX is governed by the PRNG keys you use and XLA flags; this flag does not itself change execution.

  • benchmark (bool) – Recorded for API compatibility; has no effect (XLA autotunes and there is no cuDNN benchmark switch in JAX). A warning is emitted if set.

  • seed (int | None) – Random seed for reproducibility. Seeds both NumPy and brainstate.

Notes

Only the first optimizer returned by LightningModule.configure_optimizers() is stepped by the training loop; configuring several emits a warning. Freeze parameters with LightningModule.freeze() before calling fit() to exclude them from gradients and optimizer updates.

Examples

Basic usage:

>>> trainer = Trainer(max_epochs=10)
>>> trainer.fit(model, train_loader, val_loader)

With callbacks and logging:

>>> trainer = Trainer(
...     max_epochs=100,
...     callbacks=[
...         ModelCheckpoint(dirpath='checkpoints/', monitor='val_loss'),
...         EarlyStopping(monitor='val_loss', patience=5),
...     ],
...     logger=TensorBoardLogger('logs/'),
... )
>>> trainer.fit(model, train_loader, val_loader)
property callbacks: List[Callback]#

List of callbacks.

property current_epoch: int#

Current epoch.

fit(model, train_dataloaders=None, val_dataloaders=None, ckpt_path=None)[source]#

Run the full training loop.

Parameters:

Examples

>>> trainer.fit(model, train_loader, val_loader)
property global_step: int#

Global step count.

property is_training: bool#

Whether currently in training stage.

predict(model=None, dataloaders=None, ckpt_path=None)[source]#

Run prediction.

Parameters:
Returns:

Predictions for each batch.

Return type:

List[Any]

test(model=None, dataloaders=None, ckpt_path=None, verbose=True)[source]#

Run testing.

Parameters:
Returns:

Test metrics.

Return type:

Dict[str, Any]

validate(model=None, dataloaders=None, ckpt_path=None, verbose=True)[source]#

Run validation.

Parameters:
Returns:

Validation metrics.

Return type:

Dict[str, Any]