torch_brain

Documentation | Join our Discord community

torch_brain is an end-to-end framework for building deep learning models and training pipelines for neuroscience. It pairs a lightweight, time-based data format (plus tools to preprocess existing neural datasets into it) with PyTorch-compatible building blocks: datasets, samplers, nn.Modules, and models.

Features

• Lazy, on-demand data loading that reads only the time-slices and attributes you request
• Advanced samplers for arbitrary on-the-fly slicing of recordings
• Multi-recording training across heterogeneous datasets
• Support for arbitrary neural and behavioral modalities
• Flexible collation strategies, including chaining and padding

Installation

torch_brain requires Python >= 3.10. To install a stable release:

pip install torch torch_brain

pip install git+https://github.com/neuro-galaxy/torch_brain

The data format

A recording is a Data object holding heterogeneous, time-aware modalities: regularly-sampled signals (LFP, EEG, etc.), irregular event streams (spikes), interval annotations (trials), and plain arrays.

import numpy as np
from torch_brain.data import Data, IrregularTimeSeries, RegularTimeSeries, Interval

data = Data(
    spikes=IrregularTimeSeries(  # event stream
        timestamps=[0.1, 0.2, 0.3, 2.1, 2.2, 2.3],
        unit_index=[0, 0, 1, 0, 1, 2],
        domain="auto",
    ),
    lfp=RegularTimeSeries(raw=np.zeros((1000, 3)), sampling_rate=250.0),  # 4s @ 250Hz
    trials=Interval(start=[0.0, 1.0, 2.0], end=[1.0, 2.0, 3.0]),  # annotations
    domain=Interval(0.0, 4.0),
)

The point of the format is that slicing is time-based and lazy: Every modality is sliced consistently, regardless of their different sampling rates, and the data is lazily read from disk so only the requested window and attributes are loaded.

window = data.slice(1.0, 3.0)
# spikes -> the 3 events in [1, 3)   lfp -> 500 samples   trials -> 2 trials

This is why a torch_brain Dataset is indexed by time, not by integer (see below).

Training pipelines

torch_brain leans on the standard PyTorch training loop, and most of its job is to handle the data side. You define a Dataset (built on the time-slicing above) and a Sampler that decides which slices become samples. The DataLoader, model, and loop are ordinary PyTorch.

import torch
from torch.utils.data import DataLoader
from torch_brain.datasets import PeiPandarinathNLB2021, DatasetIndex
from torch_brain.samplers import TrialSampler
from torch_brain.utils import bin_spikes


# torch_brain ships loaders for many public datasets.
# Subclass one to define the two things specific to your task:
class MyDataset(PeiPandarinathNLB2021):
    # 1. WHICH windows count as samples (here, one per behavioral trial).
    def get_sampling_intervals(self):
        sampling_intervals = {}
        for rid in self.recording_ids:
            sampling_intervals[rid] = self.get_recording(rid).trials
        return sampling_intervals

    # 2. HOW one window becomes tensors.
    def __getitem__(self, index: DatasetIndex):
        # `index` is a DatasetIndex(recording_id, start, end) handed in by the sampler;

        data = super().__getitem__(index)
        # super().__getitem__(...) returns that slice with
        # every modality (.spikes, .hand.vel, ...) lazily cropped.

        # Only attributes actually accessed will be loaded into memory from disk.
        X = bin_spikes(data.spikes, num_units=len(data.units), bin_size=0.05)
        Y = data.hand.vel
        return torch.from_numpy(X).float(), torch.from_numpy(Y).float()


dataset = MyDataset(root="data/processed", recording_ids=["jenkins_maze_train"])

# The sampler turns those intervals into per-sample DatasetIndex objects.
sampler = TrialSampler(
    sampling_intervals=dataset.get_sampling_intervals(), shuffle=True
)
loader = DataLoader(dataset, sampler=sampler, batch_size=8)

# From here on it's plain PyTorch
for X, Y in loader:
    pred = model(X)
    loss = loss_fn(pred, Y)
    ...

The key idea: unlike a standard PyTorch Dataset indexed by integers, a torch_brain Dataset is indexed by time-slices, and loads data lazily, so only the slice you ask for is read from disk. A Sampler decides what to load, the Dataset decides how, and everything downstream stays vanilla PyTorch.

See examples/ for simple and readable training implementations.

Contributing

Contributions are welcome! Get started with:

pip install -e ".[dev]"   # editable install with dev dependencies
pre-commit install        # formatting & lint hooks
pytest                    # run the test suite

See CONTRIBUTING.md for the full workflow and code-style guidelines.

Building the documentation

pip install -e ".[dev,docs]"
cd docs && make clean html

The built docs are placed in docs/build/html.