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:

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

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
model.push_to_hub(
    repo_id="username/my-duin-model",
    commit_message="Initial model upload",
)

Loading a model from the Hub:

from braindecode.models import DuIN

# Load pretrained model
model = DuIN.from_pretrained("username/my-duin-model")

# 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).

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.

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.