braindecode.models.Brant#

class braindecode.models.Brant(n_outputs=None, n_chans=None, chs_info=None, n_times=None, input_window_seconds=None, sfreq=None, *, patch_size=250, on_non_divisible='pad', embed_dim=256, ffn_dim=384, temporal_n_layers=4, spatial_n_layers=2, n_heads=8, n_freq_bands=8, band_power_sfreq=256.0, drop_prob=0.1, channel_strategy='native', channel_strategy_kwargs=None)[source]#

Brant from Zhang et al. (2023) [Brant2023].

Foundation Model Attention/Transformer

Architecture Overview

Brant models intracranial neural signals (sEEG/iEEG) in four stages [Brant2023]:

  1. Split every channel into non-overlapping temporal patches.

  2. Combine a linear patch projection with learned temporal positions and a spectral-power embedding.

  3. Apply a temporal Transformer within each channel, followed by a spatial Transformer across channels at each patch index.

  4. Mean-pool the channel-patch representation and classify it with a Braindecode downstream head.

Macro Components

Brant.patch_tokenizer

Operations. Crops an incomplete tail and reshapes the signal into non-overlapping patches of patch_size samples.

Role. Preserves the raw samples consumed by both the linear patch projection and the spectral-power calculation.

Brant.band_power and Brant.temporal_encoder

Operations. A periodogram is summed over the eight rhythmic bands in BRANT_FREQ_BANDS. Their log powers softmax-weight learned band embeddings, which are added to projected patches and temporal positions before self-attention within each channel.

Role. Fuse time- and frequency-domain information while capturing long-range dependencies between consecutive patches.

Brant.spatial_encoder

Operations. Self-attention is applied across all channels sharing a patch index.

Role. Capture spatial correlations without a fixed channel vocabulary or channel-specific parameters.

Brant.final_layer

Operations. Mean-pool the encoded channel-patch grid, then apply a single linear layer.

Role. Adapt the upstream encoder to Braindecode classification. This pooling and head are not part of the masked-reconstruction pretraining objective.

Temporal, Spatial, and Spectral Encoding

  • Temporal: learned patch positions and the temporal Transformer preserve patch order and model long-range activity within each channel.

  • Spatial: the spatial Transformer attends across the channels at every temporal patch index.

  • Spectral: log power in theta, alpha, beta, and five gamma bands weights eight learned frequency embeddings.

Additional Mechanisms

The upstream reconstruction projection is retained in Brant.spatial_encoder as part of the source architecture, although classification does not consume its output. Brant has no class token, so return_features=True returns cls_token=None. The learned temporal positions require the runtime signal length to equal the configured n_times. As no parameter is channel-specific, the number of channels may differ from the configured n_chans, between batches and within one batch when the recordings are padded to a common count and ch_mask (see forward()) flags the real ones. A window that is not a multiple of patch_size is right zero-padded to ceil(n_times / patch_size) patches by default (on_non_divisible).

The upstream model operates on signals down-sampled to 250 Hz, but its band-power features are computed at a fixed 256 Hz (upstream pre_utils.py:54 and utils.py:36), whatever the true rate; see band_power_sfreq below. At a different rate the band edges shift and the pretrained embeddings change (by up to 0.55 in our checks).

The defaults are a modest, ready-to-run configuration. The paper’s large architecture uses patch_size=1500 (6 s), embed_dim=2048, ffn_dim=3072, temporal_n_layers=12, spatial_n_layers=5, n_heads=16, and n_times=22500 (15 patches, 90 s).

Important

Pretrained weights. braindecode/brant-pretrained holds the official weights converted to this implementation. All 210 shared tensors of the official checkpoint load unchanged; the encoder outputs of the port match the official code within 1e-5 on identical inputs. It uses the large configuration above; the classification head is braindecode’s and is not pretrained:

model = Brant.from_pretrained(
    "braindecode/brant-pretrained", n_outputs=2
)

Added in version 2.0.

Parameters:

  • n_outputs (int) – Number of outputs of the model. This is the number of classes in the case of classification.

  • n_chans (int) – Number of EEG channels.

  • chs_info (list of dict) – Information about each individual EEG channel. This should be filled with info["chs"]. Refer to mne.Info for more details.

  • n_times (int) – Number of time samples of the input window.

  • input_window_seconds (float) – Length of the input window in seconds.

  • sfreq (float) – Sampling frequency of the EEG recordings.

  • patch_size (int) – Number of time samples per patch fed to the encoders. Default 250 (1 s at 250 Hz). The paper’s large architecture uses 1500 (see above).

  • on_non_divisible (str) – How PatchTokenizer treats a time axis that is not a multiple of patch_size (right zero-pad, crop, or raise). The reference implementation crops ("crop"); both give the same model and output when n_times is a multiple of patch_size.

  • embed_dim (int) – Model width D (patch embedding size). Default 256.

  • ffn_dim (int) – Inner dimension of the Transformer feed-forward blocks. Default 384.

  • temporal_n_layers (int) – Number of layers in the temporal Transformer encoder. Default 4.

  • spatial_n_layers (int) – Number of layers in the spatial Transformer encoder. Default 2.

  • n_heads (int) – Number of attention heads in both encoders. Default 8.

  • n_freq_bands (int) – Number of frequency bands used by the frequency encoding. Default 8 (must match len(BRANT_FREQ_BANDS)).

  • band_power_sfreq (float) – Sampling rate the upstream code assumes when it computes the eight band powers: 256 Hz, although the data are 250 Hz (Brant_src/utils.py:36). Keep the default to load the released weights faithfully; sfreq describes the data and is not used by the band-power features.

  • drop_prob (float) – Dropout probability. Default 0.1.

  • channel_strategy (str) – How any montage reaches the backbone (pretrained models only; see Channel strategies: any montage in). "native" keeps the model as it is. "exact", "zero", "nearest", "idw", "spline", "field", "source", "region", "wiener" (call model.channel_layer.fit first) or "latent" map the montage of chs_info (or of the chs_info given to forward) onto the backbone’s channels with a ChannelLayer. Saved in the config. model(x) and model.forward(x) both apply the layer.

  • channel_strategy_kwargs (dict | None) – Options of the strategy (e.g. {"reg": 1e-2} for "spline").

Raises:

ValueError – If some input signal-related parameters are not specified: and can not be inferred.

Notes

If some input signal-related parameters are not specified, there will be an attempt to infer them from the other parameters.

References

[Brant2023] (1,2)

Zhang, D., Yuan, Z., Yang, Y., Chen, J., Wang, J. and Li, Y., 2023. Brant: Foundation Model for Intracranial Neural Signal. In Thirty-seventh Conference on Neural Information Processing Systems, NeurIPS. Code: yzz673/Brant (Apache-2.0).

Hugging Face Hub integration

When the optional huggingface_hub package is installed, all models automatically gain the ability to be pushed to and loaded from the Hugging Face Hub. Install with:

pip install braindecode[hub]

Pushing a model to the Hub:

from braindecode.models import Brant

# Train your model
model = Brant(n_chans=22, n_outputs=4, n_times=1000)
# ... training code ...

# Push to the Hub; ``datasets`` and ``base_model`` go to the card
model.push_to_hub(
    repo_id="username/my-brant-model",
    commit_message="Initial model upload",
    model_card_kwargs={"datasets": ["BNCI2014_001"]},
)

Loading a model from the Hub:

from braindecode.models import Brant

# Load pretrained model
model = Brant.from_pretrained("username/my-brant-model")
# Data the weights were trained on, from the Hub model card
print(model.card_data.datasets if model.card_data else None)

# Load with a different number of outputs (head is rebuilt automatically)
model = Brant.from_pretrained("username/my-brant-model", n_outputs=4)

Extracting features and replacing the head:

import torch

x = torch.randn(1, model.n_chans, model.n_times)
# Extract encoder features (consistent dict across all models)
out = model(x, return_features=True)
features = out["features"]

# Replace the classification head
model.reset_head(n_outputs=10)

Saving and restoring full configuration:

import json

config = model.get_config()            # all __init__ params
with open("config.json", "w") as f:
    json.dump(config, f)

model2 = Brant.from_config(config)    # reconstruct (no weights)

All model parameters (both EEG-specific and model-specific such as dropout rates, activation functions, number of filters) are automatically saved to the Hub and restored when loading.

See Loading and Adapting Pretrained Foundation Models for a complete tutorial.

Methods

forward(x, return_features=False, ch_mask=None)[source]#

Decode a batch of signals.

Parameters:

  • x (Tensor) – Input of shape (batch, n_chans, n_times); n_chans may differ from the construction one.

  • return_features (bool) – If True, return the pooled encoder embedding instead of the class logits, as {"features": pooled, "cls_token": None} (braindecode foundation-model convention). Brant pools over channels and patches and has no class token, hence cls_token is None. A TorchScript-compiled model returns the logits instead.

  • ch_mask (Tensor | None) – (batch, n_chans) mask of real channels, see check_ch_mask(). Padded channels affect no output.

Returns:

Class logits of shape (batch, n_outputs), or the feature dict {"features", "cls_token"} when return_features is set.

Return type:

torch.Tensor or dict

reset_head(n_outputs)[source]#

Swap the classification head for a new number of outputs.

Parameters:

n_outputs (int) – New number of output classes.

Return type:

None

Examples

>>> from braindecode.models import BENDR
>>> model = BENDR(n_chans=22, n_times=1000, n_outputs=4)
>>> model.reset_head(10)
>>> model.n_outputs
10

Added in version 1.4.

Examples using braindecode.models.Brant#

Loading and Adapting Pretrained Foundation Models

Loading and Adapting Pretrained Foundation Models