braindecode.models.EpiNT#

class braindecode.models.EpiNT(patch_size=256, embed_dim=512, n_layers=6, n_heads=8, ffn_dim=2048, drop_prob=0.1, head_drop_prob=0.1, rope_theta=10000.0, activation=<class 'torch.nn.modules.activation.ReLU'>, n_outputs=None, n_chans=None, chs_info=None, n_times=None, input_window_seconds=None, sfreq=None, channel_strategy='native', channel_strategy_kwargs=None)[source]#

EpiNT from Zhang et al. (2025) [EpiNT2025].

Foundation Model Attention/Transformer

EpiNT architecture

Added in version 1.9.

EpiNT (Epilepsy Neurophysiological Transformer) is pretrained on 2,700 hours of scalp EEG and intracranial EEG (ECoG, sEEG) from 1,199 patients. It is channel independent: every channel is encoded alone, so one checkpoint serves any montage and modality.

Architecture Overview

  1. instance_norm z-scores each channel over time (affine, one scale and one offset shared by all channels).

  2. patch_embedding projects non-overlapping patches of patch_size samples to embed_dim, and a learned cls_token is prepended.

  3. encoder is a stack of post-norm Transformer layers with rotary position embeddings (RoPE) in the self-attention. As upstream, each LayerNorm spans the token and feature axes, so the number of patches (n_times // patch_size) is fixed at construction.

  4. The class token of each channel is averaged over channels and final_layer (dropout and one linear layer) classifies it.

The pretraining mask token and frequency-domain mapping quantizer are not used downstream, so they are not part of this module.

Parameters:

  • patch_size (int) – Samples per patch. Default 256.

  • embed_dim (int) – Transformer width. Default 512.

  • n_layers (int) – Number of Transformer layers. Default 6.

  • n_heads (int) – Number of attention heads; embed_dim // n_heads must be even for RoPE. Default 8.

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

  • drop_prob (float) – Dropout in the Transformer layers. Default 0.1.

  • head_drop_prob (float) – Dropout before final_layer. Default 0.1.

  • rope_theta (float) – RoPE base frequency. Default 10000.

  • activation (type[Module]) – Feed-forward activation. Default nn.ReLU, as pretrained.

  • 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.

  • 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

Input. The pretrained weights expect n_times=3072: 12 s at 256 Hz (scalp EEG) or 3 s at 1024 Hz (iEEG), in microvolts as the authors prepare them (instance_norm has eps=1e-5, so signals in volts are not scaled to unit variance). The authors classify single-channel samples (n_chans=1).

Weights. The authors’ weights are plain files in their repository, not a release (License): weights/representations.bin (the pretrained backbone) and weights/finetune_on_MAYO.bin (a linear probe of it on MAYO). Both load through mapping, the probe with strict=True:

state = torch.load("representations.bin", map_location="cpu")
model = EpiNT(n_chans=1, n_outputs=2, n_times=3072)
model.load_state_dict(state, strict=False)  # no head, mask token unused

The authors’ fine-tuning code at the time of the paper applies the cross-entropy to softmax outputs (changed upstream in August 2025); with it and seed 0 their code trains a MAYO probe with the scores of finetune_on_MAYO.bin (F1 0.916, accuracy 0.851).

Replication. With the released weights, the port matches the authors’ code within float rounding (logits within 1.4e-6 in float32). Fine-tuned with the authors’ protocol, the three MAYO cells of Table 3 (the authors’ test patients) are within 3 % of the paper. Linear probing on CHB-MIT and FNUSA gives 0.73-0.76 with both the port and the authors’ code (paper: 0.857 and 0.805); the paper does not publish its test patients for these two datasets or the runs behind its mean ± std.

References

[EpiNT2025]

Zhang, R., Yu, H., Gan, J. Q. and Wang, H., 2025. Cross-modal epileptic signal harmonization: Frequency domain mapping quantization for pre-training a unified neurophysiological transformer. arXiv:2506.17068. Code: https://github.com/RunKZhang/EpiNT

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 EpiNT

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

# Push to the Hub
model.push_to_hub(
    repo_id="username/my-epint-model",
    commit_message="Initial model upload",
)

Loading a model from the Hub:

from braindecode.models import EpiNT

# Load pretrained model
model = EpiNT.from_pretrained("username/my-epint-model")

# Load with a different number of outputs (head is rebuilt automatically)
model = EpiNT.from_pretrained("username/my-epint-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 = EpiNT.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)[source]#

Classify a batch of signals.

Parameters:

  • x (Tensor) – Input of shape (batch, n_chans, n_times).

  • return_features (bool) – If True, return {"features": pooled, "cls_token": cls}: cls holds the class token of every channel, shape (batch, n_chans, embed_dim), and pooled is its mean over channels. A scripted model ignores this flag.

Returns:

Logits of shape (batch, n_outputs), or the feature dict.

Return type:

torch.Tensor or dict

reset_head(n_outputs)[source]#

Replace the classification head for a new number of outputs.

This is called automatically by from_pretrained() when the user passes an n_outputs that differs from the saved config. Override in subclasses that need a model-specific head structure. Implementations keep changed constructor arguments in sync with self._update_init_kwargs, so that a saved model can be loaded back. Implementations requiring positive outputs can also use self._set_n_outputs to validate and record the new value.

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.