braindecode.models.NICE#
- class braindecode.models.NICE(n_outputs=None, n_chans=None, chs_info=None, n_times=None, input_window_seconds=None, sfreq=None, n_filters=40, filter_time_length=25, pool_time_length=51, pool_time_stride=5, embed_dim=768, activation=<class 'torch.nn.modules.activation.ELU'>, drop_prob=0.5, drop_prob_proj=0.5)[source]#
NICE EEG/MEG encoder from Song et al. (2024) [song2024nice].
Convolution
Added in version 2.1.
NICE (Natural Image Contrast EEG) learns an EEG/MEG embedding aligned with the image embedding of a frozen CLIP ViT-L/14 by contrastive learning, and decodes seen images zero-shot by nearest-neighbour search in that space. This class is the brain branch: the TSConv encoder (
Enc_eeg) and its projector (Proj_eeg).TSConv: a temporal convolution with
n_filterskernels of size(1, filter_time_length), average pooling of size(1, pool_time_length)and stride(1, pool_time_stride), batch norm + activation, a spatial convolution over alln_chans, batch norm + activation, dropout and a 1x1 convolution. UnlikeEEGConformer’s patch embedding, the pooling comes before the spatial convolution.Projector: the flattened
(time, filters)map goes through a linear layer toembed_dim, thenLayerNorm(x + Dropout(Linear(GELU(x)))). Its output is the aligned embedding, returned as"features"byreturn_features=True; the contrastive loss L2-normalises it.Head:
final_layeris a linear classifier on the embedding (not in the original, which only does retrieval).
The paper uses THINGS-EEG2 (63 channels, 250 Hz, 0-1 s, so
n_chans=63, n_times=250) and, in its appendix A.6, THINGS-MEG (271 channels, 200 Hz, 0-1 s, son_chans=271, n_times=200), with the same TSConv hyperparameters. The plug-in spatial modules of the paper (NICE-SA, NICE-GA) are not included.The authors’ checkpoint (eeyhsong/NICE, CC-BY-NC-ND-4.0, 10 subjects of THINGS-EEG2) keeps its original keys, which
mappingtranslates:model.load_state_dict({**enc, **proj}, strict=False)withencandprojthe state dicts ofpretrain_Enc_eeg_cls.pthandpretrain_Proj_eeg_cls.pth. At Hub revision7f98b08the encoder file is truncated andtorch.load()rejects it; the projector file loads.- 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 tomne.Infofor 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.
n_filters (
int) – Number of temporal and spatial filters (kin the paper).filter_time_length (
int) – Temporal kernel length (m1).pool_time_length (
int) – Average-pooling length (m2).pool_time_stride (
int) – Average-pooling stride (s).embed_dim (
int) – Size of the aligned embedding, the CLIP ViT-L/14 image feature size.activation (
type[Module]) – Activation of the convolutional encoder.drop_prob (
float) – Dropout of the convolutional encoder.drop_prob_proj (
float) – Dropout in the projector’s residual branch.
- 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
[song2024nice]Song, Y., Liu, B., Li, X., Shi, N., Wang, Y., & Gao, X. (2024). Decoding Natural Images from EEG for Object Recognition. International Conference on Learning Representations (ICLR). https://arxiv.org/abs/2308.13234
Hugging Face Hub integration
When the optional
huggingface_hubpackage 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 NICE # Train your model model = NICE(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-nice-model", commit_message="Initial model upload", model_card_kwargs={"datasets": ["BNCI2014_001"]}, )
Loading a model from the Hub:
from braindecode.models import NICE # Load pretrained model model = NICE.from_pretrained("username/my-nice-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 = NICE.from_pretrained("username/my-nice-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 = NICE.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, mask=None)[source]#
Encode
xof shape (batch, n_chans, n_times).Returns logits
(batch, n_outputs). In eager mode,return_features=Truereturns"features"(the aligned embedding,(batch, embed_dim)),"cls_token"(None), and"tokens"(batch, n_steps, n_filters): the projected temporal convolution map before flattening and the global projector, not per-channel tokens.n_steps = (n_times - filter_time_length - pool_time_length + 1) // pool_time_stride + 1. TorchScript always returns logits.maskis an optional boolean(batch, n_chans, n_times)tensor: True replaces the raw sample with zero before the temporal convolution. Hidden samples may contain NaN; their input gradients are zero. Mask sampling and the SSL loss are downstream, with no reconstruction head. Tokens are not forced to zero: each reads all montage channels andfilter_time_length + pool_time_length - 1input samples (75 with defaults), starting atj * pool_time_stridefor tokenj. Receptive fields overlap and are noncausal; BatchNorm during training also shares statistics across time, channels and recordings. Account for this when choosing masked spans and reconstruction targets.A
ch_mask(seecheck_ch_mask()) raises aValueError: the spatial convolution has one weight per channel of then_chansmontage, so a recording cannot run on fewer channels.