Modifying the Training Loop

NeuralBench uses a PyTorch Lightning LightningModule called BrainModule to handle training, validation, and testing. This tutorial explains how the training loop works and how to customize it by subclassing BrainModule.

How BrainModule works

BrainModule lives in neuralbench/pl_module.py. Its key components are:

• ``model_forward(batch)`` – runs the model on a batch, automatically passing subject_ids and channel_positions when the model’s forward signature requires them.

• ``_run_step(batch, step_name, batch_idx)`` – shared logic for train, validation, and test steps:

  1. Extracts y_true from batch.data["target"]

  2. Applies target scaling if configured (e.g., for regression)

  3. Converts targets for the loss function (one-hot to argmax for CrossEntropyLoss, clamping for BCEWithLogitsLoss)

  4. Calls model_forward(batch) to get predictions

  5. Computes and logs the loss

  6. Updates all metrics matching the current step name

• ``training_step``, ``validation_step``, ``test_step`` – thin wrappers around _run_step that return the appropriate values for Lightning.

• ``configure_optimizers`` – builds the optimizer and learning rate scheduler from the config, injecting total_steps / T_max automatically.

Here is _run_step:

def _run_step(
    self, batch: Batch, step_name: str, batch_idx: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    y_true = batch.data["target"]

    if self.target_scaler is not None:
        y_true = self.target_scaler.transform(y_true)
    if y_true.ndim == 3 and y_true.shape[1] == 1:
        y_true = y_true.squeeze(1)

    if isinstance(self.loss, nn.CrossEntropyLoss):
        assert y_true.ndim == 2
        y_true = y_true.argmax(dim=1)
    elif isinstance(self.loss, nn.BCEWithLogitsLoss):
        y_true = y_true.clamp(max=1.0)

    log_kwargs = {
        "on_step": step_name == "train",
        "on_epoch": True,
        "logger": True,
        "prog_bar": True,
        "batch_size": y_true.shape[0],
        "sync_dist": self.trainer.world_size > 1,
    }

    y_pred = self.model_forward(batch)

    if y_pred.ndim == 3 and y_true.ndim == 3:
        y_pred = y_pred.reshape(y_pred.shape[0], -1)
        y_true = y_true.reshape(y_true.shape[0], -1)

    loss = self.loss(y_pred, y_true)
    self.log(f"{step_name}/loss", loss, **log_kwargs)

    for metric_name, metric in self.metrics.items():
        if metric_name.startswith(step_name):
            metric.update(y_pred, y_true)
            if "confusion_matrix" not in metric_name:
                self.log(metric_name, metric, **log_kwargs)

    return loss, y_pred, y_true

The training step simply calls _run_step and returns the loss for backpropagation:

def training_step(self, batch: Batch, batch_idx: int):
    loss, _, _ = self._run_step(batch, step_name="train", batch_idx=batch_idx)
    return loss

Sequential evaluation

To feed each timeline (recording) to a model one window at a time, in time order, set data.val_batch_size=1 and/or data.test_batch_size=1. Validation and test windows are ordered by timeline then time unless data.val_shuffle/data.test_shuffle is set. Keeping validation batched (only test_batch_size=1) keeps epochs fast and multi-GPU training available. eeg _sleep_onset_stream and eeg _motor_imagery_stream stream their test split this way.

data.stream_by lists the event fields that identify a stream, e.g. [timeline] for one stream per recording or [subject, session] for one per session, whose timelines must then be listed contiguously. It adds a PyTorch Lightning callback that acts whenever a new stream starts during evaluation. It runs the stream on a fresh copy of the model as it was when evaluation began, and calls the copy’s optional reset_state(), so stateful models know where streams begin and nothing the model changes while predicting (weights, buffers, attributes) carries over to the next stream. The original model is restored when evaluation ends. reset_state() also runs at each new stream during training, which with shuffled batches is almost every batch. The wrapper forwards the hook to its backbone. GroupedMetric metrics then score each stream separately.

As an illustration only (not a NeuralBench model), a recurrent backbone could carry its hidden state across windows and clear it in reset_state():

class RecurrentBackbone(nn.Module):
    def __init__(self, n_chans: int, n_hidden: int = 64):
        super().__init__()
        self.gru = nn.GRU(n_chans, n_hidden, batch_first=True)
        self.hidden: torch.Tensor | None = None

    def reset_state(self) -> None:
        self.hidden = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:  # x: (B, C, T)
        out, hidden = self.gru(x.transpose(1, 2), self.hidden)
        self.hidden = hidden.detach()  # next backward() would hit a freed graph
        return out[:, -1]

%% Subclassing BrainModule ———————–

To customize the training loop, subclass BrainModule and override the methods you need. The rest of the pipeline (data loading, checkpointing, testing) remains unchanged.

Example: adding an L2 regularization term to the training loss.

import torch

from neuralbench.pl_module import BrainModule


class RegularizedBrainModule(BrainModule):
    """BrainModule with an L2 penalty on model weights."""

    def __init__(self, *args, l2_lambda: float = 1e-4, **kwargs):
        super().__init__(*args, **kwargs)
        self.l2_lambda = l2_lambda

    def training_step(self, batch, batch_idx):
        loss, _, _ = self._run_step(batch, step_name="train", batch_idx=batch_idx)

        l2_reg = sum(p.pow(2).sum() for p in self.model.parameters())
        loss = loss + self.l2_lambda * l2_reg

        self.log("train/l2_reg", l2_reg, prog_bar=False)
        return loss

Example: logging additional metrics during validation.

class VerboseValBrainModule(BrainModule):
    """BrainModule that logs extra info during validation."""

    def validation_step(self, batch, batch_idx):
        _, y_pred, y_true = self._run_step(batch, step_name="val", batch_idx=batch_idx)
        confidence = torch.softmax(y_pred, dim=-1).max(dim=-1).values.mean()
        self.log("val/mean_confidence", confidence, prog_bar=True)
        return y_pred, y_true

What you can customize

By overriding different methods on BrainModule, you can:

• Add auxiliary losses (e.g., contrastive terms, regularization) by overriding training_step.

• Change how predictions are computed by overriding model_forward.

• Modify target preprocessing by overriding _run_step.

• Add custom logging or callbacks in any step method.

• Change the optimizer or scheduler by overriding configure_optimizers.

Because NeuralBench is built on PyTorch Lightning, the full Lightning API is available. See the Lightning docs for more advanced hooks (on_train_epoch_end, on_before_optimizer_step, etc.).