LightningModule

class braintools.trainer.LightningModule(*args, **kwargs)

Base class for all training modules.

This class provides a PyTorch Lightning-like interface for defining models that work with the Trainer class. Users should subclass this and implement at minimum training_step and configure_optimizers.

Parameters:

None

trainer

Reference to the Trainer instance (set during training).

Type:

Trainer

current_epoch

Current training epoch (0-indexed).

Type:

int

global_step

Total number of training steps across all epochs.

Type:

int

logged_metrics

Metrics logged during the current step.

Type:

Dict[str, Any]

Examples

Basic usage:

>>> import braintools
>>> import brainstate
>>> import jax.numpy as jnp
>>>
>>> class MyModel(braintools.trainer.LightningModule):
...     def __init__(self, input_size, hidden_size, output_size):
...         super().__init__()
...         self.linear1 = brainstate.nn.Linear(input_size, hidden_size)
...         self.linear2 = brainstate.nn.Linear(hidden_size, output_size)
...
...     def __call__(self, x):
...         x = jax.nn.relu(self.linear1(x))
...         return self.linear2(x)
...
...     def training_step(self, batch, batch_idx):
...         x, y = batch['x'], batch['y']
...         logits = self(x)
...         loss = jnp.mean((logits - y) ** 2)
...         self.log('train_loss', loss)
...         return {'loss': loss}
...
...     def configure_optimizers(self):
...         return braintools.optim.Adam(lr=1e-3)

See also

Trainer

The training orchestration class.

configure_optimizers()

Configure optimizer(s) and learning rate scheduler(s).

Override this method to define your optimization setup.

Return type:

Optimizer | Tuple[Optimizer, LRScheduler] | Tuple[List[Optimizer], List[LRScheduler]] | Dict[str, Any]

Returns:

• Optimizer – A single optimizer.

• Tuple[Optimizer, LRScheduler] – An optimizer and a learning rate scheduler.

• Tuple[List[Optimizer], List[LRScheduler]] – Multiple optimizers and schedulers.

• Dict[str, Any] – A dictionary with ‘optimizer’ and optionally ‘lr_scheduler’ keys.

Examples

Simple optimizer:

>>> def configure_optimizers(self):
...     return braintools.optim.Adam(lr=1e-3)

Optimizer with scheduler:

>>> def configure_optimizers(self):
...     optimizer = braintools.optim.Adam(lr=1e-2)
...     scheduler = braintools.optim.StepLR(base_lr=1e-2, step_size=10)
...     return optimizer, scheduler

Multiple optimizers (e.g., for GANs):

>>> def configure_optimizers(self):
...     opt_g = braintools.optim.Adam(lr=1e-4)
...     opt_d = braintools.optim.Adam(lr=4e-4)
...     return [opt_g, opt_d], []

property current_epoch: int

Current training epoch (0-indexed).

property device: Any

The device this module is on (inferred from parameters).

freeze(names=None)

Freeze parameters so they are not updated during training.

brainstate.ParamState has no requires_grad flag, so freezing is recorded by name on the module. The Trainer reads these names when it sets up training and excludes the corresponding states from both gradient computation and optimizer updates. Call this before trainer.fit(model).

Parameters:

names (str | List[str] | None) – Parameter name(s) to freeze. If None (default), all parameters are frozen. Names correspond to the keys of model.states(brainstate.ParamState).

See also

unfreeze

Reverse the effect of this method.

Examples

>>> model.freeze()                 # freeze everything
>>> model.unfreeze('decoder.bias')  # then re-enable one parameter

property frozen_parameters: set

Set of parameter names currently frozen.

property global_step: int

Total number of training steps across all epochs.

is_frozen(name)

Return whether the parameter name is currently frozen.

Return type:

bool

load_state_dict(state_dict)

Load state from a state dictionary.

Parameters:

state_dict (Dict[str, Any]) – State dictionary to load.

log(name, value, prog_bar=False, logger=True, on_step=True, on_epoch=False, reduce_fx='mean', sync_dist=False)

Log a single metric.

Parameters:

• name (str) – Name of the metric.

• value (Any) – Value to log. Should be a scalar or JAX array.

• prog_bar (bool) – Whether to show this metric in the progress bar.

• logger (bool) – Whether to log to the logger(s).

• on_step (bool) – Whether to log at each step.

• on_epoch (bool) – Whether to log at epoch end (accumulated).

• reduce_fx (str) – Reduction function for epoch-level metrics (‘mean’, ‘sum’, ‘min’, ‘max’).

• sync_dist (bool) – Whether to synchronize across devices in distributed training.

Examples

>>> self.log('train_loss', loss, prog_bar=True)
>>> self.log('val_acc', accuracy, on_step=False, on_epoch=True)

log_dict(metrics, prog_bar=False, logger=True, on_step=True, on_epoch=False, reduce_fx='mean', sync_dist=False)

Log multiple metrics at once.

Parameters:

• metrics (Dict[str, Any]) – Dictionary of metric names to values.

• prog_bar (bool) – Whether to show these metrics in the progress bar.

• logger (bool) – Whether to log to the logger(s).

• on_step (bool) – Whether to log at each step.

• on_epoch (bool) – Whether to log at epoch end (accumulated).

• reduce_fx (str) – Reduction function for epoch-level metrics.

• sync_dist (bool) – Whether to synchronize across devices in distributed training.

Examples

>>> self.log_dict({'val_loss': loss, 'val_acc': acc}, prog_bar=True)

property logged_metrics: Dict[str, Any]

Metrics logged during the current step.

on_after_backward()

Called after backward pass (gradient computation).

on_after_optimizer_step(optimizer)

Called after each optimizer step.

on_before_backward(loss)

Called before backward pass (gradient computation).

on_before_optimizer_step(optimizer)

Called before each optimizer step.

on_fit_end()

Called at the very end of fit.

on_fit_start()

Called at the very beginning of fit.

on_predict_batch_end(outputs, batch, batch_idx)

Called at the end of each predict batch.

on_predict_batch_start(batch, batch_idx)

Called at the beginning of each predict batch.

on_predict_end()

Called at the end of prediction.

on_predict_start()

Called at the beginning of prediction.

on_test_batch_end(outputs, batch, batch_idx)

Called at the end of each test batch.

on_test_batch_start(batch, batch_idx)

Called at the beginning of each test batch.

on_test_end()

Called at the end of testing.

on_test_epoch_end()

Called at the end of each test epoch.

on_test_epoch_start()

Called at the beginning of each test epoch.

on_test_start()

Called at the beginning of testing.

on_train_batch_end(outputs, batch, batch_idx)

Called at the end of each training batch.

on_train_batch_start(batch, batch_idx)

Called at the beginning of each training batch.

on_train_end()

Called at the end of training.

on_train_epoch_end()

Called at the end of each training epoch.

on_train_epoch_start()

Called at the beginning of each training epoch.

on_train_start()

Called at the beginning of training.

on_validation_batch_end(outputs, batch, batch_idx)

Called at the end of each validation batch.

on_validation_batch_start(batch, batch_idx)

Called at the beginning of each validation batch.

on_validation_end()

Called at the end of validation.

on_validation_epoch_end()

Called at the end of each validation epoch.

on_validation_epoch_start()

Called at the beginning of each validation epoch.

on_validation_start()

Called at the beginning of validation.

predict_step(batch, batch_idx)

Compute predictions for a single batch.

Override this method to define your prediction logic.

Parameters:

• batch (Any) – A batch of data from the predict dataloader.

• batch_idx (int) – Index of the current batch.

Returns:

The predictions for this batch.

Return type:

Any

print_summary(input_shape=None)

Print a summary of the model architecture.

Parameters:

input_shape (Tuple[int, ...] | None) – Reserved for future shape-inference support. Currently unused; the summary reports parameter counts from the model’s existing states regardless of this argument.

state_dict()

Return the state dictionary for checkpointing.

Returns:

State dictionary containing all parameter states.

Return type:

Dict[str, Any]

test_step(batch, batch_idx)

Compute test metrics for a single batch.

Override this method to define your test logic.

Parameters:

• batch (Any) – A batch of data from the test dataloader.

• batch_idx (int) – Index of the current batch.

Returns:

A dictionary of metrics, an EvalOutput instance, or None.

Return type:

Dict[str, Any] | EvalOutput | None

property trainer: Any | None

Reference to the Trainer instance.

training_step(batch, batch_idx)

Compute the training loss for a single batch.

Override this method to define your training logic.

Parameters:

• batch (Any) – A batch of data from the train dataloader.

• batch_idx (int) – Index of the current batch.

Returns:

A dictionary with at least a ‘loss’ key, or a TrainOutput instance. The loss will be used for gradient computation.

Return type:

Dict[str, Any] | TrainOutput

Examples

>>> def training_step(self, batch, batch_idx):
...     x, y = batch['x'], batch['y']
...     logits = self(x)
...     loss = cross_entropy(logits, y)
...     self.log('train_loss', loss, prog_bar=True)
...     return {'loss': loss}

unfreeze(names=None)

Unfreeze parameters so they are trained again.

Parameters:

names (str | List[str] | None) – Parameter name(s) to unfreeze. If None (default), all parameters are unfrozen.

See also

freeze

Mark parameters as non-trainable.

validation_step(batch, batch_idx)

Compute validation metrics for a single batch.

Override this method to define your validation logic.

Parameters:

• batch (Any) – A batch of data from the validation dataloader.

• batch_idx (int) – Index of the current batch.

Returns:

A dictionary of metrics, an EvalOutput instance, or None.

Return type:

Dict[str, Any] | EvalOutput | None

Examples

>>> def validation_step(self, batch, batch_idx):
...     x, y = batch['x'], batch['y']
...     logits = self(x)
...     loss = cross_entropy(logits, y)
...     acc = (logits.argmax(-1) == y).mean()
...     self.log_dict({'val_loss': loss, 'val_acc': acc})
...     return {'val_loss': loss, 'val_acc': acc}