Codebase overview#

The repository is built from six pieces, each owning one job. One is fixed by the benchmark, so that a reported number means the same thing in every submission: the evaluation contract, and the standardized test path that reads it. The rest are yours to write or extend. This page says which piece is responsible for what, so you know where to look before reading Pretraining, Task Suite 1: Behavior Prediction, Task Suite 2: Neural Activity Prediction or Task Suite 3: Brain Region Prediction.

The six map onto src/ like this:

src/
├── ibl_bwb_eval/    # the evaluation contract: tasks, metrics, protocol, submission format
├── core/            # shared machinery, no suite of its own
│   ├── dataset.py   # IBLBrainWideBench2026, the base every dataset extends
│   ├── trainer.py   # BaseTrainer: loop, logging, checkpointing, DDP, patience
│   ├── model.py     # BaseModel, the model interface
│   └── transforms/, samplers/, nn/, finetuning/, ...
├── pretrain/        # one directory per model, each owning its trainer
├── ts1/             # dataset, task trainer, test mixin, model trainers
├── ts2/             # same layout as ts1
└── ts3/             # extractors and probes instead of a task trainer

Who owns what#

Evaluation contract (ibl_bwb_eval). Fixed, and imported by every other piece while importing none of them.

  • the eval protocol: the eval sessions, the eval seeds, the selection metric.

  • the task vocabulary of each suite, and one readout spec per task.

  • the submission format.

Dataset (IBLBrainWideBench2026, one subclass per suite, instantiated per task, split and session).

  • which recordings and which units the regime (pretrain or eval) allows.

  • the splits (train, val, and test), selected by split=.

  • the intervals a sampler may draw from, get_sampling_intervals.

  • the evaluation target, _get_target, taken before any augmentation to avoid leakage.

  • the conditioning every sample gets, dataset_transform, and the order the rest run in, __getitem__. Both are in How a sample reaches the model.

TS3 is the exception: it scores units rather than windows, so it serves whole sessions with no split, no task and no target.

Base Trainer (BaseTrainer). What depends on neither the task nor the model. Every trainer subclasses it, in all three suites and in pretraining.

  • the epoch loop, train(): when val runs, when test runs, when to stop.

  • everything around it: logging, checkpointing, DDP, early stopping.

Task trainers (TS1 and TS2, one each: TS1EvalTrainer, TS2EvalTrainer).

  • the task, its readout spec, and the loss that matches it.

  • the loaders, the optimizer and the finetuning strategy.

  • the epoch internals: train_epoch, val_epoch, loss and predict.

  • the wiring between model and dataset, link_model, the one place the two meet.

Testing is not theirs: TS1TestMixin and TS2TestMixin fix the test split, the intervals, the metrics and the prediction file, so every submission is scored the same way.

Model (BaseModel, in its own directory alongside its trainer and configs).

  • the architecture and its forward.

  • how a sample becomes model inputs, input_fn: one Data sample to the tensors this architecture expects.

  • whatever has to be built from the dataset, link_datasets: a unit or session vocabulary, say.

  • the readout head, configure_readout, matching the task’s readout spec, which is what makes one architecture evaluable on several tasks.

  • how pretrained weights land in it, load_ckpt.

Model trainers (next to the model they belong to). Where a model states what it needs differently, and nothing else.

  • on the eval side, a subclass of the suite’s trainer, overriding only what differs.

  • on the pretrain side, the objective itself, subclassing BaseTrainer directly. Its model code still satisfies the interface above, so the same class can be evaluated downstream by a suite’s trainer.

How a sample reaches the model#

Every suite and every pretraining run reads the benchmark through one base dataset class, IBLBrainWideBench2026, and each sample it hands the loader goes through the same steps, in the order __getitem__ fixes. Two of them are extension points: dataset_transform, owned by the dataset class, and transform, supplied from the config.

  sampler picks a recording and a window
        │
        ▼
  data.slice(start, end)          # the window, straight off the .h5
        │
        ▼
  dataset_transform(sample)       # dataset-owned: normalize, add, drop signals
        │
        ├──▶ _get_target(sample)       # TS1/TS2 snapshot the label here, from clean data
        ▼
  transform(sample)               # <split>_transforms from the config: augmentation
        │
        ▼
  model.input_fn(sample)          # last element of transform: sample -> model inputs
        │
        ▼
  collate_fn                      # pads and stacks the batch

data.slice#

A Data object is one whole recording, read lazily from the .h5: spikes as an irregular time series, behavior as a regular one, trials and split domains as intervals, and non-temporal attributes such as units and session metadata.

slice(start, end) cuts every time-based attribute down to the window, copies the rest through, and re-zeroes timestamps to the window start.

dataset_transform#

Owned by the dataset class, applied to every sample. It is where a dataset conditions the signals it promises to carry: IBLBrainWideBenchTS1 z-scores or bins the behavioral target, IBLBrainWideBenchTS2 drops the hold-out mask belonging to the split it is not serving. The base implementation returns the sample untouched.

It runs on a whole recording as readily as on a slice, so it can be called outside the sampler:

data = dataset.dataset_transform(dataset.get_recording(recording_id))

transform#

Unlike dataset_transform, this one comes from outside the dataset: pass it as the transform argument when you construct a dataset yourself, or let a trainer build it from the config. In link_model the trainer instantiates the train_transforms, val_transforms and test_transforms lists and composes them onto the dataset:

dataset.transform = Compose([dataset.transform, *split_transforms, model.input_fn])

The lists are empty by default; a model that wants augmentation sets them in its own trainer config:

train_transforms:
  - _target_: core.transforms.UnitDropout
    min_units: 0.6

Augmentation belongs here rather than in dataset_transform because the target is snapshotted between the two, so nothing listed here can leak into the label.

Note

Transforms come from two packages, and your own are welcome too. torch_brain.transforms carries the generic ones (UnitDropout, UnitFilter, RandomCrop, Compose), core.transforms the benchmark’s own (FilterUnits, AdditivePepperNoise). Anything taking and returning a Data can be listed.

model.input_fn#

Up to here the sample is still a Data object carrying every signal the recording holds. input_fn is where the model picks out the ones it needs and puts them in the shape it expects, with any light preprocessing that takes, binning the spikes for instance.

It returns a dict: the model_inputs entry is splatted into forward, and any other top-level key is the model’s own to read back in its trainer. It runs per item on the dataloader workers, so it returns one sample’s tensors, which collate_fn then pads and stacks into the batch.