braindecode.models.DuIN#

class braindecode.models.DuIN(n_outputs=None, n_chans=None, chs_info=None, n_times=None, input_window_seconds=None, sfreq=None, *, patch_size=100, on_non_divisible='pad', n_neural=16, n_filters=(128, 128, 16), kernel_sizes=(19, 3, 3), strides=(10, 1, 1), n_layers=8, n_heads=8, head_dim=64, ffn_dim=320, head_hidden_dim=128, activation=<class 'torch.nn.modules.activation.ReLU'>, attn_drop_prob=0.2, drop_prob=0.2, channel_strategy='native', channel_strategy_kwargs=None)[source]#

Du-IN from Zheng et al. (2024) [DuIN2024].

Foundation Model Convolution Attention/Transformer

Architecture Overview

Du-IN decodes speech from intracranial (sEEG) signals with region-level tokens: instead of one token per channel, the few channels recorded in a language-related region (ventral sensorimotor cortex, superior temporal gyrus) are fused into a single token per 100 ms patch [DuIN2024].

  1. Project the channels linearly into a small hidden “neural” space.

  2. Cut the signal into non-overlapping patches and encode each patch with a stack of strided 1D convolutions into one patch embedding.

  3. Add fixed sinusoidal temporal embeddings and run a Transformer encoder over the patch sequence.

  4. Flatten the encoded sequence and classify it with a one-hidden-layer MLP.

Macro Components

DuIN.spatial_projection

Operations. A linear layer from the n_chans input channels to n_neural hidden channels, applied at every time sample.

Role. Fuses the channels of the region into a common neural space. Its weights give the per-channel contribution scores used by the authors to select about 10 channels per subject.

DuIN.spatial_encoder

Operations. Each patch of patch_size samples goes through Conv1d + BatchNorm1d blocks (n_filters, kernel_sizes, strides); the output, time-major, is flattened into one embed_dim-dimensional patch embedding.

Role. Turns a patch of raw signal into a region-level token.

DuIN.encoder

Operations. Post-norm Transformer layers whose attention applies layer normalisation to the queries and keys of each head (as LaBraM), followed by a ReLU feed-forward block.

Role. Models the dynamics across the patch sequence.

DuIN.head_hidden and DuIN.final_layer

Operations. The encoded (n_patches, embed_dim) sequence is flattened, mapped to head_hidden_dim units with a ReLU, then to n_outputs by a linear layer.

Role. The upstream “label prediction head” of the Du-IN CLS model.

Pretraining

Du-IN is pretrained per subject in two stages: a VQ-VAE learns a discrete codebook of patch embeddings, then a masked model (Du-IN MAE) predicts the codebook indices of 50% masked patches. Fine-tuning (“Du-IN (mae)”, 62.70% top-1 accuracy over the 61 words, averaged over 12 subjects) loads the encoder of the MAE and trains it with the head above. Only the classification model is ported here: the VQ-VAE, the codebook and the MAE token head are not.

Differences from the reference

  • The upstream head ends with a sigmoid and is trained with F.cross_entropy(torch.sigmoid(logits), target). The port returns the logits, as every braindecode classifier. Train with that loss to reproduce the paper: fine-tuning the released MAE checkpoints on the 61-word task (12 subjects x 6 seeds) gives 60.2% top-1 with it and 55.6% with the cross-entropy on the logits.

  • The upstream subject layer maps a one-hot subject id to a channel-to-neural matrix. The released checkpoints are single-subject, so it is a plain linear layer (spatial_projection) here; the multi-subject variant (“Du-IN (poms)”) is not ported.

  • The paper describes the patch convolutions as depthwise; the released code and checkpoints use full Conv1d layers, which the port follows.

  • The temporal embedding table is rebuilt from n_times and not stored in the state dict (it has no parameters).

Important

Pre-trained weights. The authors release one Du-IN checkpoint per subject (001 to 012) and pretraining stage (mae, the paper’s Du-IN, and vqvae) of their 61-word dataset (CC BY 4.0). They are converted to this class on the Hugging Face Hub at braindecode/duin-pretrained (subj-<subject>/<stage>/ model.safetensors; the root model.safetensors is subj-001/mae). Each was trained on that subject’s 10 selected channels at 1000 Hz, so it only fits that subject’s montage. The head is not pretrained, and it is the only part sized by n_times: at another n_times, from_pretrained loads the encoder and the head keeps its initialization:

model = DuIN.from_pretrained(
    "braindecode/duin-pretrained",
    filename="subj-002/mae/model.safetensors",
)

The encoder matches the upstream duin_cls model to float precision on identical inputs. load_state_dict also accepts the original upstream checkpoints (liulab-repository/Du-IN, pretrains/duin/ <subject>/<stage>/model/checkpoint-399.pth) with strict=False.

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. Default 100 (100 ms at 1000 Hz), as the released model.

  • 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 input window is padded or cropped before the spatial projection. The reference implementation requires a multiple ("error"); all three give the same model and output when n_times is a multiple of patch_size.

  • n_neural (int) – Number of hidden channels of spatial_projection. Default 16.

  • n_filters (tuple[int, ...]) – Output channels of each patch convolution. Default (128, 128, 16).

  • kernel_sizes (tuple[int, ...]) – Kernel size of each patch convolution. Default (19, 3, 3).

  • strides (tuple[int, ...]) – Stride of each patch convolution. Default (10, 1, 1). Their product must divide patch_size.

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

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

  • head_dim (int) – Dimension of each attention head. Default 64 (so the attention width, 512, is larger than embed_dim, as upstream).

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

  • head_hidden_dim (int) – Hidden units of the classification head. Default 128. 0 gives a single linear layer on the flattened sequence.

  • activation (type[Module]) – Activation of the feed-forward blocks and of the head. Default ReLU, as upstream.

  • attn_drop_prob (float) – Dropout on the attention weights. Default 0.2.

  • drop_prob (float) – Dropout after the first feed-forward layer. Default 0.2.

  • 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

[DuIN2024] (1,2)

Zheng, H., Wang, H.-T., Jiang, W.-B., Chen, Z.-T., He, L., Lin, P.-Y., Wei, P.-H., Zhao, G.-G., & Liu, Y.-Z. (2024). Du-IN: Discrete units-guided mask modeling for decoding speech from intracranial neural signals. Advances in Neural Information Processing Systems 37 (NeurIPS 2024). arXiv:2405.11459. Code: https://github.com/liulab-repository/Du-IN (MIT).

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 DuIN

# Train your model
model = DuIN(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-duin-model",
    commit_message="Initial model upload",
    model_card_kwargs={"datasets": ["BNCI2014_001"]},
)

Loading a model from the Hub:

from braindecode.models import DuIN

# Load pretrained model
model = DuIN.from_pretrained("username/my-duin-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 = DuIN.from_pretrained("username/my-duin-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 = DuIN.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]#

Decode a batch of signals.

Parameters:

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

  • return_features (bool) – If True, return the flattened encoded sequence the head consumes, (batch, n_patches * embed_dim), as {"features": features, "cls_token": None} (braindecode foundation-model convention; Du-IN has no class token). A TorchScript-compiled model returns the logits instead.

Returns:

Class logits of shape (batch, n_outputs), or the feature dict when return_features is set.

Return type:

torch.Tensor or dict

load_state_dict(state_dict, *args, **kwargs)[source]#

Also accept an upstream Du-IN checkpoint (duin_mae/duin_cls) and a checkpoint saved at another n_times.

Upstream keys are renamed to this implementation. The pretraining parts (mask_emb, the VQ codebook, the MAE token head), the fixed temporal table and the upstream classification head are dropped, so load such a checkpoint with strict=False. A head that reads another number of patches keeps this model’s weights, with a warning.

Parameters:

  • state_dict – The description is missing.

  • *args – The description is missing.

  • **kwargs – The description is missing.

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.