Evaluating your own model

The CLI and run_benchmark() both take a model name, which means the model has to live in this repo with a models/*.yaml beside it. When your model lives in your own codebase, use evaluate_model() instead: it takes the model itself and needs no files here.

It takes your model already built, weights and all. One instance then serves every task in the selection, which is only possible for a model that adapts to its input, so this path asks two things of it:

• it accepts any channel count and any window length, since tasks differ in montage and epoch length – from a 3-electrode sensorimotor strip to a 128-channel cap;

• it reads channel identity in forward rather than from a montage fixed at construction. channel_positions is always passed; a forward that also names ch_names is handed the dataset’s electrode names, which is what a model keyed by name rather than by coordinate needs.

What it does not need is a classifier head. NeuralBench wraps your model in a probe sized to each task’s number of classes, so the same backbone can serve a 4-class motor imagery task and a 2-class P300 task unchanged.

A model to evaluate

This one encodes each channel’s time course into tokens, adds a per-channel embedding derived from that channel’s 3D position, then pools over channels. The output keeps a time axis – (batch, tokens, width), the usual shape a transformer encoder produces – and the number of tokens follows the input length, which is fine: the probe pools over them.

from torch import nn


class TinyFm(nn.Module):
    """A small position-conditioned encoder, adaptive in channels and time."""

    def __init__(self, width: int = 64) -> None:
        super().__init__()
        self.width = width
        # Channel identity comes from the position, so an unseen montage needs
        # no new parameters -- which is what lets one instance span datasets.
        self.pos_embedding = nn.Linear(3, width)
        self.encoder = nn.Conv1d(1, width, kernel_size=17, stride=8, padding=8)
        self.norm = nn.LayerNorm(width)

    def forward(self, x, channel_positions):
        # x: (batch, channels, samples)
        # channel_positions: (batch, channels, 3)
        batch, channels, samples = x.shape

        # Encode every channel with the same temporal filters.
        tokens = self.encoder(x.reshape(batch * channels, 1, samples))
        tokens = tokens.reshape(batch, channels, self.width, -1)
        tokens = tokens.permute(0, 1, 3, 2)  # (batch, channels, tokens, width)

        tokens = tokens + self.pos_embedding(channel_positions).unsqueeze(2)
        return self.norm(tokens.mean(1))  # (batch, tokens, width)

Check it before you queue anything

check_model() pushes a synthetic batch of the selection’s shapes through the model, at several montage widths. It reads only YAML, so it takes seconds and needs no downloaded data – which is the point, since the alternative is discovering a shape bug an hour into a real run.

from neuralbench import check_model

print(check_model(TinyFm(), "eeg", "motor_imagery"))

         task  n_spatial_locations  n_temporal_samples status  output_shape
0  motor_imagery                 3                 480     ok  (2, 60, 64)
1  motor_imagery                19                 480     ok  (2, 60, 64)
2  motor_imagery                64                 480     ok  (2, 60, 64)
3  motor_imagery               128                 480     ok  (2, 60, 64)

A model that only handles one montage width shows up as ok on one row and forward failed: ... on the others. Pass the same task and dataset you intend to run: dataset variants can shorten a task’s window, and each distinct window gets its own rows.

A forward that does not accept channel_positions at all raises instead of filling the status column, being a property of the model rather than of any one task.

Run it

device and task mean what they mean on the CLI, so task="all" runs every validated task for the device. Start with debug=True: it runs locally on a reduced config (2 epochs, 5 batches, a data subset).

from neuralbench import evaluate_model

scores = evaluate_model(
    TinyFm(),
    "eeg",
    "motor_imagery",
    name="tiny-fm",
    dataset=["schalk2004bci2000", "leeb2007"],
    debug=True,
)
print(scores[["dataset_name", "test/bal_acc", "n_total_params",
              "n_trainable_params"]])

        dataset_name  test/bal_acc  n_total_params  n_trainable_params
0  Schalk2004Bci2000         0.258             900                 132
1      Leeb2007Brain         0.475             834                  66

Those two datasets have 64 and 3 channels respectively, and one instance served both. The parameter counts show the split: the backbone is identical across the rows, and only the probe differs – 132 trainable parameters for Schalk’s 4 classes against 66 for Leeb’s 2.

Build the model however you normally would – model.load_state_dict(...) and all – and evaluate_model takes it from there. The instance is serialized once to the cache folder and reloaded fresh for each experiment, so weights that one run fine-tunes never leak into the next, and the weights themselves take part in the cache key: change them and you get new runs rather than stale cached results.

How your model is adapted to each task

downstream_wrapper chooses the adaptation strategy, and with it where the task’s classifier head comes from. The default, linear_probe_mean, freezes your model and trains a linear probe on its mean-pooled output.

scores = evaluate_model(
    TinyFm(), "eeg", "all", downstream_wrapper="finetune_mean"
)

Prefer the _mean presets here. flatten builds the probe from the flattened output instead, which for an adaptive model makes the probe’s size a function of the dataset – on the run above it would be 8196 trainable parameters for the 64-channel dataset against 194 for the 3-channel one. It works, but head capacity then varies with montage density and window length, which is not something you want varying underneath a benchmark number.

Suites

The dataset selection decides which suite you are running, exactly as it does on the CLI (see Benchmark suites):

Call	Suite
evaluate_model(m, "eeg", "all")	NeuralBench-EEG-Core (one dataset per task)
evaluate_model(m, "eeg", "all", dataset="all")	NeuralBench-EEG-Full (every dataset per task)

Downloading, caching, and where the runs happen

Two things have to happen before an experiment can train, and neither is automatic: the dataset has to be on disk, and the preprocessed data has to be in the cache. evaluate_model() does both first, in that order, because preprocessing needs the downloaded files. Once warm, both phases are a cheap status check, so leave them on; turn them off only when you know the state of the cache. Warming it for a whole device takes hours, and it is the same work the CLI does with --download and --prepare, so a cache warmed either way serves the other.

Everything then runs in the current process, which is usually what you want in a notebook, and which means a full suite can take a very long time. Pass cluster="auto" to fan the experiments out to SLURM where one is available:

scores = evaluate_model(
    TinyFm(), "eeg", "all", name="my-fm", cluster="auto",
    download=False,   # datasets already on disk
    prepare=False,    # caches already warm
)

That call returns as soon as the jobs are queued, so the frame holds whatever had already finished – nothing, the first time. Call it again with the same arguments to collect: completed experiments are not resubmitted, exactly as when re-running the CLI. With prepare=True the first call queues the cache warm-up alone and stops there, since the experiments would otherwise race it.

Results never raise on a failed experiment. You get the rows that exist, and the rest are reported in a printed table where !n marks jobs that errored out as opposed to ones still pending.

Changing the protocol

overrides takes any config key, as dotted keys or nested dicts, and is applied on top of the defaults and the task config:

scores = evaluate_model(
    TinyFm(),
    "eeg",
    "all",
    overrides={
        "data.neuro.frequency": 200.0,                    # preprocessing
        "lightning_optimizer_config.optimizer.lr": 1e-4,  # optimization
        "trainer_config.n_epochs": 50,
    },
)

Whatever you pass is recorded in the frame’s overrides column, so a run with a changed protocol is never mistaken for a stock one.

Be careful with keys that define what a task measures – the epoch window (data.start / data.duration), the target extractor, the splits, the loss, the metrics. Changing those produces a number that is not comparable to any other entry on the suite, including the published baselines.

A worked example: pretrained REVE

The foundation models the benchmark ships with are ordinary modules too, so one of them makes a good end-to-end example. REVE is a braindecode model and NtReve is the config that builds it: n_outputs=None asks for the encoder alone, since the probe supplies the head, and the width and window length only size the build – the instance reads the montage it actually gets from channel_positions on every call.

Note

The weights are pulled from the Hugging Face hub at build time, and REVE’s are gated, so the first build needs an account and an accepted licence: see Pretrained model weights. Later builds read them from the local hub cache and need no network.

from neuraltrain.models.reve import NtReve

model = NtReve(from_pretrained_name="brain-bzh/reve-base").build(
    n_spatial_locations=64,
    n_temporal_samples=400,
    n_outputs=None,
)

A pretrained model expects the preprocessing it was trained with rather than the task default, which is what overrides is for. These values are REVE’s, read off models/reve.yaml:

scores = evaluate_model(
    model,
    "eeg",
    "motor_imagery",
    name="reve-base",
    overrides={
        "data.neuro.frequency": 200.0,
        "data.neuro.filter": [0.5, 99.5],
        "data.neuro.notch_filter": None,
        "data.neuro.baseline": None,
        "data.neuro.scaler": "StandardScaler",
        "data.neuro.clamp": 15,
        # REVE's position bank is in head-frame metres, not normalised.
        "data.channel_positions.normalize": False,
    },
    debug=True,
)

Same configuration, same number as the CLI. The call above matches -m reve -w linear_probe_mean rather than a bare -m reve, which reads its adaptation strategy and its channel mapping from models/reve.yaml instead.

What the Slurm jobs need

Off debug=True, experiments run as Slurm jobs in separate processes, which rebuild the model from the serialized instance. Classes defined in a script, a notebook or an uninstalled checkout are stored in it whole, so there is nothing to package and nothing to put on the jobs’ PYTHONPATH. Code from installed packages is stored by import path instead: a model that imports, say, timm needs timm in the environment the jobs run in.