Upload 25 files

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+example_files/eeg_sample/example_eeg_file.xdf filter=lfs diff=lfs merge=lfs -text
+images/banner.png filter=lfs diff=lfs merge=lfs -text
+images/ecg.png filter=lfs diff=lfs merge=lfs -text
+images/eeg.png filter=lfs diff=lfs merge=lfs -text
+images/emg.png filter=lfs diff=lfs merge=lfs -text

NeuroRVQ/NeuroRVQ.py

@@ -0,0 +1,821 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from einops import rearrange
+import math
+from functools import partial
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.dirname(__file__)))
+from NeuroRVQ_modules import Block, trunc_normal_
+from RVQ import ResidualVectorQuantization
+
+def inverse_fft_cos_sin(fft_amp, fft_sin_pha, fft_cos_pha):
+    """
+    Inverse FFT function using sin and cos
+    :param fft_amp: amplitude
+    :param fft_sin_pha: sine
+    :param fft_cos_pha: cosine
+    :return: inverse fft in time
+    """
+    imag = fft_amp * fft_sin_pha
+    real = fft_amp * fft_cos_pha
+    fft_y = torch.complex(real, imag)
+    y = torch.fft.ifft(fft_y)
+    return y
+
+class PatchEmbed(nn.Module):
+    """
+    Project each codebook to the patch latent space
+    :param in_chans: number of input channels
+    :param embed_dim: dimension of embedding space
+    """
+    def __init__(self, in_chans=1, embed_dim=200):
+        super().__init__()
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(1, 1), stride=(1, 1))
+    def forward(self, x):
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+        
+class MultiDimentionalTemporalConv(nn.Module):
+    """
+    EEG to Patch Embedding - Multidimentional Temporal Filtering
+    :param in_chans: number of input channels
+    :param out_chans: number of output channels
+    """
+    def __init__(self, in_chans=1, out_chans=8):
+        super().__init__()
+        # Inception Style Seperate Branches - Group 1 #
+        # Branch 1: >10 Hz assuming fs=200Hz
+        self.conv1_1 = nn.Conv2d(in_chans, out_chans, kernel_size=(1, 21), padding=(0, 10))
+        self.norm1_1 = nn.GroupNorm(4, out_chans)
+        self.pool1_1 = nn.AvgPool2d(kernel_size=(1, 2))
+
+        # Branch 2: >13 Hz assuming fs=200Hz
+        self.conv1_2 = nn.Conv2d(in_chans, out_chans, kernel_size=(1, 15), padding=(0, 7))
+        self.norm1_2 = nn.GroupNorm(4, out_chans)
+        self.pool1_2 = nn.AvgPool2d(kernel_size=(1, 2))
+
+        # Branch 3: >20 Hz assuming fs=200Hz
+        self.conv1_3 = nn.Conv2d(in_chans, out_chans, kernel_size=(1, 9), padding=(0, 4))
+        self.norm1_3 = nn.GroupNorm(4, out_chans)
+        self.pool1_3 = nn.AvgPool2d(kernel_size=(1, 2))
+
+        # Branch 4: >40 Hz assuming fs=200Hz
+        self.conv1_4 = nn.Conv2d(in_chans, out_chans, kernel_size=(1, 5), padding=(0, 2))
+        self.norm1_4 = nn.GroupNorm(4, out_chans)
+        self.pool1_4 = nn.AvgPool2d(kernel_size=(1, 2))
+        self.gelu1 = nn.GELU()
+        
+        # Inception Style Seperate Branches - Group 2 #
+        self.conv2_1 = nn.Conv2d(out_chans, out_chans, kernel_size=(1, 9), padding=(0, 4))
+        self.norm2_1 = nn.GroupNorm(4, out_chans)
+        self.pool2_1 = nn.AvgPool2d(kernel_size=(1, 4))
+
+        self.conv2_2 = nn.Conv2d(out_chans, out_chans, kernel_size=(1, 7), padding=(0, 3))
+        self.norm2_2 = nn.GroupNorm(4, out_chans)
+        self.pool2_2 = nn.AvgPool2d(kernel_size=(1, 4))
+
+        self.conv2_3 = nn.Conv2d(out_chans, out_chans, kernel_size=(1, 5), padding=(0, 2))
+        self.norm2_3 = nn.GroupNorm(4, out_chans)
+        self.pool2_3 = nn.AvgPool2d(kernel_size=(1, 4))
+        
+        self.conv2_4 = nn.Conv2d(out_chans, out_chans, kernel_size=(1, 3), padding=(0, 1))
+        self.norm2_4 = nn.GroupNorm(4, out_chans)
+        self.pool2_4 = nn.AvgPool2d(kernel_size=(1, 4))
+        self.gelu2 = nn.GELU()
+                                
+    def forward(self, x):
+        x = rearrange(x, 'B N A T -> B (N A) T')
+        x = x.unsqueeze(1)
+        
+        # First layer of filtering - Group 1
+        x1 = self.pool1_1(self.gelu1(self.norm1_1(self.conv1_1(x))))
+        x2 = self.pool1_2(self.gelu1(self.norm1_2(self.conv1_2(x))))
+        x3 = self.pool1_3(self.gelu1(self.norm1_3(self.conv1_3(x))))
+        x4 = self.pool1_4(self.gelu1(self.norm1_4(self.conv1_4(x))))
+        
+        # First layer of filtering - Group 2
+        x1 = self.pool2_1(self.gelu2(self.norm2_1(self.conv2_1(x1))))
+        x2 = self.pool2_2(self.gelu2(self.norm2_2(self.conv2_2(x2))))
+        x3 = self.pool2_3(self.gelu2(self.norm2_3(self.conv2_3(x3))))
+        x4 = self.pool2_4(self.gelu2(self.norm2_4(self.conv2_4(x4))))
+        
+        # Re-arrange
+        x1 = rearrange(x1, 'B C NA T -> B NA (T C)')
+        x2 = rearrange(x2, 'B C NA T -> B NA (T C)')
+        x3 = rearrange(x3, 'B C NA T -> B NA (T C)')
+        x4 = rearrange(x4, 'B C NA T -> B NA (T C)')
+        return x1, x2, x3, x4
+
+class NeuroRVQFM(nn.Module):
+    """
+    NeuroRVQ Foundation Model Class
+    """
+    def __init__(self, n_patches=256, patch_size=200, in_chans=1, out_chans=8, num_classes=5,
+                 embed_dim=200, depth=12, num_heads=10, mlp_ratio=4., qkv_bias=False, qk_norm=None, drop_rate=0.,
+                 attn_drop_rate=0., drop_path_rate=0., init_values=None, init_scale=0.001,
+                 n_global_electrodes=127, vocab_size=8192, use_as_encoder=True, use_for_pretraining=False):
+       
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+        self.use_for_pretraining = use_for_pretraining
+        self.use_as_encoder = use_as_encoder
+        # Not necessary - legacy code
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+
+        # To identify whether patch_embed layer is used as tokenizer/encoder or as a decoder
+        if use_as_encoder or use_for_pretraining:
+            self.patch_embed = MultiDimentionalTemporalConv(out_chans=out_chans)
+        else:
+            self.patch_embed_1 = PatchEmbed(in_chans=in_chans, embed_dim=embed_dim)
+            self.patch_embed_2 = PatchEmbed(in_chans=in_chans, embed_dim=embed_dim)
+            self.patch_embed_3 = PatchEmbed(in_chans=in_chans, embed_dim=embed_dim)
+            self.patch_embed_4 = PatchEmbed(in_chans=in_chans, embed_dim=embed_dim)
+
+        self.pos_embed = nn.Parameter(torch.zeros(n_global_electrodes + 1, embed_dim), requires_grad=True)
+        self.time_embed = nn.Parameter(torch.zeros(n_patches, embed_dim), requires_grad=True)
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_norm=qk_norm,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=nn.LayerNorm,
+                init_values=init_values, window_size=None)
+            for i in range(depth)])
+
+        # If used for pre-training we need heads to predict 8 * 4 tokens per input sample
+        # TODO: Modular
+        if use_for_pretraining:
+            self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+            self.norm_pre = nn.LayerNorm(embed_dim)
+            self.head_pre_1 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_2 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_3 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_4 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_5 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_6 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_7 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_8 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_9 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_10 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_11 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_12 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_13 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_14 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_15 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_16 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_17 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_18 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_19 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_20 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_21 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_22 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_23 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_24 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_25 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_26 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_27 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_28 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_29 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_30 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_31 = nn.Linear(embed_dim, vocab_size)
+            self.head_pre_32 = nn.Linear(embed_dim, vocab_size)
+        else:
+            self.norm = nn.Identity()
+            self.fc_norm_1 = nn.LayerNorm(embed_dim)
+            self.head_1 = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+            self.fc_norm_2 = nn.LayerNorm(embed_dim)
+            self.head_2 = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+            self.fc_norm_3 = nn.LayerNorm(embed_dim)
+            self.head_3 = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+            self.fc_norm_4 = nn.LayerNorm(embed_dim)
+            self.head_4 = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        # Initialize the weights of the network
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.time_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+
+        # Initialization ONLY during pretraining
+        if use_for_pretraining:
+            trunc_normal_(self.mask_token, std=.02)
+            trunc_normal_(self.head_pre_1.weight, std=.02)
+            trunc_normal_(self.head_pre_2.weight, std=.02)
+            trunc_normal_(self.head_pre_3.weight, std=.02)
+            trunc_normal_(self.head_pre_4.weight, std=.02)
+            trunc_normal_(self.head_pre_5.weight, std=.02)
+            trunc_normal_(self.head_pre_6.weight, std=.02)
+            trunc_normal_(self.head_pre_7.weight, std=.02)
+            trunc_normal_(self.head_pre_8.weight, std=.02)
+            trunc_normal_(self.head_pre_9.weight, std=.02)
+            trunc_normal_(self.head_pre_10.weight, std=.02)
+            trunc_normal_(self.head_pre_11.weight, std=.02)
+            trunc_normal_(self.head_pre_12.weight, std=.02)
+            trunc_normal_(self.head_pre_13.weight, std=.02)
+            trunc_normal_(self.head_pre_14.weight, std=.02)
+            trunc_normal_(self.head_pre_15.weight, std=.02)
+            trunc_normal_(self.head_pre_16.weight, std=.02)
+            trunc_normal_(self.head_pre_17.weight, std=.02)
+            trunc_normal_(self.head_pre_18.weight, std=.02)
+            trunc_normal_(self.head_pre_19.weight, std=.02)
+            trunc_normal_(self.head_pre_20.weight, std=.02)
+            trunc_normal_(self.head_pre_21.weight, std=.02)
+            trunc_normal_(self.head_pre_22.weight, std=.02)
+            trunc_normal_(self.head_pre_23.weight, std=.02)
+            trunc_normal_(self.head_pre_24.weight, std=.02)
+            trunc_normal_(self.head_pre_25.weight, std=.02)
+            trunc_normal_(self.head_pre_26.weight, std=.02)
+            trunc_normal_(self.head_pre_27.weight, std=.02)
+            trunc_normal_(self.head_pre_28.weight, std=.02)
+            trunc_normal_(self.head_pre_29.weight, std=.02)
+            trunc_normal_(self.head_pre_30.weight, std=.02)
+            trunc_normal_(self.head_pre_31.weight, std=.02)
+            trunc_normal_(self.head_pre_32.weight, std=.02)
+        else:
+            if isinstance(self.head_1, nn.Linear):
+                trunc_normal_(self.head_1.weight, std=.02)
+            if isinstance(self.head_1, nn.Linear):
+                self.head_1.weight.data.mul_(init_scale)
+                self.head_1.bias.data.mul_(init_scale)
+            if isinstance(self.head_2, nn.Linear):
+                trunc_normal_(self.head_2.weight, std=.02)
+            if isinstance(self.head_2, nn.Linear):
+                self.head_2.weight.data.mul_(init_scale)
+                self.head_2.bias.data.mul_(init_scale)
+            if isinstance(self.head_3, nn.Linear):
+                trunc_normal_(self.head_3.weight, std=.02)
+            if isinstance(self.head_3, nn.Linear):
+                self.head_3.weight.data.mul_(init_scale)
+                self.head_3.bias.data.mul_(init_scale)
+            if isinstance(self.head_4, nn.Linear):
+                trunc_normal_(self.head_4.weight, std=.02)
+            if isinstance(self.head_4, nn.Linear):
+                self.head_4.weight.data.mul_(init_scale)
+                self.head_4.bias.data.mul_(init_scale)
+
+        self.apply(self._init_weights)
+        self.fix_init_weight()
+
+    # Function to initialize the weights of the network
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    # Function to initialize the weights of the network
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            trunc_normal_(m.weight, std=.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    # Get number of layers from the transformer
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    # Get classification head
+    def get_classifier(self):
+        return self.head
+
+    # Reset the classification head
+    def reset_classifier(self, num_classes):
+        self.num_classes = num_classes
+        self.fc_norm = nn.LayerNorm(self.embed_dim * 4) # multiply dim by 4 for concat [x1,x2,x3,x4] 
+        self.head = nn.Linear(self.embed_dim * 4, num_classes) if num_classes > 0 else nn.Identity()
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token', 'time_embed'}
+
+    def forward(self, x, temporal_embedding_ix, spatial_embedding_ix, return_patch_tokens=False, bool_masked_pos=None, use_for_pretraining=False, branch_idx=0):
+        '''
+        x: [batch size, number of electrodes, number of patches, patch size]
+        For example, for an EEG sample of 4 seconds with 64 electrodes, x will be [batch size, 64, 4, 200]
+        '''
+        if (self.use_as_encoder):
+            x1, x2, x3, x4 = self.patch_embed(x)
+            x = x1
+        else:
+            if (branch_idx==0):
+                x = self.patch_embed_1(x)
+            elif (branch_idx==1):
+                x = self.patch_embed_2(x)
+            elif (branch_idx==2):
+                x = self.patch_embed_3(x)
+            elif (branch_idx==3):
+                x = self.patch_embed_4(x)
+            
+        batch_size, seq_len, _ = x.size()
+
+        # Concatenate the cls token - Legacy code
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+
+        # ONLY in Pre-Training - Masking technique used in LaBraM
+        if (use_for_pretraining or bool_masked_pos is not None):
+            mask_token = self.mask_token.expand(batch_size, seq_len, -1)
+            w = bool_masked_pos.unsqueeze(-1).type_as(mask_token)
+            symmetric_bool_masked_pos = ~bool_masked_pos
+            w_symmetric = symmetric_bool_masked_pos.unsqueeze(-1).type_as(mask_token)
+            
+            x_symmetric_1 = x1 * (1 - w_symmetric) + mask_token * w_symmetric
+            x1 = x1 * (1 - w) + mask_token * w
+            x_symmetric_1 = torch.cat((cls_tokens, x_symmetric_1), dim=1)
+
+            x_symmetric_2 = x2 * (1 - w_symmetric) + mask_token * w_symmetric
+            x2 = x2 * (1 - w) + mask_token * w
+            x_symmetric_2 = torch.cat((cls_tokens, x_symmetric_2), dim=1)
+
+            x_symmetric_3 = x3 * (1 - w_symmetric) + mask_token * w_symmetric
+            x3 = x3 * (1 - w) + mask_token * w
+            x_symmetric_3 = torch.cat((cls_tokens, x_symmetric_3), dim=1)
+
+            x_symmetric_4 = x4 * (1 - w_symmetric) + mask_token * w_symmetric
+            x4 = x4 * (1 - w) + mask_token * w
+            x_symmetric_4 = torch.cat((cls_tokens, x_symmetric_4), dim=1)
+
+        if (self.use_as_encoder):
+            x1 = torch.cat((cls_tokens, x1), dim=1)
+            x2 = torch.cat((cls_tokens, x2), dim=1)
+            x3 = torch.cat((cls_tokens, x3), dim=1)
+            x4 = torch.cat((cls_tokens, x4), dim=1)
+        else:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        # Pad the spatial_embedding_ix - spatial_embedding_ix: (batch, n_patches), spatial_embedding: (n_electrodes + 1, embed_dim)
+        spatial_embedding_ix = F.pad(input=spatial_embedding_ix, pad=(1, 0), mode='constant', value=0)  # for cls token (batch_size, n_patches + 1)
+        # Gets the corresponding pos_embed
+        spatial_embedding = self.pos_embed[spatial_embedding_ix.reshape(-1), :]  # (batch_size * (n_patches + 1), embed_dim)
+        spatial_embedding = spatial_embedding.reshape(spatial_embedding_ix.shape[0], spatial_embedding_ix.shape[1], spatial_embedding.shape[-1])  # (batch_size, n_patches + 1, embed_dim)
+        
+        if (self.use_as_encoder):
+            x1 = x1 + spatial_embedding
+            x2 = x2 + spatial_embedding
+            x3 = x3 + spatial_embedding
+            x4 = x4 + spatial_embedding
+        else:
+            x = x + spatial_embedding
+
+        #  temporal_embedding_ix: (batch, n_patches), temporal_embedding: (n_patches, embed_dim)
+        temporal_embedding = self.time_embed[temporal_embedding_ix.reshape(-1), :]  # (batch_size * (n_patches), embed_dim)
+        temporal_embedding = temporal_embedding.reshape(temporal_embedding_ix.shape[0], temporal_embedding_ix.shape[1], temporal_embedding.shape[-1])  # (batch_size, n_patches, embed_dim)
+        
+        if (self.use_as_encoder):
+            x1[:, 1:, :] += temporal_embedding
+            x1 = self.pos_drop(x1)
+            x2[:, 1:, :] += temporal_embedding
+            x2 = self.pos_drop(x2)
+            x3[:, 1:, :] += temporal_embedding
+            x3 = self.pos_drop(x3)
+            x4[:, 1:, :] += temporal_embedding
+            x4 = self.pos_drop(x4)
+        else:
+            x[:, 1:, :] += temporal_embedding
+            x = self.pos_drop(x)
+            
+        if (self.use_as_encoder):
+            # Pass the transformer blocks
+            for i, x in enumerate([x1, x2, x3, x4]):
+                for blk in self.blocks:
+                    x = blk(x)
+                    
+                if (use_for_pretraining or bool_masked_pos is not None):
+                    x = self.norm_pre(x)
+                else:
+                    x = self.norm(x)
+                # All except cls token
+                if i == 0:
+                    x1 = x[:, 1:, :]
+                elif i == 1:
+                    x2 = x[:, 1:, :]
+                elif i == 2:
+                    x3 = x[:, 1:, :]
+                else:
+                    x4 = x[:, 1:, :]
+        else:
+            # Pass the transformer blocks
+            for blk in self.blocks:
+                x = blk(x)
+            if (use_for_pretraining or bool_masked_pos is not None):
+                x = self.norm_pre(x)
+            else:
+                x = self.norm(x)
+
+            # All except cls token
+            x = x[:, 1:, :]
+
+        if (use_for_pretraining or bool_masked_pos is not None):
+            for i, x_symmetric in enumerate([x_symmetric_1, x_symmetric_2, x_symmetric_3, x_symmetric_4]):
+                x_symmetric += spatial_embedding
+                x_symmetric[:, 1:, :] += temporal_embedding
+                x_symmetric = self.pos_drop(x_symmetric)
+                for blk in self.blocks:
+                    x_symmetric = blk(x_symmetric)
+                x_symmetric = self.norm_pre(x_symmetric)
+                # All except cls token
+                if i == 0:
+                    x_symmetric_1 = x_symmetric[:, 1:, :]
+                elif i == 1:
+                    x_symmetric_2 = x_symmetric[:, 1:, :]
+                elif i == 2:
+                    x_symmetric_3 = x_symmetric[:, 1:, :]
+                else:
+                    x_symmetric_4 = x_symmetric[:, 1:, :]
+
+        # ONLY in Pre-Training
+        if (use_for_pretraining or bool_masked_pos is not None):
+        
+            logits = torch.stack([
+                # x1 heads
+                self.head_pre_1(x1[bool_masked_pos]),
+                self.head_pre_2(x1[bool_masked_pos]),
+                self.head_pre_3(x1[bool_masked_pos]),
+                self.head_pre_4(x1[bool_masked_pos]),
+                self.head_pre_5(x1[bool_masked_pos]),
+                self.head_pre_6(x1[bool_masked_pos]),
+                self.head_pre_7(x1[bool_masked_pos]),
+                self.head_pre_8(x1[bool_masked_pos]),
+
+                # x2 heads
+                self.head_pre_9(x2[bool_masked_pos]),
+                self.head_pre_10(x2[bool_masked_pos]),
+                self.head_pre_11(x2[bool_masked_pos]),
+                self.head_pre_12(x2[bool_masked_pos]),
+                self.head_pre_13(x2[bool_masked_pos]),
+                self.head_pre_14(x2[bool_masked_pos]),
+                self.head_pre_15(x2[bool_masked_pos]),
+                self.head_pre_16(x2[bool_masked_pos]),
+
+                # x3 heads
+                self.head_pre_17(x3[bool_masked_pos]),
+                self.head_pre_18(x3[bool_masked_pos]),
+                self.head_pre_19(x3[bool_masked_pos]),
+                self.head_pre_20(x3[bool_masked_pos]),
+                self.head_pre_21(x3[bool_masked_pos]),
+                self.head_pre_22(x3[bool_masked_pos]),
+                self.head_pre_23(x3[bool_masked_pos]),
+                self.head_pre_24(x3[bool_masked_pos]),
+
+                # x4 heads
+                self.head_pre_25(x4[bool_masked_pos]),
+                self.head_pre_26(x4[bool_masked_pos]),
+                self.head_pre_27(x4[bool_masked_pos]),
+                self.head_pre_28(x4[bool_masked_pos]),
+                self.head_pre_29(x4[bool_masked_pos]),
+                self.head_pre_30(x4[bool_masked_pos]),
+                self.head_pre_31(x4[bool_masked_pos]),
+                self.head_pre_32(x4[bool_masked_pos]),
+            ], dim=0)
+            
+            logits_symmetric = torch.stack([
+                # x_symmetric_1 heads
+                self.head_pre_1(x_symmetric_1[~bool_masked_pos]),
+                self.head_pre_2(x_symmetric_1[~bool_masked_pos]),
+                self.head_pre_3(x_symmetric_1[~bool_masked_pos]),
+                self.head_pre_4(x_symmetric_1[~bool_masked_pos]),
+                self.head_pre_5(x_symmetric_1[~bool_masked_pos]),
+                self.head_pre_6(x_symmetric_1[~bool_masked_pos]),
+                self.head_pre_7(x_symmetric_1[~bool_masked_pos]),
+                self.head_pre_8(x_symmetric_1[~bool_masked_pos]),
+
+                # x_symmetric_2 heads
+                self.head_pre_9(x_symmetric_2[~bool_masked_pos]),
+                self.head_pre_10(x_symmetric_2[~bool_masked_pos]),
+                self.head_pre_11(x_symmetric_2[~bool_masked_pos]),
+                self.head_pre_12(x_symmetric_2[~bool_masked_pos]),
+                self.head_pre_13(x_symmetric_2[~bool_masked_pos]),
+                self.head_pre_14(x_symmetric_2[~bool_masked_pos]),
+                self.head_pre_15(x_symmetric_2[~bool_masked_pos]),
+                self.head_pre_16(x_symmetric_2[~bool_masked_pos]),
+
+                # x_symmetric_3 heads
+                self.head_pre_17(x_symmetric_3[~bool_masked_pos]),
+                self.head_pre_18(x_symmetric_3[~bool_masked_pos]),
+                self.head_pre_19(x_symmetric_3[~bool_masked_pos]),
+                self.head_pre_20(x_symmetric_3[~bool_masked_pos]),
+                self.head_pre_21(x_symmetric_3[~bool_masked_pos]),
+                self.head_pre_22(x_symmetric_3[~bool_masked_pos]),
+                self.head_pre_23(x_symmetric_3[~bool_masked_pos]),
+                self.head_pre_24(x_symmetric_3[~bool_masked_pos]),
+
+                # x_symmetric_4 heads
+                self.head_pre_25(x_symmetric_4[~bool_masked_pos]),
+                self.head_pre_26(x_symmetric_4[~bool_masked_pos]),
+                self.head_pre_27(x_symmetric_4[~bool_masked_pos]),
+                self.head_pre_28(x_symmetric_4[~bool_masked_pos]),
+                self.head_pre_29(x_symmetric_4[~bool_masked_pos]),
+                self.head_pre_30(x_symmetric_4[~bool_masked_pos]),
+                self.head_pre_31(x_symmetric_4[~bool_masked_pos]),
+                self.head_pre_32(x_symmetric_4[~bool_masked_pos]),
+            ], dim=0)
+        
+            return  logits, logits_symmetric
+        else:
+            # ONLY in RVQ
+            if return_patch_tokens:
+                if (self.use_as_encoder):
+                    return self.head_1(self.fc_norm_1(x1)), self.head_2(self.fc_norm_2(x2)), self.head_3(self.fc_norm_3(x3)), self.head_4(self.fc_norm_4(x4)), _
+                else:
+                    if (branch_idx==0):
+                        return self.head_1(self.fc_norm_1(x)), _
+                    elif (branch_idx==1):
+                        return self.head_2(self.fc_norm_2(x)), _
+                    elif (branch_idx==2):
+                        return self.head_3(self.fc_norm_3(x)), _
+                    elif (branch_idx==3):
+                        return self.head_4(self.fc_norm_4(x)), _
+            else:
+                # ONLY in Fine-Tune
+                x = torch.concat([x1,x2,x3,x4], dim=-1)
+                return self.head(self.fc_norm(x.mean(1))), _
+
+
+
+class NeuroRVQTokenizer(nn.Module):
+    """
+    NeuroRVQ Tokenizer
+    """
+    def __init__(self,
+                 encoder_config,
+                 decoder_config,
+                 n_code,
+                 code_dim,
+                 decoder_out_dim
+                 ):
+
+        super().__init__()
+        self.patch_size = encoder_config['patch_size']
+        self.code_dim = code_dim
+
+        # Encoder layer of NeuroRVQFM
+        self.encoder = NeuroRVQFM(n_patches=encoder_config['n_patches'], patch_size=encoder_config['patch_size'],
+            in_chans=encoder_config['in_chans'], out_chans=encoder_config['out_chans_encoder'],
+            num_classes = encoder_config['num_classes'], embed_dim=encoder_config['embed_dim'],
+            depth=encoder_config['depth'], num_heads=encoder_config['num_heads'],
+            mlp_ratio=encoder_config['mlp_ratio'], qkv_bias=encoder_config['qkv_bias'],
+            qk_norm=partial(nn.LayerNorm, eps=1e-6), drop_rate=encoder_config['drop_rate'],
+            attn_drop_rate=encoder_config['attn_drop_rate'], drop_path_rate=encoder_config['drop_path_rate'],
+            init_values=encoder_config['init_values'], init_scale=encoder_config['init_scale'],
+            n_global_electrodes=encoder_config['n_global_electrodes'], vocab_size=n_code,
+            use_as_encoder=True, use_for_pretraining = False)
+
+        # Decoder layer of NeuroRVQFM
+        self.decoder = NeuroRVQFM(n_patches=decoder_config['n_patches'], patch_size=decoder_config['patch_size'],
+            in_chans=decoder_config['in_chans'], out_chans=0,
+            num_classes = decoder_config['num_classes'], embed_dim=decoder_config['embed_dim'],
+            depth=decoder_config['depth'], num_heads=decoder_config['num_heads'],
+            mlp_ratio=decoder_config['mlp_ratio'], qkv_bias=decoder_config['qkv_bias'],
+            qk_norm=partial(nn.LayerNorm, eps=1e-6), drop_rate=decoder_config['drop_rate'],
+            attn_drop_rate=decoder_config['attn_drop_rate'], drop_path_rate=decoder_config['drop_path_rate'],
+            init_values=decoder_config['init_values'], init_scale=decoder_config['init_scale'],
+            n_global_electrodes=decoder_config['n_global_electrodes'], vocab_size=n_code,
+            use_as_encoder=False, use_for_pretraining = False)
+        
+        self.quantize_1 = ResidualVectorQuantization(num_quantizers = 8,
+            n_embed=n_code, embedding_dim=code_dim, beta=1.0, kmeans_init=True, decay=0.99,
+        )
+        self.quantize_2 = ResidualVectorQuantization(num_quantizers = 8,
+            n_embed=n_code, embedding_dim=code_dim, beta=1.0, kmeans_init=True, decay=0.99,
+        )
+        self.quantize_3 = ResidualVectorQuantization(num_quantizers = 8,
+            n_embed=n_code, embedding_dim=code_dim, beta=1.0, kmeans_init=True, decay=0.99,
+        )
+        self.quantize_4 = ResidualVectorQuantization(num_quantizers = 8,
+            n_embed=n_code, embedding_dim=code_dim, beta=1.0, kmeans_init=True, decay=0.99,
+        )
+
+        # Output dimension of the decoder layer
+        self.decoder_out_dim = decoder_out_dim
+
+        # Encoding head after the encoder transformer
+        self.encode_task_layer_1 = nn.Sequential(
+            nn.Linear(encoder_config['embed_dim'], encoder_config['embed_dim']),
+            nn.Tanh(),
+            nn.Linear(encoder_config['embed_dim'], code_dim)
+        )
+        self.encode_task_layer_2 = nn.Sequential(
+            nn.Linear(encoder_config['embed_dim'], encoder_config['embed_dim']),
+            nn.Tanh(),
+            nn.Linear(encoder_config['embed_dim'], code_dim)
+        )
+        self.encode_task_layer_3 = nn.Sequential(
+            nn.Linear(encoder_config['embed_dim'], encoder_config['embed_dim']),
+            nn.Tanh(),
+            nn.Linear(encoder_config['embed_dim'], code_dim)
+        )
+        self.encode_task_layer_4 = nn.Sequential(
+            nn.Linear(encoder_config['embed_dim'], encoder_config['embed_dim']),
+            nn.Tanh(),
+            nn.Linear(encoder_config['embed_dim'], code_dim)
+        )
+        self.encode_task_layer_1.apply(self._init_weights)
+        self.encode_task_layer_2.apply(self._init_weights)
+        self.encode_task_layer_3.apply(self._init_weights)
+        self.encode_task_layer_4.apply(self._init_weights)
+
+        # Decoding heads after the decoder transformer
+        self.decode_task_layer_amplitude = nn.Sequential(
+            nn.Linear(4*decoder_config['embed_dim'], decoder_config['embed_dim']),
+            nn.GELU(),
+            nn.Linear(decoder_config['embed_dim'], self.decoder_out_dim),
+        )
+        self.decode_task_layer_angle_sin = nn.Sequential(
+            nn.Linear(4*decoder_config['embed_dim'], decoder_config['embed_dim']),
+            nn.Tanh(),
+            nn.Linear(decoder_config['embed_dim'], self.decoder_out_dim),
+            nn.Tanh()
+        )
+        self.decode_task_layer_angle_cos = nn.Sequential(
+            nn.Linear(4*decoder_config['embed_dim'], decoder_config['embed_dim']),
+            nn.Tanh(),
+            nn.Linear(decoder_config['embed_dim'], self.decoder_out_dim),
+            nn.Tanh()
+        )
+
+        # Initialize model weights
+        self.decode_task_layer_amplitude.apply(self._init_weights)
+        self.decode_task_layer_angle_sin.apply(self._init_weights)
+        self.decode_task_layer_angle_cos.apply(self._init_weights)
+
+        # MSE loss function
+        self.loss_fn = F.mse_loss
+
+    # Function to initialize the weights of the network
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'quantize.embedding.weight', 'decoder.cls_token', 'decoder.pos_embed', 'decoder.time_embed',
+                'encoder.cls_token', 'encoder.pos_embed', 'encoder.time_embed'}
+
+    def get_number_of_tokens(self):
+        return self.quantize.n_e
+
+    def get_tokens(self, data, temporal_embedding_ix, spatial_embedding_ix):
+        quantize, code_ind, loss, usage_ratios = self.encode(data, temporal_embedding_ix, spatial_embedding_ix)
+        # Convert [8, B*P] to [8, B, P]
+        code_inds = [code_ind_ix.view(8, data.shape[0], -1) for code_ind_ix in code_ind]
+        # Stack all codebooks [4, 8, B, P]
+        stacked_code_inds = torch.stack(code_inds, dim=0)
+        quantize_vecs = [rearrange(quantize_ix, 'b d a c -> b (a c) d').contiguous() for quantize_ix in quantize]
+        output = {}
+        output['token'] = stacked_code_inds
+        output['input_img'] = data
+        output['quantize'] = quantize_vecs
+        return output
+
+    def encode(self, x, temporal_embedding_ix, spatial_embedding_ix):
+        batch_size, n, a, t = x.shape
+        encoder_features_1, encoder_features_2, encoder_features_3, encoder_features_4, _ = self.encoder(x, temporal_embedding_ix=temporal_embedding_ix, spatial_embedding_ix=spatial_embedding_ix, return_patch_tokens=True)
+            
+        with torch.cuda.amp.autocast(enabled=False):
+            to_quantizer_features_1 = self.encode_task_layer_1(encoder_features_1.type_as(self.encode_task_layer_1[-1].weight))
+            to_quantizer_features_2 = self.encode_task_layer_2(encoder_features_2.type_as(self.encode_task_layer_2[-1].weight))
+            to_quantizer_features_3 = self.encode_task_layer_3(encoder_features_3.type_as(self.encode_task_layer_3[-1].weight))
+            to_quantizer_features_4 = self.encode_task_layer_4(encoder_features_4.type_as(self.encode_task_layer_4[-1].weight))
+            
+        N = to_quantizer_features_1.shape[1]
+        h, w = n, N // n
+            
+        # reshape tokens to feature maps for patch embed in decoder
+        to_quantizer_features_1 = rearrange(to_quantizer_features_1, 'b (h w) c -> b c h w', h=h,
+                                            w=w).contiguous()  # reshape for quantizer
+        quantize_1, code_ind_1, loss_1, usage_ratios_1  = self.quantize_1(to_quantizer_features_1)
+            
+        to_quantizer_features_2 = rearrange(to_quantizer_features_2, 'b (h w) c -> b c h w', h=h,
+                                            w=w).contiguous()  # reshape for quantizer
+        quantize_2, code_ind_2, loss_2, usage_ratios_2  = self.quantize_2(to_quantizer_features_2)
+            
+        to_quantizer_features_3 = rearrange(to_quantizer_features_3, 'b (h w) c -> b c h w', h=h,
+                                            w=w).contiguous()  # reshape for quantizer
+        quantize_3, code_ind_3, loss_3, usage_ratios_3  = self.quantize_3(to_quantizer_features_3)
+            
+        to_quantizer_features_4 = rearrange(to_quantizer_features_4, 'b (h w) c -> b c h w', h=h,
+                                            w=w).contiguous()  # reshape for quantizer
+        quantize_4, code_ind_4, loss_4, usage_ratios_4  = self.quantize_4(to_quantizer_features_4)
+            
+        #  Combine loss
+        loss = loss_1 + loss_2 + loss_3 + loss_4
+
+        return [quantize_1, quantize_2, quantize_3, quantize_4], [code_ind_1, code_ind_2, code_ind_3, code_ind_4], loss, [usage_ratios_1, usage_ratios_2, usage_ratios_3, usage_ratios_4]
+
+    def decode(self, quantize, temporal_embedding_ix, spatial_embedding_ix):
+        
+        for i, quantize_i in enumerate(quantize):
+            if i == 0:
+                decoder_features_1, _ = self.decoder(quantize_i, temporal_embedding_ix=temporal_embedding_ix,
+                                spatial_embedding_ix=spatial_embedding_ix, return_patch_tokens=True, branch_idx = 0)
+            elif i == 1:
+                decoder_features_2, _ = self.decoder(quantize_i, temporal_embedding_ix=temporal_embedding_ix,
+                                spatial_embedding_ix=spatial_embedding_ix, return_patch_tokens=True, branch_idx = 1)
+            elif i == 2:
+                decoder_features_3, _ = self.decoder(quantize_i, temporal_embedding_ix=temporal_embedding_ix,
+                                spatial_embedding_ix=spatial_embedding_ix, return_patch_tokens=True, branch_idx = 2)
+            else:
+                decoder_features_4, _ = self.decoder(quantize_i, temporal_embedding_ix=temporal_embedding_ix,
+                                spatial_embedding_ix=spatial_embedding_ix, return_patch_tokens=True, branch_idx = 3)
+        decoder_features = torch.cat([decoder_features_1, decoder_features_2, decoder_features_3, decoder_features_4], dim=2)
+        
+        # Reconstruct Amplitude, Sine and Cosine
+        rec_amplitude = self.decode_task_layer_amplitude(decoder_features)
+        rec_angle_sin = self.decode_task_layer_angle_sin(decoder_features)
+        rec_angle_cos = self.decode_task_layer_angle_cos(decoder_features)
+
+        return rec_amplitude, rec_angle_sin, rec_angle_cos
+
+    def get_codebook_indices(self, x, temporal_embedding_ix, spatial_embedding_ix):
+        return self.get_tokens(x, temporal_embedding_ix, spatial_embedding_ix)['token']
+        
+    def calculate_phase_loss(self, rec_sin, target_sin, rec_cos, target_cos):
+        target_sin = rearrange(target_sin, 'b n a c -> b (n a) c').contiguous()
+        target_cos = rearrange(target_cos, 'b n a c -> b (n a) c').contiguous()
+        rec = torch.stack((rec_cos, rec_sin), dim=-1)
+        target = torch.stack((target_cos, target_sin), dim=-1)
+        #  Cosine Similarity for direction and Enforcing Magnitude loss
+        phase_loss = 1.0 - F.cosine_similarity(rec, target, dim=-1).mean() + 0.1 * ((rec_sin**2 + rec_cos**2 - 1) ** 2).mean()
+        return phase_loss
+
+    def calculate_rec_loss(self, rec, target):
+        target = rearrange(target, 'b n a c -> b (n a) c').contiguous()
+        rec_loss = self.loss_fn(rec, target)
+        return rec_loss
+
+    def calculate_signal_rec_loss(self, rec, target):
+        target = rearrange(target, 'b n a c -> b (n a) c').contiguous()
+        rec = rearrange(rec, 'b n a c -> b (n a) c').contiguous()
+        mse = self.loss_fn(rec, target)
+        return mse
+
+    def std_norm(self, x):
+        mean = torch.mean(x, dim=(1, 2, 3), keepdim=True)
+        std = torch.sqrt(torch.var(x, dim=(1, 2, 3), keepdim=True).clamp(min=1e-8))
+        x = (x - mean) / std
+        return x, mean, std
+
+    def forward(self, x, temporal_embedding_ix, spatial_embedding_ix):
+        """
+        x: shape [B, N, T]
+        """
+        x = rearrange(x, 'B N (A T) -> B N A T', T=self.patch_size).contiguous()
+        x_fft = torch.fft.fft(x, dim=-1)
+
+        # Get the log ampltitude
+        amplitude = torch.abs(x_fft)
+        amplitude = torch.log1p(amplitude)
+        amplitude, amp_mean, amp_std = self.std_norm(amplitude)
+        
+        # Get the sine / cosine of the phase
+        angle = torch.angle(x_fft)
+        sin_angle = torch.sin(angle)
+        cos_angle = torch.cos(angle)
+
+        # Encoding and Quantize
+        quantize, code_ind, code_loss, usage_ratios = self.encode(x, temporal_embedding_ix, spatial_embedding_ix)
+
+        # Decoding
+        xrec_amp, xrec_angle_sin, xrec_angle_cos = self.decode(quantize, temporal_embedding_ix, spatial_embedding_ix)
+        
+        # Reconstruct raw signal from amplitude and sine / cosine
+        ustd_xrec = (rearrange(xrec_amp, 'B N (A T) -> B N A T', T=self.patch_size).contiguous() * amp_std) + amp_mean  # unstandardize
+        ustd_xrec = torch.expm1(ustd_xrec)
+        ustd_xrec = rearrange(ustd_xrec, 'b n a c -> b (n a) c').contiguous()
+        xrec_signal = torch.real(inverse_fft_cos_sin(ustd_xrec, xrec_angle_sin, xrec_angle_cos))
+
+        # Standardize sample and Reconstructed signal for MSE
+        std_x, _, _ = self.std_norm(x)
+        std_xrec_signal, _, _ = self.std_norm(rearrange(xrec_signal, 'B N (A T) -> B N A T', T=self.patch_size).contiguous())
+        signal_rec_loss = self.calculate_signal_rec_loss(std_xrec_signal, std_x)
+
+        # Calculate losses from decoder
+        rec_amplitude_loss = self.calculate_rec_loss(xrec_amp, amplitude)
+        phase_loss = self.calculate_phase_loss(xrec_angle_sin, sin_angle, xrec_angle_cos, cos_angle)
+
+        # Total loss
+        loss = code_loss + rec_amplitude_loss + phase_loss + signal_rec_loss
+
+        std_x = std_x.view(std_x.size(0), -1, 1,std_x.size(-1)).squeeze(2)
+        std_xrec_signal = std_xrec_signal.squeeze(2)
+
+        return std_x, std_xrec_signal
+

NeuroRVQ/NeuroRVQ_modules.py

@@ -0,0 +1,254 @@
+import math
+import torch
+from torch import nn
+import warnings
+import os
+from functools import partial
+from torch.nn import functional as F
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+'''
+Get dictionary for shared encoder and decoder configurations based on the yaml file params.
+'''
+def get_encoder_decoder_params(args):
+    config = dict(patch_size=args['patch_size'], n_patches=args['n_patches'],
+                  n_global_electrodes=args['n_global_electrodes'],
+                  embed_dim=args['embed_dim'], num_heads=args['num_heads_tokenizer'],
+                  mlp_ratio=args['mlp_ratio_tokenizer'],
+                  qkv_bias=args['qkv_bias_tokenizer'], drop_rate=args['drop_rate_tokenizer'],
+                  attn_drop_rate=args['attn_drop_rate_tokenizer'], drop_path_rate=args['drop_path_rate_tokenizer'],
+                  norm_layer=partial(nn.LayerNorm, eps=1e-6), init_values=args['init_values_tokenizer'],
+                  init_scale=args['init_scale_tokenizer'])
+
+    encoder_config = config.copy()
+    encoder_config['in_chans'] = args['in_chans_encoder']
+    encoder_config['depth'] = args['depth_encoder']
+    encoder_config['num_classes'] = 0
+    encoder_config['out_chans_encoder'] = args['out_chans_encoder']
+
+    decoder_config = config.copy()
+    decoder_config['in_chans'] = args['code_dim']
+    decoder_config['depth'] = args['depth_decoder']
+    decoder_config['num_classes'] = 0
+
+    return encoder_config, decoder_config
+
+'''
+Code taken from: https://github.com/huggingface/pytorch-image-models/tree/v0.4.12/timm/models/layers
+'''
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+'''
+Code taken from: https://github.com/huggingface/pytorch-image-models/tree/v0.4.12/timm/models/layers
+'''
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+'''
+Code taken from: https://github.com/huggingface/pytorch-image-models/tree/v0.4.12/timm/models/layers
+'''
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+class Mlp(nn.Module):
+    """
+    MLP module of Transformer based on https://github.com/935963004/LaBraM/
+    """
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+class Attention(nn.Module):
+    """
+    Attention module of Transformer based on https://github.com/935963004/LaBraM/
+    """
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_norm=None, attn_drop=0., proj_drop=0., window_size=None):
+        super().__init__()
+
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        all_head_dim = head_dim * self.num_heads
+        self.scale = head_dim ** -0.5
+        self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        if qk_norm is not None:
+            self.q_norm = qk_norm(head_dim)
+            self.k_norm = qk_norm(head_dim)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+
+        self.window_size = None
+        self.relative_position_bias_table = None
+        self.relative_position_index = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, return_attention=False, return_qkv=False):
+        B, N, C = x.shape
+
+        if self.q_bias is not None:
+            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+        else:
+            qkv_bias = None
+
+        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+
+        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)  # (3, B, H, N, C)
+
+        q, k, v = qkv[0], qkv[1], qkv[2]  # (B, H, N, C)
+        if self.q_norm is not None:
+            q = self.q_norm(q).type_as(v)
+        if self.k_norm is not None:
+            k = self.k_norm(k).type_as(v)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        if return_attention:
+            return attn
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        if return_qkv:
+            return x, qkv
+
+        return x
+
+class Block(nn.Module):
+    """
+    Block of Transformer based on https://github.com/935963004/LaBraM/
+    """
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_norm=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 window_size=None, attn_head_dim=None):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm, attn_drop=attn_drop,
+            proj_drop=drop, window_size=window_size)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x):
+        if self.gamma_1 is None:
+            x = x + self.drop_path(self.attn(self.norm1(x)))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x)))
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x

NeuroRVQ/RVQ.py

@@ -0,0 +1,67 @@
+'''
+Residual Vector Quantization Implementation.
+Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from norm_ema_quantizer import NormEMAVectorQuantizer
+
+class ResidualVectorQuantization(nn.Module):
+    def __init__(self, *, num_quantizers, **kwargs):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [NormEMAVectorQuantizer(**kwargs) for _ in range(num_quantizers)]
+        )
+
+    def forward(self, x):
+        quantized_out = torch.zeros_like(x)
+        residual = x
+
+        all_losses = []
+        all_indices = []
+        n_q = len(self.layers)
+        
+        usage_ratios = []  # Track usage per quantizer
+        total_codes = self.layers[0].num_tokens
+
+        for layer in self.layers[:n_q]:
+            quantized, loss, indices = layer(residual)
+            residual = residual - quantized
+            quantized_out = quantized_out + quantized
+            # Auxilatory Loss
+            loss = loss + 0.4 * F.mse_loss(quantized, residual.detach())
+
+            all_indices.append(indices)
+            all_losses.append(loss)
+           
+            # --- Codebook usage tracking ---
+            unique_codes = torch.unique(indices)
+            usage_ratio = unique_codes.numel() / total_codes
+            usage_ratios.append(float(usage_ratio))
+
+        out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
+        out_losses = out_losses.mean()
+        
+        return quantized_out, out_indices, out_losses, usage_ratios
+
+    def encode(self, x):
+        residual = x
+        all_indices = []
+        n_q = len(self.layers)
+        for layer in self.layers[:n_q]:
+            indices = layer.encode(residual)
+            quantized = layer.decode(indices)
+            residual = residual - quantized
+            all_indices.append(indices)
+        out_indices = torch.stack(all_indices)
+        return out_indices
+
+    def decode(self, q_indices):
+        quantized_out = torch.tensor(0.0, device=q_indices.device)
+        for i, indices in enumerate(q_indices):
+            layer = self.layers[i]
+            quantized = layer.decode(indices)
+            quantized_out = quantized_out + quantized
+        return quantized_out

NeuroRVQ/__init__.py

@@ -0,0 +1,2 @@
+from .NeuroRVQ import NeuroRVQTokenizer
+from .NeuroRVQ import NeuroRVQFM

NeuroRVQ/norm_ema_quantizer.py

@@ -0,0 +1,204 @@
+'''
+This file is based on LaBraM, BEiT-v2, timm, DeiT, and DINO code bases
+https://github.com/935963004/LaBraM/blob/main/norm_ema_quantizer.py
+https://github.com/microsoft/unilm/tree/master/beitv2
+https://github.com/rwightman/pytorch-image-models/tree/master/timm
+https://github.com/facebookresearch/deit/
+https://github.com/facebookresearch/dino
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as distributed
+from einops import rearrange, repeat
+
+def l2norm(t):
+    return F.normalize(t, p = 2, dim = -1)
+
+def ema_inplace(moving_avg, new, decay):
+    moving_avg.data.mul_(decay).add_(new, alpha = (1 - decay))
+
+def sample_vectors(samples, num):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device = device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device = device)
+
+    return samples[indices]
+
+def kmeans(samples, num_clusters, num_iters = 10, use_cosine_sim = False):
+    dim, dtype, device = samples.shape[-1], samples.dtype, samples.device
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        if use_cosine_sim:
+            dists = samples @ means.t()
+        else:
+            diffs = rearrange(samples, 'n d -> n () d') \
+                    - rearrange(means, 'c d -> () c d')
+            dists = -(diffs ** 2).sum(dim = -1)
+
+        buckets = dists.max(dim = -1).indices
+        bins = torch.bincount(buckets, minlength = num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype = dtype)
+        new_means.scatter_add_(0, repeat(buckets, 'n -> n d', d = dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        if use_cosine_sim:
+            new_means = l2norm(new_means)
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+class EmbeddingEMA(nn.Module):
+    def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5, kmeans_init=True, codebook_init_path=''):
+        super().__init__()
+        self.num_tokens = num_tokens
+        self.codebook_dim = codebook_dim
+        self.decay = decay
+        self.eps = eps 
+        if codebook_init_path == '':   
+            if not kmeans_init:
+                weight = torch.randn(num_tokens, codebook_dim)
+                weight = l2norm(weight)
+            else:
+                weight = torch.zeros(num_tokens, codebook_dim)
+            self.register_buffer('initted', torch.Tensor([not kmeans_init]))
+        else:
+            print(f"load init codebook weight from {codebook_init_path}")
+            codebook_ckpt_weight = torch.load(codebook_init_path, map_location='cpu')
+            weight = codebook_ckpt_weight.clone()
+            self.register_buffer('initted', torch.Tensor([True]))
+            
+        self.weight = nn.Parameter(weight, requires_grad = False)
+        self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad = False)
+        self.embed_avg = nn.Parameter(weight.clone(), requires_grad = False)
+        # self.register_buffer('initted', torch.Tensor([not kmeans_init]))
+        self.update = True
+
+    @torch.jit.ignore
+    def init_embed_(self, data):
+        if self.initted:
+            return
+#        print("Performing Kmeans init for codebook")
+        embed, cluster_size = kmeans(data, self.num_tokens, 10, use_cosine_sim = True)
+        self.weight.data.copy_(embed)
+        self.cluster_size.data.copy_(cluster_size)
+        self.initted.data.copy_(torch.Tensor([True]))
+        
+    def forward(self, embed_id):
+        return F.embedding(embed_id, self.weight)
+
+    def cluster_size_ema_update(self, new_cluster_size):
+        self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay)
+
+    def embed_avg_ema_update(self, new_embed_avg): 
+        self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay)
+
+    def weight_update(self, num_tokens):
+        n = self.cluster_size.sum()
+        smoothed_cluster_size = (
+                (self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n
+            )
+        #normalize embedding average with smoothed cluster size
+        embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1)
+        # embed_normalized = l2norm(self.embed_avg / smoothed_cluster_size.unsqueeze(1))
+        self.weight.data.copy_(embed_normalized)   
+
+def norm_ema_inplace(moving_avg, new, decay):
+    moving_avg.data.mul_(decay).add_(new, alpha = (1 - decay))
+    moving_avg.data.copy_(l2norm(moving_avg.data))
+
+class NormEMAVectorQuantizer(nn.Module):
+    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5, 
+                statistic_code_usage=True, kmeans_init=False, codebook_init_path=''):
+        super().__init__()
+        self.codebook_dim = embedding_dim
+        self.num_tokens = n_embed
+        self.beta = beta
+        self.decay = decay
+        
+        # learnable = True if orthogonal_reg_weight > 0 else False
+        self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps, kmeans_init, codebook_init_path)
+        
+        self.statistic_code_usage = statistic_code_usage
+        if statistic_code_usage:
+            self.register_buffer('cluster_size', torch.zeros(n_embed))
+        if distributed.is_available() and distributed.is_initialized():
+            print("ddp is enable, so use ddp_reduce to sync the statistic_code_usage for each gpu!")
+            self.all_reduce_fn = distributed.all_reduce
+        else:
+            self.all_reduce_fn = nn.Identity()
+    
+    def reset_cluster_size(self, device):
+        if self.statistic_code_usage:
+            self.register_buffer('cluster_size', torch.zeros(self.num_tokens))
+            self.cluster_size = self.cluster_size.to(device)
+
+    def forward(self, z):
+        # reshape z -> (batch, height, width, channel) and flatten
+        #z, 'b c h w -> b h w c'
+        z = rearrange(z, 'b c h w -> b h w c')
+        z = l2norm(z)
+        z_flattened = z.reshape(-1, self.codebook_dim)
+        self.embedding.init_embed_(z_flattened)
+        
+        d = z_flattened.pow(2).sum(dim=1, keepdim=True) + \
+            self.embedding.weight.pow(2).sum(dim=1) - 2 * \
+            torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight) # 'n d -> d n'
+        
+        encoding_indices = torch.argmin(d, dim=1)
+
+        z_q = self.embedding(encoding_indices).view(z.shape)
+        
+        encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)     
+        
+        if not self.training:
+            with torch.no_grad():
+                cluster_size = encodings.sum(0)
+                self.all_reduce_fn(cluster_size)
+                ema_inplace(self.cluster_size, cluster_size, self.decay)
+        
+        if self.training and self.embedding.update:
+            #EMA cluster size
+
+            bins = encodings.sum(0)
+            self.all_reduce_fn(bins)
+
+            # self.embedding.cluster_size_ema_update(bins)
+            ema_inplace(self.cluster_size, bins, self.decay)
+
+            zero_mask = (bins == 0)
+            bins = bins.masked_fill(zero_mask, 1.)
+
+            embed_sum = z_flattened.t() @ encodings
+            self.all_reduce_fn(embed_sum)
+                        
+            embed_normalized = (embed_sum / bins.unsqueeze(0)).t()
+            embed_normalized = l2norm(embed_normalized)
+            
+            embed_normalized = torch.where(zero_mask[..., None], self.embedding.weight,
+                                           embed_normalized)
+
+            norm_ema_inplace(self.embedding.weight, embed_normalized, self.decay)
+
+        # compute loss for embedding
+        loss = self.beta * F.mse_loss(z_q.detach(), z) 
+        
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # reshape back to match original input shape
+        #z_q, 'b h w c -> b c h w'
+        z_q = rearrange(z_q, 'b h w c -> b c h w')
+        return z_q, loss, encoding_indices
+    

README.md

@@ -0,0 +1,132 @@
+<div align="center">
+
+<img src="images/banner.png" width="600">
+
+# 🧠NeuroRVQ: Multi-Scale EEG Tokenization for Generative Large Brainwave Models
+
+<a href='https://arxiv.org/abs/2510.13068'><img src='https://img.shields.io/badge/Paper-arXiv-red'></a>
+<a href=''><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-orange'></a>
+
+[Konstantinos Barmpas](https://www.barmpas.com)<sup>1,2</sup> &emsp; [Na Lee](https://www.linkedin.com/in/na-lee-57777387/)<sup>1,2</sup> &emsp; [Alexandros Koliousis](https://akoliousis.com)<sup>3</sup>   
+
+[Yannis Panagakis](http://users.uoa.gr/~yannisp/)<sup>2,4,5</sup> &emsp; [Dimitrios Adamos](https://profiles.imperial.ac.uk/d.adamos)<sup>1,2</sup> &emsp; [Nikolaos Laskaris](https://people.auth.gr/laskaris/?lang=en)<sup>2,6</sup>  &emsp; [Stefanos Zafeiriou](https://profiles.imperial.ac.uk/s.zafeiriou)<sup>1,2</sup>  
+
+<sup>1</sup>Imperial College London, United Kingdom <br>
+<sup>2</sup>Cogitat, United Kingdom <br>
+<sup>3</sup>Northeastern University London, United Kingdom <br>
+<sup>4</sup>National and Kapodistrian University of Athens, Greece <br>
+<sup>5</sup>Archimedes Research Unit, Greece <br>
+<sup>6</sup>Aristotle University of Thessaloniki, Greece 
+
+This is the official implementation of **NeuroRVQ**, a foundation model for biosignals powered by a state-of-the-art biosignal tokenizer
+
+</div>
+
+## 🌟 Overview
+
+**NeuroRVQ Tokenizer** is a specialized network designed to convert raw biosignals into a sequence of compact and informative neural tokens. This transformation reduces the inherently high-dimensional and noisy nature of biosginals into a structured lower-dimensional representation that preserves essential temporal–spectral patterns. In doing so, the tokenizer provides a kind of "neural grammar" for neural activity. The input multi-variate timeseries is first segmented into patches. These patches are encoded by the multi-scale temporal encoder, that captures features in multiple resolutions and are then combined via the transfromer encoder. For each scale, RVQ codebooks discretize the embeddings into a sequence of neural tokens. Finally, these tokens are combined and passed through the tokenizer decoder to reconstruct the input patches using the Fourier spectrum.
+
+**NeuroRVQ Foundation Model** is a scalable foundation model that operates on the tokenized representation. By working at the token level rather than raw signals, this model can better capture long-range dependencies, learn abstract neural dynamics and enable efficient pretraining across diverse EEG datasets. The model leverages the learned codebooks from the tokenizer stage and is trained using a masked-token prediction strategy, where a subset of input patches is randomly masked. This objective encourages the network to infer missing tokens from their surrounding context.
+
+## Model and Modalities
+
+| Modality | Support |
+| :--- | :--- |
+| **EEG** | ✅ |
+| **EMG** | ✅ |
+| **ECG** | ✅ |
+
+| Model Version | Parameters | Modality | Trained Models <a href=''><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-orange'></a> | 
+| :--- | :--- | :--- | :--- |
+| **NeuroRVQ-EEG-tokenizer-v1** | 76 Million | EEG | NeuroRVQ_EEG_tokenizer_v1.pt |
+| **NeuroRVQ-EEG-foundation-model-v1** | 6 Million | EEG | NeuroRVQ_EEG_foundation_model_v1.pt |
+| **NeuroRVQ-EMG-tokenizer-v1** | - | EMG | "To be released soon..." |
+| **NeuroRVQ-EMG-foundation-model-v1** | - | EMG | "Training in Progress..." |
+| **NeuroRVQ-ECG-tokenizer-v1** | - | ECG | "To be released soon..." |
+| **NeuroRVQ-ECG-foundation-model-v1** | - | ECG | "Training in Progress..." |
+
+## Tokenization / Reconstruction Capabilities
+
+| EEG | ECG | EMG |
+|:---:|:---:|:---:|
+| <img src="images/eeg.png" width="300"/> | <img src="images/ecg.png" width="300"/> | <img src="images/emg.png" width="350"/> |
+
+## Downstream Performance
+
+### EEG
+
+| Model      | Motor | ERP | Memory | Sleep* | Eyes | Mean | Size |
+|-----------|-------|------|---------|---------|-------|--------|-------|
+| NeuroGPT  | <u>0.682±0.083</u> | 0.757±0.048 | **0.597±0.029** | <u>0.674±0.033</u> | 0.827±0.036 | <u>0.707±0.046</u> | 79.5M |
+| CBraMod   | 0.614±0.104 | 0.777±0.052 | <u>0.574±0.038</u> | 0.635±0.041 | <u>0.839±0.041</u> | 0.688±0.055 | 4.9M |
+| BIOT      | 0.443±0.079 | 0.500±0.000 | 0.510±0.018 | -- | 0.763±0.049 | -- | 3.2M |
+| MIRepNet  | 0.689±0.086 | -- | -- | -- | -- | -- | -- |
+| LaBraM    | 0.630±0.076 | <u>0.822±0.040</u> | 0.526±0.026 | 0.652±0.037 | 0.799±0.047 | 0.686±0.045 | 5.8M |
+| EEGPT     | 0.313±0.035 | 0.668±0.146 | 0.520±0.017 | 0.634±0.044 | 0.797±0.037 | 0.587±0.056 | 25.7M |
+| **NeuroRVQ**  | **0.700±0.073** | **0.876±0.033** | <u>0.574±0.027</u> | **0.728±0.028** | **0.869±0.026** | **0.749±0.037** | 5.9M |
+
+We used the benchmark presented in IEEE MLSP 2025 Paper [Assessing the Capabilities of Large Brainwave Foundation Models](https://ieeexplore.ieee.org/document/11204282).
+
+#### About the Benchmark
+Over the last decade, deep learning models have been widely used for automatic feature extraction and classification in various Brain-Computer Interface (BCI) tasks. However, their performance and generalization capabilities are often not adequately assessed, as these models are frequently trained and tested under flawed setups and / or influenced by spurious correlations. Recently, these limitations have also been observed in the training and evaluation of Large Brainwave Foundation Models (LBMs). In this work, we employ causal reasoning and careful consideration for task-discriminative artifacts in various EEG datasets covering diverse BCI paradigms and propose a benchmarking protocol to properly evaluate the decoding performance and generalization capabilities of LBMs. Utilising a subject-independent cross-validation approach for each curated benchmark dataset, we showcase that LBMs achieve marginal performance gains over conventional deep learning baselines.
+
+[Open-Source Benchmark Code](https://github.com/dykestra/EEG-Benchmarking)
+
+## Installation
+```bash
+conda create -n neurorvq python=3.10
+conda activate neurorvq
+
+# Install requirements
+pip install -r requirements.txt
+```
+
+## Download Models
+The models and the sample biosignal for reconstruction demos can be downloaded manually from [HuggingFace]() or using python:
+```python
+from huggingface_hub import hf_hub_download
+
+hf_hub_download(repo_id="", filename="pretrained_models/tokenizers/NeuroRVQ_EEG_tokenizer_v1.pt", local_dir="./pretrained_models/tokenizers")
+hf_hub_download(repo_id="", filename="pretrained_models/foundation_models/NeuroRVQ_EEG_foundation_model_v1.pt", local_dir="./pretrained_models/foundation_models")
+hf_hub_download(repo_id="", filename="example_files/eeg_sample/example_eeg_file.xdf", local_dir="./example_files/eeg_sample")
+```
+
+## Model Loading / Usage
+
+Load tokenizer and see reconstruction results. Example for EEG tokenizer:
+```python
+
+from inference.run.NeuroRVQ_EEG_tokenizer_example import load_neurorqv_tokenizer
+
+# Set run_example=True and plot_results=True to see reconstruction results
+# Checkout the load_neurorqv_tokenizer() function to load and use tokenizer
+
+load_neurorqv_tokenizer(run_example=True, plot_results=True, verbose=True,
+                            model_path='./pretrained_models/tokenizers/NeuroRVQ_EEG_tokenizer_v1.pt')
+```
+
+Load foundation model and see an example for fine-tuning. Example for EEG foundation model:
+```python
+
+from inference.run.NeuroRVQ_EEG_FM_example import load_neurorqv_fm
+
+# Checkout the load_neurorqv_fm() function with fine_tuning=False to see the correct model loading
+# See the instructions in data.py for your custom dataset before setting fine_tuning=True
+
+load_neurorqv_fm(fine_tuning=False, verbose=True,
+                     model_path = './pretrained_models/foundation_models/NeuroRVQ_EEG_foundation_model_v1.pt')
+```
+
+## Citation
+```
+@misc{neurorvq,
+      title={NeuroRVQ: Multi-Scale EEG Tokenization for Generative Large Brainwave Models}, 
+      author={Konstantinos Barmpas and Na Lee and Alexandros Koliousis and Yannis Panagakis and Dimitrios A. Adamos and Nikolaos Laskaris and Stefanos Zafeiriou},
+      year={2025},
+      eprint={2510.13068},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2510.13068}, 
+}
+```
+

example_files/eeg_sample/example_eeg_file.xdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a94b2b99371a32ecc7386ef5eb84555079616609c45fbb50ffe1a3b90c19ba3
+size 125751622

fine_tuning/NeuroRVQ_EEG_FM_FineTuning.py

@@ -0,0 +1,163 @@
+import numpy as np
+import torch
+from inference.modules.NeuroRVQ_EEG_tokenizer_inference_modules import ch_names_global, create_embedding_ix, check_model_eval_mode
+import torch.nn.functional as F
+from sklearn.metrics import accuracy_score, balanced_accuracy_score
+from torch.utils.data import DataLoader
+import warnings
+from tqdm import tqdm
+
+def get_class_weights(y, n_cls):
+    y = torch.Tensor(y)
+    class_weights = torch.unique(y, return_counts=True)[1]
+    class_weights = 1 / class_weights
+    class_weights = class_weights / class_weights.sum()
+    class_weights = class_weights * len(torch.unique(y))  # (n_classes,)
+    if len(class_weights) < n_cls:
+        tmp = class_weights
+        class_weights = torch.zeros(n_cls)
+        class_weights[:len(tmp)] = tmp
+    class_weights = class_weights.cuda()
+    return class_weights
+
+class NeuroRVQModule():
+    '''
+    Module that performs fine-tuning of NeuroRVQ
+    '''
+    def __init__(self, sample_length, chnames, n_out, train_head_only, args, foundation_model):
+        self.n_time = sample_length // args['patch_size']
+        chnames = np.array([c.lower().encode() for c in chnames])
+        self.chmask = np.isin(chnames, ch_names_global)
+        self.chnames = chnames[self.chmask]
+        self.n_out = n_out
+        self.model = foundation_model
+        self.train_head_only = train_head_only
+        self.criterion = F.cross_entropy if self.n_out > 2 else F.binary_cross_entropy_with_logits
+        self.results = {'train_accuracy': [], 'val_accuracy': [], 'train_bacc': [], 'val_bacc': []}
+        self.weight_decay = args['weight_decay_finetuning']
+        self.warmup_epochs = args['warmup_epochs_finetuning']
+        self.amp_dtype = torch.bfloat16
+        self.lr = float(args['lr_finetuning'])
+        self.layer_decay = float(args['layer_decay_finetuning'])
+        self.n_patches = args['n_patches']
+        self.patch_size = args['patch_size']
+
+    def size(self):
+        """ Returns number of trainable parameters in model """
+        return sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+
+    def fit(self, train_dataset, validation_dataset, batch_size, epochs):
+        d_out = self.n_out if self.n_out > 2 else 1
+        self.model.reset_classifier(d_out)
+        self.model.cuda()
+        # Set model parameter groups with layer_decay on the learning rate
+        if self.train_head_only:
+            for name, param in self.model.named_parameters():
+                if 'head.' in name or 'fc_norm.' in name:
+                    continue
+                else:
+                    param.requires_grad = False
+
+        param_groups = {}
+        for i_m, (p_name, param) in enumerate(self.model.named_parameters()):  # model layers
+            if not param.requires_grad:
+                continue
+            if ('head.' in p_name) or ('fc_norm.' in p_name):  # normal lr for classification head
+                param_groups[p_name] = {'params': [param],
+                                        'weight_decay': self.weight_decay,
+                                        'lr': self.lr}
+            else:
+                param_groups[p_name] = {'params': [param],
+                                        'weight_decay': self.weight_decay,
+                                        'lr': self.lr * self.layer_decay ** (
+                                                len(list(self.model.named_parameters())) - i_m)}
+
+        # Optimizer and lr_scheduler
+        optimizer = torch.optim.AdamW(list(param_groups.values()))
+        n_batches_train = int(np.ceil(len(train_dataset) / batch_size))
+        if epochs < self.warmup_epochs + 1:
+            lr_scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1e-1, end_factor=1,
+                                                             total_iters=epochs * n_batches_train)
+        else:
+            scheduler1 = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1e-1, end_factor=1,
+                                                           total_iters=self.warmup_epochs * n_batches_train)
+            scheduler2 = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1, end_factor=1e-1,
+                                                           total_iters=(epochs - self.warmup_epochs) * n_batches_train)
+            lr_scheduler = torch.optim.lr_scheduler.SequentialLR(optimizer, [scheduler1, scheduler2],
+                                                                 milestones=[self.warmup_epochs * n_batches_train])
+        warnings.filterwarnings('ignore', category=UserWarning, module='torch.optim.lr_scheduler')
+        # Prepare automatic mixed precision training
+        scaler = torch.cuda.amp.GradScaler()
+
+        y_train = [ys for _, ys in train_dataset]
+        y_val = [ys for _, ys in validation_dataset]
+        y = y_train + y_val
+        class_weights = get_class_weights(y, self.n_out)
+
+        train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+        val_dataloader = DataLoader(validation_dataset, batch_size=batch_size, shuffle=False)
+
+        temp_embed_ix, spat_embed_ix = create_embedding_ix(self.n_time, self.n_patches,
+                                                           self.chnames, ch_names_global)
+
+        # Loop over epochs
+        for i_epoch in range(epochs):
+            print(f"Epoch {i_epoch}")
+            # Loop over training batches
+            self.model.train()
+            e_pred_train = []  # collect predictions
+            y_true_train = []  # y in order seen
+            for x_b, y_b in tqdm(train_dataloader):
+                x_b = x_b[:, self.chmask, :]
+                n, c, t = x_b.shape
+                x_b = x_b.reshape(n, c, self.n_time, self.patch_size).cuda()
+                y_b = y_b.long() if self.n_out > 2 else y_b.float()
+                with torch.amp.autocast(device_type='cuda', dtype=self.amp_dtype):
+                    optimizer.zero_grad()
+                    p, _ = self.model(x_b, temp_embed_ix, spat_embed_ix)
+                    p = p.squeeze(-1)  # remove class dim if binary task
+                    loss_weight = class_weights if p.ndim == 2 else class_weights[y_b.long()]
+                    loss = self.criterion(p, y_b.cuda(), weight=loss_weight)
+
+
+                scaler.scale(loss).backward()
+                scaler.step(optimizer)
+                scaler.update()
+                lr_scheduler.step()
+
+                # Collect class predictions to compute metrics on the full epoch
+                p = p.detach().cpu().float()
+                p = p.argmax(dim=-1) if p.ndim == 2 else torch.round(torch.sigmoid(p))
+                e_pred_train += [p.numpy()]
+                y_true_train += [y_b.numpy()]
+
+            # Loop over validation batches
+            self.model.eval()
+            e_pred_val = []  # collect predictions
+            y_true_val = []  # y in order seen
+            for x_b, y_b in tqdm(val_dataloader):
+                x_b = x_b[:, self.chmask, :]
+                n, c, t = x_b.shape
+                x_b = x_b.reshape(n, c, self.n_time, self.patch_size).cuda()
+                with torch.amp.autocast(device_type='cuda', dtype=self.amp_dtype):
+                    p, _ = self.model(x_b, temp_embed_ix, spat_embed_ix)
+                    p = p.squeeze(-1)  # remove class dim if binary task
+
+                # Collect class predictions to compute metrics on the full epoch
+                p = p.detach().cpu().float()
+                p = p.argmax(dim=-1) if p.ndim == 2 else torch.round(torch.sigmoid(p))
+                e_pred_val += [p.numpy()]
+                y_true_val += [y_b.numpy()]
+
+            # Compute accuracy and balanced accuracy
+            e_pred_train = np.concatenate(e_pred_train)
+            e_pred_val = np.concatenate(e_pred_val)
+            y_true_train = np.concatenate(y_true_train)
+            y_true_val = np.concatenate(y_true_val)
+
+            self.results['train_accuracy'] += [accuracy_score(y_true_train, e_pred_train)]
+            self.results['val_accuracy'] += [accuracy_score(y_true_val, e_pred_val)]
+            self.results['train_bacc'] += [balanced_accuracy_score(y_true_train, e_pred_train)]
+            self.results['val_bacc'] += [balanced_accuracy_score(y_true_val, e_pred_val)]
+            if len(validation_dataset) > 1:
+                print(f"VAL ACC: {self.results['val_accuracy'][-1]}, VAL BACC: {self.results['val_bacc'][-1]}")

fine_tuning/data.py

@@ -0,0 +1,87 @@
+import numpy as np
+import pandas as pd
+from abc import ABC
+
+class Benchmark(ABC):
+    """
+    Class for benchmark dataset with expected properties:
+        eeg: array of EEG data (samples, channels, time)
+        subject_ids: array of subject ID for each data sample (samples,)
+        labels: array of target class labels for each data sample (samples,)
+        chnames: array of electrode channel names (channels,)
+    """
+    def __init__(self):
+        self.eeg = None
+        self.subject_ids = None
+        self.labels = None
+        self.chnames = None
+
+    def get_data(self):
+        return self.eeg, self.subject_ids, self.labels, self.chnames
+
+    def sample_balanced_set(self, idx, seed):
+        """
+        Performs a random sampling of indices to balance classes for each subject
+            idx: array of sample indices relative to self.eeg
+            seed: random seed for sampling
+        Returns:
+            filtered indices after random sampling
+        """
+        rng = np.random.default_rng(seed)
+        subj_all = self.subject_ids[idx]
+        y_all = self.labels[idx]
+        sampled = []
+
+        for s in np.unique(subj_all):
+            mask_s = (subj_all == s)
+            idx_s = idx[mask_s]
+            y_s = y_all[mask_s]
+
+            labels = np.unique(y_s)
+
+            idx_by_label = [idx_s[y_s == label] for label in labels]
+
+            # minority per subject
+            n = min([len(idx_l) for idx_l in idx_by_label])
+            if n == 0:
+                continue
+
+            take_by_label = [rng.choice(idx_l, size=n, replace=False) for idx_l in idx_by_label]
+            sampled.append(np.concatenate(take_by_label))
+
+        sampled_idx = np.concatenate(sampled)
+        return sampled_idx
+
+class YourCustomBenchmark(Benchmark):
+    """
+    Custom Class Example where your eeg trials are in stored in .npy file
+    The labels and other info in the .pd file
+    And your dasaset has 4-classes
+    """
+    def __init__(self, root, subdir, apply_car):
+        super().__init__()
+        print("Loading Your Data...")
+        eeg = np.load('./fine_tuning/data/data_eeg.npy', mmap_mode='r')
+        tf = pd.read_pickle('./fine_tuning/data/trial_features.pd')
+        subject_ids = tf['subject_id'].to_numpy()
+        chnames = np.array([c.upper() for c in tf.attrs['channel_names']])
+        labels = tf['task'].replace({'class_1': 0, 'class_2': 1, 'class_3': 2, 'class_4': 3}).to_numpy()
+
+        self.eeg = eeg
+        self.subject_ids = subject_ids
+        self.labels = labels
+        self.chnames = chnames
+
+    def sample_balanced_set(self, idx, seed):
+        print("Classes are already balanced for High Gamma")
+        return idx
+
+def load_benchmark(benchmark, root, subdir, apply_car=False) -> Benchmark:
+    BENCHMARK_CLASSES = {
+        "Custom Benchmark": YourCustomBenchmark
+    }
+
+    assert (benchmark in BENCHMARK_CLASSES), f"Unsupported benchmark {benchmark}. Make sure load function is added to BENCHMARK_LOADERS."
+
+    benchmark_cls = BENCHMARK_CLASSES[benchmark]
+    return benchmark_cls(root, subdir, apply_car)

fine_tuning/utils.py

@@ -0,0 +1,51 @@
+import os
+import csv
+from fine_tuning.wrappers import NeuroRVQWrapper
+
+class CSVLogger():
+    def __init__(self, output_dir, ex_id):
+        self.log_dir = os.path.join(output_dir, f"{ex_id}_log")
+        if not os.path.exists(self.log_dir):
+            os.makedirs(self.log_dir)
+        self._files = set()
+    
+    def report_scalar(self, title, series, value, iteration):
+        '''
+        Mimics clearml report_scalar() function to log values to CSV file
+        '''
+        if 'train' in series:
+            filepath = os.path.join(self.log_dir, f"{title}_train.csv")
+        else:
+            filepath = os.path.join(self.log_dir, f"{title}_val.csv")
+
+        write_header = filepath not in self._files
+
+        with open(filepath, mode="a", newline="") as f:
+            writer = csv.writer(f)
+            if 'MEAN' in title:
+                if write_header:
+                    writer.writerow(["Series", "Iteration", "Value"])
+                    self._files.add(filepath)
+                writer.writerow([series, iteration, value])
+            else:
+                if write_header:
+                    writer.writerow(["Fold", "Iteration", "Value"])
+                    self._files.add(filepath)
+                writer.writerow([series.split(' ')[-1], iteration, value])
+
+def get_logger():
+    logger = CSVLogger("results", 0)
+    return logger
+
+def get_model(ch_names, n_times, n_outputs, args, foundation_model, train_head_only=False):
+    """
+    Returns: FinetuningWrapper for the specified model
+    """
+    return NeuroRVQWrapper(
+        n_time=n_times,
+        ch_names=ch_names,
+        n_outputs=n_outputs,
+        train_head_only=train_head_only,
+        args = args,
+        foundation_model = foundation_model
+    )

fine_tuning/wrappers.py

@@ -0,0 +1,45 @@
+'''
+Wrapper classes of foundation model modules for use in main benchmarking script
+'''
+import torch
+from abc import ABC, abstractmethod
+from fine_tuning.NeuroRVQ_EEG_FM_FineTuning import NeuroRVQModule
+
+class FinetuningWrapper(ABC):
+    """
+    Wrapper class for initializing model, fitting and evaluating on benchmark data, and storing results
+    """
+    def __init__(self):
+        self.model = None
+        self.results = {}
+
+    @abstractmethod
+    def fit(self, train_dataset, validation_dataset, batch_size, epochs):
+        print("fit function not implemented")
+
+    def size(self):
+        """ Returns number of trainable parameters in model """
+        if self.model is None:
+            print("model not initialised")
+        else:
+            return self.model.size()
+
+class NeuroRVQWrapper(FinetuningWrapper):
+    def __init__(self, n_time, ch_names, n_outputs, train_head_only, args, foundation_model):
+        super().__init__()
+        self.model = NeuroRVQModule(
+            sample_length=n_time, 
+            chnames=ch_names,
+            n_out=n_outputs,
+            train_head_only=train_head_only,
+            args = args,
+            foundation_model = foundation_model
+            )
+
+    def fit(self, train_dataset, validation_dataset, batch_size, epochs):
+        self.model.fit(train_dataset, validation_dataset, batch_size, epochs)
+        self.results = self.model.results
+        
+    def save_model(self, path):
+        print(f'Saving checkpoint to {path}...')
+        torch.save(self.model.model.state_dict(), path)

flags/NeuroRVQ_EEG_v1.yml

@@ -0,0 +1,47 @@
+patch_size: 200
+n_patches: 256
+num_classes: 5
+n_code: 8192
+code_dim: 128
+embed_dim: 200
+
+weight_decay_finetuning: 1e-2
+warmup_epochs_finetuning: 4
+lr_finetuning: 5e-4
+layer_decay_finetuning: 0.975
+train_head_only_finetuning: False
+batch_size_finetuning: 32
+epoch_finetuning: 20
+model_name: NeuroRVQ
+
+in_chans_encoder: 1
+out_chans_encoder: 8
+depth_encoder: 12
+depth_decoder: 3
+decoder_out_dim: 200
+
+num_heads_tokenizer: 10
+mlp_ratio_tokenizer: 4
+qkv_bias_tokenizer: True
+drop_rate_tokenizer: 0.
+attn_drop_rate_tokenizer: 0.
+drop_path_rate_tokenizer: 0.
+init_values_tokenizer: 0.
+init_scale_tokenizer: 0.001
+
+use_for_pretraining: True
+
+in_chans_second_stage: 1
+out_chans_second_stage: 8
+depth_second_stage: 12
+num_heads_second_stage: 10
+mlp_ratio_second_stage: 4.
+qkv_bias_second_stage: True
+drop_rate_second_stage: 0.
+attn_drop_rate_second_stage: 0.
+drop_path_rate_second_stage: 0.
+init_values_second_stage: 1.e-5
+init_scale_second_stage: 0.001
+embed_dim_second_stage: 200
+
+

inference/modules/NeuroRVQ_EEG_FM_inference_modules.py

@@ -0,0 +1,48 @@
+import numpy as np
+import torch
+
+# List of all channels used in the pre-training of "NeuroRVQ_EEG_v1" model
+ch_names_global = np.array([b'a1', b'a2', b'af3', b'af4', b'af7', b'af8', b'afz', b'c1', b'c2',
+                            b'c3', b'c4', b'c5', b'c6', b'ccp1', b'ccp2', b'ccp3', b'ccp4',
+                            b'ccp5', b'ccp6', b'ccp7', b'ccp8', b'cfc1', b'cfc2', b'cfc3',
+                            b'cfc4', b'cfc5', b'cfc6', b'cfc7', b'cfc8', b'cp1', b'cp2',
+                            b'cp3', b'cp4', b'cp5', b'cp6', b'cpz', b'cz', b'eog', b'f1',
+                            b'f10', b'f2', b'f3', b'f4', b'f5', b'f6', b'f7', b'f8', b'f9',
+                            b'fc1', b'fc2', b'fc3', b'fc4', b'fc5', b'fc6', b'fcz', b'fp1',
+                            b'fp2', b'fpz', b'ft7', b'ft8', b'fz', b'iz', b'loc', b'o1', b'o2',
+                            b'oz', b'p08', b'p1', b'p10', b'p2', b'p3', b'p4', b'p5', b'p6',
+                            b'p7', b'p8', b'p9', b'po1', b'po10', b'po2', b'po3', b'po4',
+                            b'po7', b'po8', b'po9', b'poz', b'pz', b'roc', b'sp1', b'sp2',
+                            b't1', b't10', b't2', b't3', b't4', b't5', b't6', b't7', b't8',
+                            b't9', b'tp10', b'tp7', b'tp8', b'tp9'])
+
+def check_model_eval_mode(model):
+    for name, module in model.named_modules():
+        if hasattr(module, 'training'):
+            if module.training:
+                print(f"[WARNING] Module {name} is still in training mode.")
+            else:
+                print(f"[OK] Module {name} is in eval mode.")
+
+def create_embedding_ix(n_time, max_n_patches, ch_names_sample, ch_names_global):
+    """Creates temporal and spatial embedding indices for a sample with given regular shape.
+    Args:
+        n_time: Int. Number of patches along the time dimension
+        max_n_patches: The maximum number of patches, for aligning the current time-point to the right.
+        ch_names_sample (n_channels_sample,): The specific channel names of the sample
+        ch_names_global (n_channels_global): The reference channel names of the model
+    Returns:
+        temp_embed_ix (1, n_patches): tensor
+        spat_embed_ix (1, n_patches): tensor
+    """
+    # Temporal embedding ix
+    temp_embed_ix = torch.arange(max_n_patches - n_time, max_n_patches)
+    temp_embed_ix = temp_embed_ix.repeat(len(ch_names_sample))
+    temp_embed_ix = temp_embed_ix.reshape(1, -1)
+
+    # Spatial embedding ix
+    spat_embed_ix = torch.tensor([np.where(ch_names_global == c)[0][0] for c in ch_names_sample])
+    spat_embed_ix = torch.repeat_interleave(spat_embed_ix, n_time)
+    spat_embed_ix = spat_embed_ix.reshape(1, -1)
+
+    return temp_embed_ix, spat_embed_ix

inference/modules/NeuroRVQ_EEG_tokenizer_inference_modules.py

@@ -0,0 +1,48 @@
+import numpy as np
+import torch
+
+# List of all channels used in the pre-training of "NeuroRVQ_EEG_v1" model
+ch_names_global = np.array([b'a1', b'a2', b'af3', b'af4', b'af7', b'af8', b'afz', b'c1', b'c2',
+                            b'c3', b'c4', b'c5', b'c6', b'ccp1', b'ccp2', b'ccp3', b'ccp4',
+                            b'ccp5', b'ccp6', b'ccp7', b'ccp8', b'cfc1', b'cfc2', b'cfc3',
+                            b'cfc4', b'cfc5', b'cfc6', b'cfc7', b'cfc8', b'cp1', b'cp2',
+                            b'cp3', b'cp4', b'cp5', b'cp6', b'cpz', b'cz', b'eog', b'f1',
+                            b'f10', b'f2', b'f3', b'f4', b'f5', b'f6', b'f7', b'f8', b'f9',
+                            b'fc1', b'fc2', b'fc3', b'fc4', b'fc5', b'fc6', b'fcz', b'fp1',
+                            b'fp2', b'fpz', b'ft7', b'ft8', b'fz', b'iz', b'loc', b'o1', b'o2',
+                            b'oz', b'p08', b'p1', b'p10', b'p2', b'p3', b'p4', b'p5', b'p6',
+                            b'p7', b'p8', b'p9', b'po1', b'po10', b'po2', b'po3', b'po4',
+                            b'po7', b'po8', b'po9', b'poz', b'pz', b'roc', b'sp1', b'sp2',
+                            b't1', b't10', b't2', b't3', b't4', b't5', b't6', b't7', b't8',
+                            b't9', b'tp10', b'tp7', b'tp8', b'tp9'])
+
+def check_model_eval_mode(model):
+    for name, module in model.named_modules():
+        if hasattr(module, 'training'):
+            if module.training:
+                print(f"[WARNING] Module {name} is still in training mode.")
+            else:
+                print(f"[OK] Module {name} is in eval mode.")
+
+def create_embedding_ix(n_time, max_n_patches, ch_names_sample, ch_names_global):
+    """Creates temporal and spatial embedding indices for a sample with given regular shape.
+    Args:
+        n_time: Int. Number of patches along the time dimension
+        max_n_patches: The maximum number of patches, for aligning the current time-point to the right.
+        ch_names_sample (n_channels_sample,): The specific channel names of the sample
+        ch_names_global (n_channels_global): The reference channel names of the model
+    Returns:
+        temp_embed_ix (1, n_patches): tensor
+        spat_embed_ix (1, n_patches): tensor
+    """
+    # Temporal embedding ix
+    temp_embed_ix = torch.arange(max_n_patches - n_time, max_n_patches)
+    temp_embed_ix = temp_embed_ix.repeat(len(ch_names_sample))
+    temp_embed_ix = temp_embed_ix.reshape(1, -1)
+
+    # Spatial embedding ix
+    spat_embed_ix = torch.tensor([np.where(ch_names_global == c)[0][0] for c in ch_names_sample])
+    spat_embed_ix = torch.repeat_interleave(spat_embed_ix, n_time)
+    spat_embed_ix = spat_embed_ix.reshape(1, -1)
+
+    return temp_embed_ix, spat_embed_ix

inference/run/NeuroRVQ_EEG_FM_example.py

@@ -0,0 +1,107 @@
+import yaml
+import numpy as np
+import torch
+from NeuroRVQ.NeuroRVQ import NeuroRVQFM
+from NeuroRVQ.NeuroRVQ_modules import get_encoder_decoder_params
+from inference.modules.NeuroRVQ_EEG_tokenizer_inference_modules import ch_names_global, create_embedding_ix, check_model_eval_mode
+from functools import partial
+from torch import nn
+from fine_tuning.utils import get_logger, get_model
+from fine_tuning.data import load_benchmark
+import skorch
+
+def perform_finetuning(benchmarks, metrics, args, foundation_model):
+    '''
+    Performs full finetuning on benchmarks using all data for training (no folds, no validation set)
+    Saves finetuned model, no metrics returned
+    '''
+    logger = get_logger()
+    results = {}
+    for benchmark in benchmarks:
+        # Load data stored in "./fine_tuning/data" folder - see example in data.py
+        b = load_benchmark(benchmark,  './', "NeuroRVQ")
+        X, sbj_id, y, ch_names = b.get_data()
+        n_outputs = len(np.unique(y))
+        n, c, t = X.shape
+        dataset = skorch.dataset.Dataset(X[:-1], y[:-1])
+        dummy_val = skorch.dataset.Dataset(np.array([X[0]]), np.array([y[0]]))
+
+        # Make model
+        model = get_model(
+            ch_names=ch_names,
+            n_times=t,
+            n_outputs=n_outputs,
+            args = args,
+            foundation_model = foundation_model,
+            train_head_only=args['train_head_only_finetuning']
+        )
+        print(f"No. Trainable Parameters: {model.size()}")
+
+        # Finetune model
+        model.fit(
+            dataset,
+            dummy_val,
+            batch_size=args['batch_size_finetuning'],
+            epochs=args['epoch_finetuning']
+        )
+
+        # Log training results (per epoch)
+        for m in metrics:
+            results = model.results[f'train_{m}']
+            for i in range(args['epoch_finetuning']):
+                logger.report_scalar(title="Fine-Tuning NeuroRVQ", series=f'train',
+                                     value=results[i], iteration=i)
+
+        # Save model
+        torch.save(model.state_dict(), './fine_tuned_model.pt')
+
+    return
+
+def load_neurorqv_fm(fine_tuning=False, verbose=False,
+                     model_path='./pretrained_models/foundation_models/NeuroRVQ_EEG_foundation_model_v1.pt'):
+    # Load experiment parameters from config file
+    config_stream = open("./flags/NeuroRVQ_EEG_v1.yml", 'r')
+    args = yaml.safe_load(config_stream)
+
+    # Fix the seeds for reproducibility
+    seed = 123
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # Get configuration params
+    args['n_global_electrodes'] = len(ch_names_global)
+    encoder_config, decoder_config = get_encoder_decoder_params(args)
+
+    # Load the foundation model
+    foundation_model = NeuroRVQFM(n_patches=args['n_patches'],
+                                    patch_size=args['patch_size'],
+                                    in_chans=args['in_chans_second_stage'],
+                                    out_chans=args['out_chans_second_stage'],
+                                    num_classes=0,
+                                    embed_dim=args['embed_dim_second_stage'],
+                                    depth=args['depth_second_stage'],
+                                    num_heads=args['num_heads_second_stage'],
+                                    mlp_ratio=args['mlp_ratio_second_stage'], qkv_bias=args['qkv_bias_second_stage'],
+                                    qk_norm=partial(nn.LayerNorm, eps=1e-6), drop_rate=args['drop_rate_second_stage'],
+                                    attn_drop_rate=args['attn_drop_rate_second_stage'],
+                                    drop_path_rate=args['drop_path_rate_second_stage'],
+                                    init_values=args['init_values_second_stage'],
+                                    init_scale=args['init_scale_second_stage'],
+                                    n_global_electrodes=args['n_global_electrodes'],
+                                    use_as_encoder = True, vocab_size=args['n_code'],
+                                    use_for_pretraining = False).to(device)
+
+    missing_keys, unexpected_keys = foundation_model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
+    total_params = sum(p.numel() for p in foundation_model.parameters())
+    print(f"Total parameters: {total_params}")
+
+    if (verbose):
+        print(f"Missing keys: {missing_keys},\nUnexpected keys: {unexpected_keys}")
+
+    if (fine_tuning):
+        # Select benchmark datasets
+        benchmarks = ["High Gamma"]
+        # Select evaluation metrics
+        # NOTE: metrics not included in this list will need to be implemented in the module for each model
+        metrics = ["accuracy", "bacc"]
+        perform_finetuning(benchmarks, metrics, args, foundation_model)

inference/run/NeuroRVQ_EEG_tokenizer_example.py

@@ -0,0 +1,50 @@
+import yaml
+import numpy as np
+import torch
+from NeuroRVQ.NeuroRVQ import NeuroRVQTokenizer
+from NeuroRVQ.NeuroRVQ_modules import get_encoder_decoder_params
+from preprocessing.preprocessing_eeg_example import preprocessing_eeg
+from plotting.plotting_example import process_and_plot
+from preprocessing.preprocessing_eeg_example import create_patches
+from inference.modules.NeuroRVQ_EEG_tokenizer_inference_modules import ch_names_global, create_embedding_ix, check_model_eval_mode
+
+def load_neurorqv_tokenizer(run_example=False, plot_results=False, verbose=False,
+                            model_path='./pretrained_models/tokenizers/NeuroRVQ_EEG_tokenizer_v1.pt'):
+    # Load experiment parameters from config file
+    config_stream = open("./flags/NeuroRVQ_EEG_v1.yml", 'r')
+    args = yaml.safe_load(config_stream)
+
+    # Fix the seeds for reproducibility
+    seed = 123
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Get configuration params for tokenizer
+    args['n_global_electrodes'] = len(ch_names_global)
+    encoder_config, decoder_config = get_encoder_decoder_params(args)
+
+    # Load the tokenizer
+    tokenizer = NeuroRVQTokenizer(encoder_config, decoder_config, n_code=args['n_code'],
+                                  code_dim=args['code_dim'], decoder_out_dim=args['decoder_out_dim']).to(device)
+
+    tokenizer.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+    tokenizer.eval()
+
+    if (verbose):
+        check_model_eval_mode(tokenizer)
+
+    if (run_example):
+        x, ch_names = preprocessing_eeg('example_files/eeg_sample/example_eeg_file.xdf')
+        ch_mask = np.isin(ch_names, ch_names_global)
+        ch_names = ch_names[ch_mask]
+        x = x[:, ch_mask, :]
+
+        x, n_time = create_patches(x, maximum_patches=args['n_patches'], patch_size=args['patch_size'], channels_use=ch_mask)
+        x = torch.from_numpy(x).float().to(device)
+
+        temporal_embedding_ix, spatial_embedding_ix = create_embedding_ix(n_time, args['n_patches'], ch_names, ch_names_global)
+        oringal_signal_std, reconstructed_signal_std = tokenizer(x, temporal_embedding_ix.int().to(device), spatial_embedding_ix.int().to(device))
+
+        if (plot_results):
+            process_and_plot(oringal_signal_std, reconstructed_signal_std, fs=args['patch_size'])

plotting/plotting_example.py

@@ -0,0 +1,84 @@
+import numpy as np
+from scipy.signal import butter, filtfilt
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import os
+
+# Define EEG bands
+bands = {
+    "Delta (0.5–4 Hz)": (0.5, 4),
+    "Theta (4–8 Hz)": (4, 8),
+    "Alpha (8–13 Hz)": (8, 13),
+    "Beta (13–30 Hz)": (13, 30),
+    "Gamma (30–45 Hz)": (30, 45),
+}
+
+# Bandpass filter function
+def bandpass_filter(data, lowcut, highcut, fs, order=2):
+    nyq = 0.5 * fs
+    low, high = lowcut / nyq, highcut / nyq
+    b, a = butter(order, [low, high], btype='band')
+    return filtfilt(b, a, data)
+
+def plot_reconstructions(originals_list, reconstructions_list, fs,
+                         labels=["NeuroRVQ"], save_dir="./figures"):
+
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    N, T = originals_list[0].shape
+    time = np.linspace(0, T / fs, T)
+
+    for i in tqdm(range(N), desc="Samples"):
+        plt.figure(figsize=(10, 12))
+
+        # Plot raw signals
+        plt.subplot(6, 1, 1)
+        orig = originals_list[0][i]
+        recon = reconstructions_list[0][i]
+        label = labels[0]
+
+        plt.plot(time, orig, label=f"Original Signal", alpha=0.7)
+        plt.plot(time, recon, linestyle='--', label=f"{label} Reconstruction", alpha=0.7)
+
+        plt.title(f"Raw Signal")
+        plt.legend()
+        plt.ylabel("Amplitude")
+
+        # Plot filtered bands
+        for j, (band_name, (low, high)) in enumerate(bands.items()):
+            plt.subplot(6, 1, j + 2)
+            orig = originals_list[0][i]
+            recon = reconstructions_list[0][i]
+            label = labels[0]
+
+            orig_band = bandpass_filter(orig, low, high, fs)
+            recon_band = bandpass_filter(recon, low, high, fs)
+
+            plt.plot(time, orig_band, label=f"{label} Original Signal", alpha=0.7)
+            plt.plot(time, recon_band, linestyle='--', label=f"{label} Reconstruction", alpha=0.7)
+
+            plt.title(f"{band_name} Band")
+            plt.ylabel("Amplitude")
+
+        plt.xlabel("Time (s)")
+        plt.tight_layout()
+
+        plt.savefig(f"{save_dir}/sample_{i}.png")
+        plt.close()
+
+
+def process_and_plot(originals, reconstructions, fs):
+    P, T = reconstructions[0].shape
+
+    originals_np = [
+        original.detach().cpu().numpy().reshape(P, T)
+        for original in originals
+    ]
+    reconstructions_np = [
+        reconstruction.detach().cpu().numpy().reshape(P, T)
+        for reconstruction in reconstructions
+    ]
+
+    # Plot
+    plot_reconstructions(originals_np, reconstructions_np, fs)

preprocessing/preprocessing_eeg_example.py

@@ -0,0 +1,55 @@
+import pyxdf
+from scipy import signal
+import numpy as np
+
+'''
+Example of how to preprocess an EEG recording
+Notch filter: 50., 60., 100.Hz
+High Pass: 45Hz
+Low Pass: 0.5Hz
+Resample at: 200Hz
+'''
+def preprocessing_eeg(data_path):
+    notch = [50., 60., 100.]
+    target_fs = 200
+    highpass = 0.5
+    lowpass = 45
+    clip = 500
+    streams, header = pyxdf.load_xdf(data_path, verbose=False,
+                                     synchronize_clocks=True, dejitter_timestamps=True,
+                                     select_streams=[{'name': 'Quick-32r_R2_EEG'}])
+    # Get sample rate
+    fs = float(streams[0]['info']['nominal_srate'][0])
+    channel_information = streams[0]["info"]["desc"][0]["channels"][0]["channel"]
+    ch_names = [x["label"][0] for x in channel_information][:29]
+
+    # Get EEG
+    x = streams[0]["time_series"][:, :29].T.astype(np.float64)  # (channels, time) comes as float32
+
+    # Filter and clip
+    for f_notch in notch:
+        if fs / 2 > f_notch:
+            [b_notch, a_notch] = signal.iirnotch(w0=f_notch, Q=f_notch / 2, fs=fs)
+            x = signal.filtfilt(b_notch, a_notch, x, axis=-1)
+    lowpass_applied = min(lowpass, fs / 2) - 0.5
+    [b, a] = signal.butter(N=3, Wn=[highpass, lowpass_applied], btype='bandpass', fs=fs)
+    x = signal.filtfilt(b, a, x, axis=-1)
+    x = x.clip(min=-clip, max=clip)
+    # Resampling
+    if target_fs != fs:
+        x = signal.resample(x, num=int(x.shape[-1] / fs * target_fs), axis=-1)
+    # Convert to float16 only after filtering
+    x = x.astype('float16')
+    x = x.reshape(1, x.shape[0], x.shape[1])
+    ch_names = np.array([c.lower().encode() for c in ch_names])
+    return x, ch_names
+
+'''
+Function to create patches for NeuroRVQ
+'''
+def create_patches(eeg_signal, maximum_patches, patch_size, channels_use):
+    n, c, t = eeg_signal.shape  # Batch / trials, channels, time
+    n_time = (maximum_patches // len(channels_use))
+    eeg_signal = eeg_signal[:, :, :n_time * patch_size]
+    eeg_signal_patches = eeg_signal[:, channels_use, :]
+    return eeg_signal_patches, n_time

pretrained_models/foundation_models/NeuroRVQ_EEG_foundation_model_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d14ef4eae2d4594d5f1f709226a3bb5f362afd84b68050d5da0fa0586aea886a
+size 234352952

pretrained_models/tokenizers/NeuroRVQ_EEG_tokenizer_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3e850e160b1a558c529466b1c0b8cebe8888a5e02824979886dd226187b2725
+size 304130106

requirements.txt

@@ -0,0 +1,189 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+# created-by: conda 25.3.1
+_libgcc_mutex=0.1=main
+_openmp_mutex=5.1=1_gnu
+attrs=25.3.0=pypi_0
+blas=1.0=mkl
+bottleneck=1.4.2=py311hf4808d0_0
+brotli-python=1.0.9=py311h6a678d5_9
+bzip2=1.0.8=h5eee18b_6
+c-ares=1.19.1=h5eee18b_0
+ca-certificates=2025.2.25=h06a4308_0
+certifi=2025.4.26=py311h06a4308_0
+charset-normalizer=3.3.2=pyhd3eb1b0_0
+clearml=1.18.0=pypi_0
+contourpy=1.3.1=py311hdb19cb5_0
+cuda-cudart=11.8.89=0
+cuda-cupti=11.8.87=0
+cuda-libraries=11.8.0=0
+cuda-nvrtc=11.8.89=0
+cuda-nvtx=11.8.86=0
+cuda-runtime=11.8.0=0
+cuda-version=12.9=3
+cycler=0.11.0=pyhd3eb1b0_0
+cyrus-sasl=2.1.28=h52b45da_1
+decorator=5.2.1=pypi_0
+einops=0.8.1=pypi_0
+expat=2.7.1=h6a678d5_0
+ffmpeg=4.3=hf484d3e_0
+filelock=3.17.0=py311h06a4308_0
+fontconfig=2.14.1=h55d465d_3
+fonttools=4.55.3=py311h5eee18b_0
+freetype=2.13.3=h4a9f257_0
+furl=2.1.4=pypi_0
+giflib=5.2.2=h5eee18b_0
+gmp=6.3.0=h6a678d5_0
+gmpy2=2.2.1=py311h5eee18b_0
+gnutls=3.6.15=he1e5248_0
+h5py=3.12.1=py311h5842655_1
+hdf5=1.14.5=h2b7332f_2
+icu=73.1=h6a678d5_0
+idna=3.7=py311h06a4308_0
+intel-openmp=2023.1.0=hdb19cb5_46306
+jinja2=3.1.6=py311h06a4308_0
+joblib=1.4.2=py311h06a4308_0
+jpeg=9e=h5eee18b_3
+jsonschema=4.23.0=pypi_0
+jsonschema-specifications=2025.4.1=pypi_0
+kiwisolver=1.4.8=py311h6a678d5_0
+krb5=1.20.1=h143b758_1
+lame=3.100=h7b6447c_0
+lazy-loader=0.4=pypi_0
+lcms2=2.16=h92b89f2_1
+ld_impl_linux-64=2.40=h12ee557_0
+lerc=4.0.0=h6a678d5_0
+libabseil=20250127.0=cxx17_h6a678d5_0
+libcublas=11.11.3.6=0
+libcufft=10.9.0.58=0
+libcufile=1.14.0.30=4
+libcups=2.4.2=h2d74bed_1
+libcurand=10.3.10.19=0
+libcurl=8.12.1=hc9e6f67_0
+libcusolver=11.4.1.48=0
+libcusparse=11.7.5.86=0
+libdeflate=1.22=h5eee18b_0
+libedit=3.1.20230828=h5eee18b_0
+libev=4.33=h7f8727e_1
+libffi=3.4.4=h6a678d5_1
+libgcc-ng=11.2.0=h1234567_1
+libgfortran-ng=11.2.0=h00389a5_1
+libgfortran5=11.2.0=h1234567_1
+libglib=2.78.4=hdc74915_0
+libgomp=11.2.0=h1234567_1
+libiconv=1.16=h5eee18b_3
+libidn2=2.3.4=h5eee18b_0
+libjpeg-turbo=2.0.0=h9bf148f_0
+libnghttp2=1.57.0=h2d74bed_0
+libnpp=11.8.0.86=0
+libnvjpeg=11.9.0.86=0
+libpng=1.6.39=h5eee18b_0
+libpq=17.4=hdbd6064_0
+libprotobuf=5.29.3=hc99497a_0
+libssh2=1.11.1=h251f7ec_0
+libstdcxx-ng=11.2.0=h1234567_1
+libtasn1=4.19.0=h5eee18b_0
+libtiff=4.7.0=hde9077f_0
+libunistring=0.9.10=h27cfd23_0
+libuuid=1.41.5=h5eee18b_0
+libwebp=1.3.2=h9f374a3_1
+libwebp-base=1.3.2=h5eee18b_1
+libxcb=1.17.0=h9b100fa_0
+libxkbcommon=1.9.1=h69220b7_0
+libxml2=2.13.8=hfdd30dd_0
+llvm-openmp=14.0.6=h9e868ea_0
+lz4-c=1.9.4=h6a678d5_1
+markupsafe=3.0.2=py311h5eee18b_0
+matplotlib=3.10.0=py311h06a4308_0
+matplotlib-base=3.10.0=py311hbfdbfaf_0
+mkl=2023.1.0=h213fc3f_46344
+mkl-service=2.4.0=py311h5eee18b_2
+mkl_fft=1.3.11=py311h5eee18b_0
+mkl_random=1.2.8=py311ha02d727_0
+mne=1.9.0=pypi_0
+mpc=1.3.1=h5eee18b_0
+mpfr=4.2.1=h5eee18b_0
+mpmath=1.3.0=py311h06a4308_0
+mysql=8.4.0=h721767e_2
+ncurses=6.4=h6a678d5_0
+nettle=3.7.3=hbbd107a_1
+networkx=3.4.2=py311h06a4308_0
+numexpr=2.10.1=py311h3c60e43_0
+numpy=2.0.1=py311h08b1b3b_1
+numpy-base=2.0.1=py311hf175353_1
+openh264=2.1.1=h4ff587b_0
+openjpeg=2.5.2=h0d4d230_1
+openldap=2.6.4=h42fbc30_0
+openssl=3.0.16=h5eee18b_0
+orderedmultidict=1.0.1=pypi_0
+packaging=24.2=py311h06a4308_0
+pandas=2.2.3=py311h6a678d5_0
+pathlib2=2.3.7.post1=pypi_0
+pcre2=10.42=hebb0a14_1
+pillow=11.1.0=py311hac6e08b_1
+pip=25.1=pyhc872135_2
+platformdirs=4.3.8=pypi_0
+pooch=1.8.2=pypi_0
+psutil=5.9.0=py311h5eee18b_1
+pthread-stubs=0.3=h0ce48e5_1
+pyjwt=2.9.0=pypi_0
+pyparsing=3.2.0=py311h06a4308_0
+pyqt=6.7.1=py311h6a678d5_1
+pyqt6-sip=13.9.1=py311h5eee18b_1
+pysocks=1.7.1=py311h06a4308_0
+python=3.11.11=he870216_0
+python-dateutil=2.9.0post0=py311h06a4308_2
+python-tzdata=2025.2=pyhd3eb1b0_0
+pytorch=2.5.1=py3.11_cuda11.8_cudnn9.1.0_0
+pytorch-cuda=11.8=h7e8668a_6
+pytorch-mutex=1.0=cuda
+pytz=2024.1=py311h06a4308_0
+pyyaml=6.0.2=py311h5eee18b_0
+qtbase=6.7.3=hdaa5aa8_0
+qtdeclarative=6.7.3=h6a678d5_0
+qtsvg=6.7.3=he621ea3_0
+qttools=6.7.3=h80c7b02_0
+qtwebchannel=6.7.3=h6a678d5_0
+qtwebsockets=6.7.3=h6a678d5_0
+readline=8.2=h5eee18b_0
+referencing=0.36.2=pypi_0
+requests=2.32.3=py311h06a4308_1
+rpds-py=0.25.0=pypi_0
+scikit-learn=1.6.1=py311h6a678d5_0
+scipy=1.15.3=py311h525edd1_0
+setuptools=78.1.1=py311h06a4308_0
+sip=6.10.0=py311h6a678d5_0
+six=1.17.0=py311h06a4308_0
+sqlite=3.45.3=h5eee18b_0
+sympy=1.13.3=py311h06a4308_1
+tbb=2021.8.0=hdb19cb5_0
+threadpoolctl=3.5.0=py311h92b7b1e_0
+tk=8.6.14=h39e8969_0
+torchaudio=2.5.1=py311_cu118
+torchtriton=3.1.0=py311
+torchvision=0.20.1=py311_cu118
+tornado=6.4.2=py311h5eee18b_0
+tqdm=4.67.1=py311h92b7b1e_0
+typing_extensions=4.12.2=py311h06a4308_0
+tzdata=2025b=h04d1e81_0
+unicodedata2=15.1.0=py311h5eee18b_1
+urllib3=2.3.0=py311h06a4308_0
+wheel=0.45.1=py311h06a4308_0
+xcb-util=0.4.1=h5eee18b_2
+xcb-util-cursor=0.1.5=h5eee18b_0
+xcb-util-image=0.4.0=h5eee18b_2
+xcb-util-renderutil=0.3.10=h5eee18b_0
+xkeyboard-config=2.44=h5eee18b_0
+xorg-libx11=1.8.12=h9b100fa_1
+xorg-libxau=1.0.12=h9b100fa_0
+xorg-libxdmcp=1.1.5=h9b100fa_0
+xorg-xorgproto=2024.1=h5eee18b_1
+xz=5.6.4=h5eee18b_1
+yaml=0.2.5=h7b6447c_0
+zlib=1.2.13=h5eee18b_1
+zstd=1.5.6=hc292b87_0
+numpy
+pandas
+skorch==0.15.0
+braindecode==0.8.1