EEG-DINO

.gitattributes

@@ -33,3 +33,4 @@ 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
+assets/eeg-dino.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,34 @@
----
-license: apache-2.0
----
+<div align="center">
+<h1>EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation</h1>
+
+<div align="center">
+    <img src="assets/eeg-dino.png" alt="positions" width="500"/>
+</div>
+
+We propose EEG-DINO, a novel foundation model for EEG encoding based on a hierarchical self-distillation framework. By multi-view semantic alignment, the model is able to extract multi-level semantic features from EEG data, which captures a wide range of semantic information, increasing the robustness against noise and variances inherent in complex EEG signals.
+Moreover, acknowledging the unique heterogeneous spatial-temporal dependencies in EEG signals, we design a channel-aware sampling mechanism and a decoupled positional coding scheme. They independently address spatial and temporal dimensions, enabling the model to capture the intricate structural characteristics of EEG signals. We pre-train EEG-DINO on a large-scale EEG corpus spanning over 9000 hours, which consistently achieves state-of-the-art performance on multiple downstream tasks. These results demonstrate the great effectiveness of our self-distillation framework for EEG encoding.
+
+## Pre-trained Models
+
+| Model           | Params |
+|:----------------|-------:|
+| EEG-DINO-Small |    4.6M|
+| EEG-DINO-Medium |   33M |
+| EEG-DINO-Large |   201M |
+
+### Usage
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python /path/to/run_finetuning.py
+```
+The default settings are for EEG-DINO-Small, if you want to use medium or large, you could change the embedding model in /path/to/models/eeg_encoder.py:
+```python
+from models.embedding_small import PatchEmbedding
+```
+and change the default settings in /path/to/run_finetuning.py:
+```python
+    parser.add_argument('--feature_size', default=200, type=int)
+    parser.add_argument('--num_layers', default=12, type=int)
+    parser.add_argument('--dim_feedforward', default=512, type=int)
+```
+512/16/1024 for medium and 1024/24/2048 for large.

engine_finetuning.py

@@ -0,0 +1,232 @@
+# --------------------------------------------------------
+# EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation
+# Based on BEiT-v2, timm, DeiT, DINO v2, LaBraM and CBraMod code bases
+# 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/dinov2
+# https://github.com/935963004/LaBraM
+# https://github.com/wjq-learning/CBraMod
+# ---------------------------------------------------------
+import math
+import sys
+from typing import Iterable, Optional
+import torch
+from timm.utils import ModelEma
+import utils
+from einops import rearrange
+import os
+import numpy as np
+import pandas as pd
+from sklearn.metrics import confusion_matrix
+
+def train_class_batch(model, samples, target, criterion):
+    outputs = model(samples)
+    loss = criterion(outputs, target)
+    return loss, outputs
+
+
+def get_loss_scale_for_deepspeed(model):
+    optimizer = model.optimizer
+    return optimizer.loss_scale if hasattr(optimizer, "loss_scale") else optimizer.cur_scale
+
+
+def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
+                    model_ema: Optional[ModelEma] = None, log_writer=None,
+                    start_steps=None, lr_schedule_values=None, wd_schedule_values=None,
+                    num_training_steps_per_epoch=None, update_freq=None, is_binary=True):
+    model.train(True)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    metric_logger.add_meter('min_lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 10
+
+    if loss_scaler is None:
+        model.zero_grad()
+        model.micro_steps = 0
+    else:
+        optimizer.zero_grad()
+
+    for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+        step = data_iter_step // update_freq
+        if step >= num_training_steps_per_epoch:
+            continue
+        it = start_steps + step  # global training iteration
+        # Update LR & WD for the first acc
+        if lr_schedule_values is not None or wd_schedule_values is not None and data_iter_step % update_freq == 0:
+            for i, param_group in enumerate(optimizer.param_groups):
+                if lr_schedule_values is not None:
+                    param_group["lr"] = lr_schedule_values[it] * param_group.get("lr_scale", 1.0)
+                if wd_schedule_values is not None and param_group["weight_decay"] > 0:
+                    param_group["weight_decay"] = wd_schedule_values[it]
+
+        # print("before", samples.shape)
+        samples = samples.float().to(device, non_blocking=True) / 100
+        samples = rearrange(samples, 'B N (A T) -> B N A T', T=200)
+        # print("after rearrange", samples.shape)
+        
+        targets = targets.to(device, non_blocking=True)
+        if is_binary:
+            targets = targets.float().unsqueeze(-1)
+
+        if loss_scaler is None:
+            samples = samples.half()
+            loss, output = train_class_batch(
+                model, samples, targets, criterion)
+        else:
+            with torch.amp.autocast(device_type='cuda'):
+                loss, output = train_class_batch(
+                    model, samples, targets, criterion)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        if loss_scaler is None:
+            loss /= update_freq
+            model.backward(loss)
+            model.step()
+
+            if (data_iter_step + 1) % update_freq == 0:
+                # model.zero_grad()
+                # Deepspeed will call step() & model.zero_grad() automatic
+                if model_ema is not None:
+                    model_ema.update(model)
+            grad_norm = None
+            loss_scale_value = get_loss_scale_for_deepspeed(model)
+        else:
+            # this attribute is added by timm on one optimizer (adahessian)
+            is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
+            loss /= update_freq
+            grad_norm = loss_scaler(loss, optimizer, clip_grad=max_norm,
+                                    parameters=model.parameters(), create_graph=is_second_order,
+                                    update_grad=(data_iter_step + 1) % update_freq == 0)
+            if (data_iter_step + 1) % update_freq == 0:
+                optimizer.zero_grad()
+                if model_ema is not None:
+                    model_ema.update(model)
+            loss_scale_value = loss_scaler.state_dict()["scale"]
+
+        torch.cuda.synchronize()
+
+        if is_binary:
+            class_acc = utils.get_metrics(torch.sigmoid(output).detach().cpu().numpy(), targets.detach().cpu().numpy(), ["accuracy"], is_binary)["accuracy"]
+        else:
+            class_acc = (output.max(-1)[-1] == targets.squeeze()).float().mean()
+            
+        metric_logger.update(loss=loss_value)
+        metric_logger.update(class_acc=class_acc)
+        metric_logger.update(loss_scale=loss_scale_value)
+        min_lr = 10.
+        max_lr = 0.
+        for group in optimizer.param_groups:
+            min_lr = min(min_lr, group["lr"])
+            max_lr = max(max_lr, group["lr"])
+
+        metric_logger.update(lr=max_lr)
+        metric_logger.update(min_lr=min_lr)
+        weight_decay_value = None
+        for group in optimizer.param_groups:
+            if group["weight_decay"] > 0:
+                weight_decay_value = group["weight_decay"]
+        metric_logger.update(weight_decay=weight_decay_value)
+        metric_logger.update(grad_norm=grad_norm)
+
+        if log_writer is not None:
+            log_writer.update(loss=loss_value, head="loss")
+            log_writer.update(class_acc=class_acc, head="loss")
+            log_writer.update(loss_scale=loss_scale_value, head="opt")
+            log_writer.update(lr=max_lr, head="opt")
+            log_writer.update(min_lr=min_lr, head="opt")
+            log_writer.update(weight_decay=weight_decay_value, head="opt")
+            log_writer.update(grad_norm=grad_norm, head="opt")
+
+            log_writer.set_step()
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def evaluate(data_loader, model, device, output_dir=None, header='Test:', metrics=['acc'], is_binary=True, epoch=None):
+    if is_binary:
+        criterion = torch.nn.BCEWithLogitsLoss()
+    else:
+        criterion = torch.nn.CrossEntropyLoss()
+
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    
+    # 新增：初始化存储预测和真实值的列表
+    all_outputs = []
+    all_targets = []
+    
+    model.eval()
+    for step, batch in enumerate(metric_logger.log_every(data_loader, 10, header)):
+        EEG = batch[0]
+        target = batch[-1]
+        EEG = EEG.float().to(device, non_blocking=True) / 100
+        EEG = rearrange(EEG, 'B N (A T) -> B N A T', T=200)
+        target = target.to(device, non_blocking=True)
+        if is_binary:
+            target = target.float().unsqueeze(-1)
+        
+        # compute output
+        with torch.amp.autocast(device_type='cuda'):
+            output = model(EEG)
+            loss = criterion(output, target)
+        
+        if is_binary:
+            output = torch.sigmoid(output).cpu()
+        else:
+            output = output.cpu()
+        target = target.cpu()
+
+        results = utils.get_metrics(output.numpy(), target.numpy(), metrics, is_binary)
+        pred = output.numpy()
+        true = target.numpy()
+
+        # 新增：收集原始输出
+        all_outputs.append(pred)
+        all_targets.append(true)
+
+        batch_size = EEG.shape[0]
+        metric_logger.update(loss=loss.item())
+        for key, value in results.items():
+            metric_logger.meters[key].update(value, n=batch_size)
+        #metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print('* loss {losses.global_avg:.3f}'
+          .format(losses=metric_logger.loss))
+    
+    # 新增：计算混淆矩阵
+    all_outputs = np.concatenate(all_outputs)
+    all_targets = np.concatenate(all_targets)
+    
+    if is_binary:
+        y_pred = (all_outputs > 0.5).astype(int)
+    else:
+        y_pred = np.argmax(all_outputs, axis=1)
+    y_true = all_targets.squeeze().astype(int)
+    
+    cm = confusion_matrix(y_true, y_pred)
+    ret = utils.get_metrics(all_outputs, all_targets, metrics, is_binary, 0.5)
+    ret['loss'] = metric_logger.loss.global_avg
+    ret['confusion_matrix'] = cm.tolist()  # 转换为列表方便保存
+
+    # 新增：保存预测结果和混淆矩阵
+    if output_dir and epoch is not None:
+        os.makedirs(output_dir, exist_ok=True)
+        # 保存分类头原始输出
+        np.save(os.path.join(output_dir, f'epoch{epoch}_predictions.npy'), all_outputs)
+        # 保存混淆矩阵
+        pd.DataFrame(cm).to_csv(os.path.join(output_dir, f'epoch{epoch}_confusion_matrix.csv'))
+        
+    return ret

models/eeg_encoder.py

@@ -0,0 +1,56 @@
+# --------------------------------------------------------
+# EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation
+# Based on BEiT-v2, timm, DeiT, DINO v2, LaBraM and CBraMod code bases
+# 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/dinov2
+# https://github.com/935963004/LaBraM
+# https://github.com/wjq-learning/CBraMod
+# ---------------------------------------------------------
+import torch
+import torch.nn as nn
+from monai.networks.nets.swin_unetr import *
+import torch.nn.functional as F
+from models.embedding_small import PatchEmbedding
+from models.transformer import TransformerEncoderLayer
+
+class EEGEncoder(nn.Module):
+    def __init__(self, args):
+        super(EEGEncoder, self).__init__()
+        self.patch_embedding = PatchEmbedding(
+            d_model=args.feature_size
+        )
+        
+        self.encoder_layers = nn.ModuleList([
+            TransformerEncoderLayer(
+                d_model=args.feature_size, 
+                nhead=args.num_heads,
+                dim_feedforward=args.dim_feedforward,
+            ) for _ in range(args.num_layers)
+        ])
+        
+        self.global_tokens = nn.Parameter(
+            torch.randn(1, args.num_global_tokens, args.feature_size)
+        )
+        self.global_token_layer = args.global_token_layer
+        
+    def forward(self, x_in):
+        B, C, P, L = x_in.shape
+        if hasattr(self.patch_embedding, 'in_dim'):
+            self.patch_embedding.in_dim = C
+        
+        # 1. Patch Embedding
+        x = self.patch_embedding(x_in)  # [B, C, P, D]
+        b = x.shape[0]
+        
+        x = x.reshape(b, -1, x.shape[-1])  # [B, C*P, D]
+
+        global_tokens = self.global_tokens.expand(b, -1, -1)  # [B, num_global, D]
+        
+        for i, encoder_layer in enumerate(self.encoder_layers):
+            x = encoder_layer(x)
+            if i + 1 == self.global_token_layer:
+                x = torch.cat([global_tokens, x], dim=1)  # [B, num_global+C*P, D]
+        
+        return x

models/embedding_large.py

@@ -0,0 +1,80 @@
+# --------------------------------------------------------
+# EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation
+# Based on BEiT-v2, timm, DeiT, DINO v2, LaBraM and CBraMod code bases
+# 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/dinov2
+# https://github.com/935963004/LaBraM
+# https://github.com/wjq-learning/CBraMod
+# ---------------------------------------------------------
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.d_model = d_model
+        self.time_encoding = nn.Sequential(
+            nn.Conv2d(in_channels=d_model, out_channels=d_model, kernel_size=(1, 5), stride=(1, 1), padding=(0, 2),
+                      groups=d_model),
+        )
+
+        self.proj_in = nn.Sequential(
+            nn.Conv2d(in_channels=1, out_channels=128, kernel_size=(1, 49), stride=(1, 25), padding=(0, 24)),
+            nn.GroupNorm(16, 128),
+            nn.GELU(),
+
+            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1)),
+            nn.GroupNorm(16, 256),
+            nn.GELU(),
+
+            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1)),
+            nn.GroupNorm(16, 128),
+            nn.GELU(),
+        )
+        self.spectral_proj = nn.Sequential(
+            nn.Linear(101, d_model),
+            nn.Dropout(0.1),
+        )
+
+        self.num_channels = 19
+        self.channel_embedding = nn.Linear(self.num_channels, d_model)
+
+    def forward(self, x):
+        bz, ch_num, patch_num, patch_size = x.shape
+        channel_in = torch.arange(self.num_channels+1).cuda()
+
+        x = x.contiguous().view(bz, 1, ch_num * patch_num, patch_size)
+        patch_emb = self.proj_in(x)
+        patch_emb = patch_emb.permute(0, 2, 1, 3).contiguous().view(bz, ch_num, patch_num, self.d_model)
+
+        x = x.contiguous().view(bz*ch_num*patch_num, patch_size)
+        spectral = torch.fft.rfft(x, dim=-1, norm='forward')
+        spectral = torch.abs(spectral).contiguous().view(bz, ch_num, patch_num, 101)
+        spectral_emb = self.spectral_proj(spectral)
+
+        patch_emb = patch_emb + spectral_emb
+
+        channel_embeddings = []
+        start_idx = 0
+
+        group_channels = channel_in[start_idx:start_idx + ch_num]
+        group_one_hot = F.one_hot(group_channels, num_classes=self.num_channels).float()
+        group_emb = self.channel_embedding(group_one_hot)
+        group_emb = group_emb.unsqueeze(0).unsqueeze(2)  # [1, ch_num, 1, d_model]
+        group_emb = group_emb.expand(bz, -1, patch_num, -1)
+        channel_embeddings.append(group_emb)
+        start_idx += ch_num
+        
+        channel_pos = torch.cat(channel_embeddings, dim=0)  # [total_bz, ch_num, patch_num, d_model]
+
+        patch_emb = patch_emb + channel_pos
+
+        time_embedding = self.time_encoding(patch_emb.permute(0, 3, 1, 2))
+        time_embedding = time_embedding.permute(0, 2, 3, 1)
+
+        patch_emb = patch_emb + time_embedding
+
+        return patch_emb

models/embedding_medium.py

@@ -0,0 +1,80 @@
+# --------------------------------------------------------
+# EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation
+# Based on BEiT-v2, timm, DeiT, DINO v2, LaBraM and CBraMod code bases
+# 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/dinov2
+# https://github.com/935963004/LaBraM
+# https://github.com/wjq-learning/CBraMod
+# ---------------------------------------------------------
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.d_model = d_model
+        self.time_encoding = nn.Sequential(
+            nn.Conv2d(in_channels=d_model, out_channels=d_model, kernel_size=(1, 5), stride=(1, 1), padding=(0, 2),
+                      groups=d_model),
+        )
+
+        self.proj_in = nn.Sequential(
+            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(1, 49), stride=(1, 25), padding=(0, 24)),
+            nn.GroupNorm(8, 64),
+            nn.GELU(),
+
+            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1)),
+            nn.GroupNorm(8, 128),
+            nn.GELU(),
+
+            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1)),
+            nn.GroupNorm(8, 64),
+            nn.GELU(),
+        )
+        self.spectral_proj = nn.Sequential(
+            nn.Linear(101, d_model),
+            nn.Dropout(0.1),
+        )
+
+        self.num_channels = 19
+        self.channel_embedding = nn.Linear(self.num_channels, d_model)
+
+    def forward(self, x):
+        bz, ch_num, patch_num, patch_size = x.shape
+        channel_in = torch.arange(self.num_channels+1).cuda()
+
+        x = x.contiguous().view(bz, 1, ch_num * patch_num, patch_size)
+        patch_emb = self.proj_in(x)
+        patch_emb = patch_emb.permute(0, 2, 1, 3).contiguous().view(bz, ch_num, patch_num, self.d_model)
+
+        x = x.contiguous().view(bz*ch_num*patch_num, patch_size)
+        spectral = torch.fft.rfft(x, dim=-1, norm='forward')
+        spectral = torch.abs(spectral).contiguous().view(bz, ch_num, patch_num, 101)
+        spectral_emb = self.spectral_proj(spectral)
+
+        patch_emb = patch_emb + spectral_emb
+
+        channel_embeddings = []
+        start_idx = 0
+
+        group_channels = channel_in[start_idx:start_idx + ch_num]
+        group_one_hot = F.one_hot(group_channels, num_classes=self.num_channels).float()
+        group_emb = self.channel_embedding(group_one_hot)
+        group_emb = group_emb.unsqueeze(0).unsqueeze(2)  # [1, ch_num, 1, d_model]
+        group_emb = group_emb.expand(bz, -1, patch_num, -1)
+        channel_embeddings.append(group_emb)
+        start_idx += ch_num
+        
+        channel_pos = torch.cat(channel_embeddings, dim=0)  # [total_bz, ch_num, patch_num, d_model]
+
+        patch_emb = patch_emb + channel_pos
+
+        time_embedding = self.time_encoding(patch_emb.permute(0, 3, 1, 2))
+        time_embedding = time_embedding.permute(0, 2, 3, 1)
+
+        patch_emb = patch_emb + time_embedding
+
+        return patch_emb

models/embedding_small.py

@@ -0,0 +1,80 @@
+# --------------------------------------------------------
+# EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation
+# Based on BEiT-v2, timm, DeiT, DINO v2, LaBraM and CBraMod code bases
+# 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/dinov2
+# https://github.com/935963004/LaBraM
+# https://github.com/wjq-learning/CBraMod
+# ---------------------------------------------------------
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.d_model = d_model
+        self.time_encoding = nn.Sequential(
+            nn.Conv2d(in_channels=d_model, out_channels=d_model, kernel_size=(1, 5), stride=(1, 1), padding=(0, 2),
+                      groups=d_model),
+        )
+
+        self.proj_in = nn.Sequential(
+            nn.Conv2d(in_channels=1, out_channels=25, kernel_size=(1, 49), stride=(1, 25), padding=(0, 24)),
+            nn.GroupNorm(5, 25),
+            nn.GELU(),
+
+            nn.Conv2d(in_channels=25, out_channels=25, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1)),
+            nn.GroupNorm(5, 25),
+            nn.GELU(),
+
+            nn.Conv2d(in_channels=25, out_channels=25, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1)),
+            nn.GroupNorm(5, 25),
+            nn.GELU(),
+        )
+        self.spectral_proj = nn.Sequential(
+            nn.Linear(101, d_model),
+            nn.Dropout(0.1),
+        )
+
+        self.num_channels = 19
+        self.channel_embedding = nn.Linear(self.num_channels, d_model)
+
+    def forward(self, x):
+        bz, ch_num, patch_num, patch_size = x.shape
+        channel_in = torch.arange(self.num_channels+1).cuda()
+
+        x = x.contiguous().view(bz, 1, ch_num * patch_num, patch_size)
+        patch_emb = self.proj_in(x)
+        patch_emb = patch_emb.permute(0, 2, 1, 3).contiguous().view(bz, ch_num, patch_num, self.d_model)
+
+        x = x.contiguous().view(bz*ch_num*patch_num, patch_size)
+        spectral = torch.fft.rfft(x, dim=-1, norm='forward')
+        spectral = torch.abs(spectral).contiguous().view(bz, ch_num, patch_num, 101)
+        spectral_emb = self.spectral_proj(spectral)
+
+        patch_emb = patch_emb + spectral_emb
+
+        channel_embeddings = []
+        start_idx = 0
+
+        group_channels = channel_in[start_idx:start_idx + ch_num]
+        group_one_hot = F.one_hot(group_channels, num_classes=self.num_channels).float()
+        group_emb = self.channel_embedding(group_one_hot)
+        group_emb = group_emb.unsqueeze(0).unsqueeze(2)  # [1, ch_num, 1, d_model]
+        group_emb = group_emb.expand(bz, -1, patch_num, -1)
+        channel_embeddings.append(group_emb)
+        start_idx += ch_num
+        
+        channel_pos = torch.cat(channel_embeddings, dim=0)  # [total_bz, ch_num, patch_num, d_model]
+
+        patch_emb = patch_emb + channel_pos
+
+        time_embedding = self.time_encoding(patch_emb.permute(0, 3, 1, 2))
+        time_embedding = time_embedding.permute(0, 2, 3, 1)
+
+        patch_emb = patch_emb + time_embedding
+
+        return patch_emb

models/transformer.py

@@ -0,0 +1,191 @@
+# --------------------------------------------------------
+# EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation
+# Based on BEiT-v2, timm, DeiT, DINO v2, LaBraM and CBraMod code bases
+# 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/dinov2
+# https://github.com/935963004/LaBraM
+# https://github.com/wjq-learning/CBraMod
+# ---------------------------------------------------------
+from typing import Union, Callable
+
+import torch
+import torch.nn as nn
+
+from torch import Tensor
+from torch.nn import functional as F
+from timm.models.layers import drop_path
+
+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 TransformerEncoderLayer(nn.Module):
+    __constants__ = ['norm_first']
+
+    def __init__(self, d_model: int, nhead: int, dim_feedforward: int = 2048, dropout: float = 0.1,
+                 activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+                 layer_norm_eps: float = 1e-5, batch_first: bool = False, norm_first: bool = False,
+                 bias: bool = True, device=None, dtype=None) -> None:
+        super().__init__()
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        
+        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+        self.attn = Attention(
+            dim=d_model,
+            num_heads=nhead,
+            qkv_bias=bias,
+            qk_norm=None,
+            qk_scale=None,
+            attn_drop=dropout,
+            proj_drop=dropout,
+            window_size=None,
+            attn_head_dim=None,
+            **factory_kwargs
+        )
+        
+        self.drop_path = DropPath(dropout) if dropout > 0. else nn.Identity()
+        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+        mlp_hidden_dim = dim_feedforward
+        self.mlp = Mlp(
+            in_features=d_model,
+            hidden_features=mlp_hidden_dim,
+            act_layer=nn.GELU,
+            drop=dropout
+        )
+
+        # 添加gamma参数支持
+        init_values = 0.0  # 可以通过参数传入
+        if init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((d_model)), requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((d_model)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x, rel_pos_bias=None, return_attention=False, return_qkv=False):
+        if return_attention:
+            return self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, return_attention=True)
+        if return_qkv:
+            y, qkv = self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, return_qkv=return_qkv)
+            x = x + self.drop_path(self.gamma_1 * y)
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+            return x, qkv
+
+        if self.gamma_1 is None:
+            x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_norm=None, qk_scale=None, attn_drop=0.,
+                 proj_drop=0., window_size=None, attn_head_dim=None, **kwargs):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or 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
+
+        if window_size:
+            self.window_size = window_size
+            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+            self.relative_position_bias_table = nn.Parameter(
+                torch.zeros(self.num_relative_distance, num_heads))
+            # 添加window_size相关的代码...
+        else:
+            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, rel_pos_bias=None, return_attention=False, return_qkv=False):
+        B, N, C = x.shape
+        qkv_bias = None
+        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))
+        
+        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)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        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))
+
+        if self.relative_position_bias_table is not None:
+            relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                self.window_size[0] * self.window_size[1] + 1,
+                self.window_size[0] * self.window_size[1] + 1, -1)
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        if rel_pos_bias is not None:
+            attn = attn + rel_pos_bias
+
+        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 Mlp(nn.Module):
+    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

optim_factory.py

@@ -0,0 +1,189 @@
+# --------------------------------------------------------
+# Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI
+# By Wei-Bang Jiang
+# Based on BEiT-v2, timm, DeiT, and DINO code bases
+# 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
+from torch import optim as optim
+
+from timm.optim.adafactor import Adafactor
+from timm.optim.adahessian import Adahessian
+from timm.optim.adamp import AdamP
+from timm.optim.lookahead import Lookahead
+# from timm.optim.nadam import Nadam
+from timm.optim.nvnovograd import NvNovoGrad
+# from timm.optim.radam import RAdam
+from timm.optim.rmsprop_tf import RMSpropTF
+from timm.optim.sgdp import SGDP
+
+import json
+
+try:
+    from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD
+    has_apex = True
+except ImportError:
+    has_apex = False
+
+
+def get_num_layer_for_vit(var_name, num_max_layer):
+    if var_name in ("cls_token", "mask_token", "pos_embed"):
+        return 0
+    elif var_name.startswith("patch_embed"):
+        return 0
+    elif var_name.startswith("rel_pos_bias"):
+        return num_max_layer - 1
+    elif var_name.startswith("blocks"):
+        layer_id = int(var_name.split('.')[1])
+        return layer_id + 1
+    else:
+        return num_max_layer - 1
+
+
+class LayerDecayValueAssigner(object):
+    def __init__(self, values):
+        self.values = values
+
+    def get_scale(self, layer_id):
+        return self.values[layer_id]
+
+    def get_layer_id(self, var_name):
+        return get_num_layer_for_vit(var_name, len(self.values))
+
+
+def get_parameter_groups(model, weight_decay=1e-5, skip_list=(), get_num_layer=None, get_layer_scale=None, **kwargs):
+    parameter_group_names = {}
+    parameter_group_vars = {}
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue  # frozen weights
+        if len(kwargs.get('filter_name', [])) > 0:
+            flag = False
+            for filter_n in kwargs.get('filter_name', []):
+                if filter_n in name:
+                    print(f"filter {name} because of the pattern {filter_n}")
+                    flag = True
+            if flag:
+                continue
+        if param.ndim <= 1 or name.endswith(".bias") or name in skip_list: # param.ndim <= 1 len(param.shape) == 1
+            group_name = "no_decay"
+            this_weight_decay = 0.
+        else:
+            group_name = "decay"
+            this_weight_decay = weight_decay
+        if get_num_layer is not None:
+            layer_id = get_num_layer(name)
+            group_name = "layer_%d_%s" % (layer_id, group_name)
+        else:
+            layer_id = None
+
+        if group_name not in parameter_group_names:
+            if get_layer_scale is not None:
+                scale = get_layer_scale(layer_id)
+            else:
+                scale = 1.
+
+            parameter_group_names[group_name] = {
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr_scale": scale
+            }
+            parameter_group_vars[group_name] = {
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr_scale": scale
+            }
+
+        parameter_group_vars[group_name]["params"].append(param)
+        parameter_group_names[group_name]["params"].append(name)
+    print("Param groups = %s" % json.dumps(parameter_group_names, indent=2))
+    return list(parameter_group_vars.values())
+
+
+def create_optimizer(args, model, get_num_layer=None, get_layer_scale=None, filter_bias_and_bn=True, skip_list=None, **kwargs):
+    opt_lower = args.opt.lower()
+    weight_decay = args.weight_decay
+    if weight_decay and filter_bias_and_bn:
+        skip = {}
+        if skip_list is not None:
+            skip = skip_list
+        elif hasattr(model, 'no_weight_decay'):
+            skip = model.no_weight_decay()
+        print(f"Skip weight decay name marked in model: {skip}")
+        parameters = get_parameter_groups(model, weight_decay, skip, get_num_layer, get_layer_scale, **kwargs)
+        weight_decay = 0.
+    else:
+        parameters = model.parameters()
+
+    if 'fused' in opt_lower:
+        assert has_apex and torch.cuda.is_available(), 'APEX and CUDA required for fused optimizers'
+
+    opt_args = dict(lr=args.lr, weight_decay=weight_decay)
+    if hasattr(args, 'opt_eps') and args.opt_eps is not None:
+        opt_args['eps'] = args.opt_eps
+    if hasattr(args, 'opt_betas') and args.opt_betas is not None:
+        opt_args['betas'] = args.opt_betas
+    
+    print('Optimizer config:', opt_args)
+    opt_split = opt_lower.split('_')
+    opt_lower = opt_split[-1]
+    if opt_lower == 'sgd' or opt_lower == 'nesterov':
+        opt_args.pop('eps', None)
+        optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+    elif opt_lower == 'momentum':
+        opt_args.pop('eps', None)
+        optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=False, **opt_args)
+    elif opt_lower == 'adam':
+        optimizer = optim.Adam(parameters, **opt_args)
+    elif opt_lower == 'adamw':
+        optimizer = optim.AdamW(parameters, **opt_args)
+    # elif opt_lower == 'nadam':
+    #     optimizer = Nadam(parameters, **opt_args)
+    # elif opt_lower == 'radam':
+    #     optimizer = RAdam(parameters, **opt_args)
+    elif opt_lower == 'adamp':
+        optimizer = AdamP(parameters, wd_ratio=0.01, nesterov=True, **opt_args)
+    elif opt_lower == 'sgdp':
+        optimizer = SGDP(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+    elif opt_lower == 'adadelta':
+        optimizer = optim.Adadelta(parameters, **opt_args)
+    elif opt_lower == 'adafactor':
+        if not args.lr:
+            opt_args['lr'] = None
+        optimizer = Adafactor(parameters, **opt_args)
+    elif opt_lower == 'adahessian':
+        optimizer = Adahessian(parameters, **opt_args)
+    elif opt_lower == 'rmsprop':
+        optimizer = optim.RMSprop(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
+    elif opt_lower == 'rmsproptf':
+        optimizer = RMSpropTF(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
+    elif opt_lower == 'nvnovograd':
+        optimizer = NvNovoGrad(parameters, **opt_args)
+    elif opt_lower == 'fusedsgd':
+        opt_args.pop('eps', None)
+        optimizer = FusedSGD(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+    elif opt_lower == 'fusedmomentum':
+        opt_args.pop('eps', None)
+        optimizer = FusedSGD(parameters, momentum=args.momentum, nesterov=False, **opt_args)
+    elif opt_lower == 'fusedadam':
+        optimizer = FusedAdam(parameters, adam_w_mode=False, **opt_args)
+    elif opt_lower == 'fusedadamw':
+        optimizer = FusedAdam(parameters, adam_w_mode=True, **opt_args)
+    elif opt_lower == 'fusedlamb':
+        optimizer = FusedLAMB(parameters, **opt_args)
+    elif opt_lower == 'fusednovograd':
+        opt_args.setdefault('betas', (0.95, 0.98))
+        optimizer = FusedNovoGrad(parameters, **opt_args)
+    else:
+        assert False and "Invalid optimizer"
+        raise ValueError
+
+    if len(opt_split) > 1:
+        if opt_split[0] == 'lookahead':
+            optimizer = Lookahead(optimizer)
+
+    return optimizer

run_finetuning.py

@@ -0,0 +1,565 @@
+# --------------------------------------------------------
+# EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation
+# Based on BEiT-v2, timm, DeiT, DINO v2, LaBraM and CBraMod code bases
+# 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/dinov2
+# https://github.com/935963004/LaBraM
+# https://github.com/wjq-learning/CBraMod
+# ---------------------------------------------------------
+import argparse
+import datetime
+import numpy as np
+import time
+import torch
+import torch.backends.cudnn as cudnn
+import json
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from pathlib import Path
+from collections import OrderedDict
+from timm.loss import LabelSmoothingCrossEntropy
+from optim_factory import create_optimizer, LayerDecayValueAssigner
+
+from engine_finetuning import train_one_epoch, evaluate
+from utils import NativeScalerWithGradNormCount as NativeScaler
+import utils
+import torch.nn as nn
+
+from models.eeg_encoder import EEGEncoder
+
+def get_args():
+    parser = argparse.ArgumentParser('EEG-DINO finetuning args', add_help=False)
+    parser.add_argument('--batch_size', default=512, type=int)
+    parser.add_argument('--epochs', default=50, type=int)
+    parser.add_argument('--update_freq', default=1, type=int)
+    parser.add_argument('--save_ckpt_freq', default=5, type=int)
+
+    parser.add_argument('--feature_size', default=200, type=int)
+    parser.add_argument('--num_global_tokens', default=1, type=int)
+    parser.add_argument('--num_heads', default=8, type=int)
+    parser.add_argument('--num_layers', default=12, type=int)
+    parser.add_argument('--dim_feedforward', default=512, type=int)
+    parser.add_argument('--global_token_layer', default=1, type=int)
+
+    parser.add_argument('--layer_scale_init_value', default=0.1, type=float, 
+                        help="0.1 for base, 1e-5 for large. set 0 to disable layer scale")
+
+    parser.add_argument('--input_size', default=200, type=int,
+                        help='EEG input size')
+
+    parser.add_argument('--drop', type=float, default=0.0, metavar='PCT',
+                        help='Dropout rate (default: 0.)')
+    parser.add_argument('--attn_drop_rate', type=float, default=0.0, metavar='PCT',
+                        help='Attention dropout rate (default: 0.)')
+    parser.add_argument('--drop_path', type=float, default=0.1, metavar='PCT',
+                        help='Drop path rate (default: 0.1)')
+
+    # Optimizer parameters
+    parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER',
+                        help='Optimizer (default: "adamw"')
+    parser.add_argument('--opt_eps', default=1e-8, type=float, metavar='EPSILON',
+                        help='Optimizer Epsilon (default: 1e-8)')
+    parser.add_argument('--opt_betas', default=None, type=float, nargs='+', metavar='BETA',
+                        help='Optimizer Betas (default: None, use opt default)')
+    parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM',
+                        help='Clip gradient norm (default: None, no clipping)')
+    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
+                        help='SGD momentum (default: 0.9)')
+    parser.add_argument('--weight_decay', type=float, default=0.05,
+                        help='weight decay (default: 0.05)')
+    parser.add_argument('--weight_decay_end', type=float, default=None, help="""Final value of the
+        weight decay. We use a cosine schedule for WD and using a larger decay by
+        the end of training improves performance for ViTs.""")
+
+    parser.add_argument('--lr', type=float, default=5e-4, metavar='LR',
+                        help='learning rate (default: 1e-4)')
+    parser.add_argument('--layer_decay', type=float, default=0.9)
+
+    parser.add_argument('--warmup_lr', type=float, default=1e-6, metavar='LR',
+                        help='warmup learning rate (default: 1e-6)')
+    parser.add_argument('--min_lr', type=float, default=1e-6, metavar='LR',
+                        help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')
+
+    parser.add_argument('--warmup_epochs', type=int, default=5, metavar='N',
+                        help='epochs to warmup LR, if scheduler supports')
+    parser.add_argument('--warmup_steps', type=int, default=-1, metavar='N',
+                        help='num of steps to warmup LR, will overload warmup_epochs if set > 0')
+
+    parser.add_argument('--smoothing', type=float, default=0.1,
+                        help='Label smoothing (default: 0.1)')
+
+    # * Random Erase params
+    parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
+                        help='Random erase prob (default: 0.25)')
+    parser.add_argument('--remode', type=str, default='pixel',
+                        help='Random erase mode (default: "pixel")')
+    parser.add_argument('--recount', type=int, default=1,
+                        help='Random erase count (default: 1)')
+    parser.add_argument('--resplit', action='store_true', default=False,
+                        help='Do not random erase first (clean) augmentation split')
+
+    # * Finetuning params
+    parser.add_argument('--finetune', default="/path/to/ckpt",
+                        help='finetune from checkpoint')
+    parser.add_argument('--model_prefix', default='', type=str)
+    parser.add_argument('--init_scale', default=0.001, type=float)
+    parser.add_argument('--disable_weight_decay_on_rel_pos_bias', action='store_true', default=False)
+    parser.add_argument('--freeze_all_except_head', action='store_true', default=False)
+
+    # Dataset parameters
+    parser.add_argument('--nb_classes', default=0, type=int,
+                        help='number of the classification types')
+    parser.add_argument('--output_dir', default="/path/to/output",
+                        help='path where to save, empty for no saving')
+    parser.add_argument('--log_dir', default="/path/to/log",
+                        help='path where to tensorboard log')
+    parser.add_argument('--device', default='cuda',
+                        help='device to use for training / testing')
+    parser.add_argument('--seed', default=0, type=int)
+    parser.add_argument('--resume', default='',
+                        help='resume from checkpoint')
+    parser.add_argument('--auto_resume', action='store_true')
+    parser.add_argument('--no_auto_resume', action='store_false', dest='auto_resume')
+    parser.set_defaults(auto_resume=True)
+
+    parser.add_argument('--save_ckpt', action='store_true')
+    parser.add_argument('--no_save_ckpt', action='store_false', dest='save_ckpt')
+    parser.set_defaults(save_ckpt=True)
+
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+                        help='start epoch')
+    parser.add_argument('--eval', action='store_true',
+                        help='Perform evaluation only')
+    parser.add_argument('--dist_eval', action='store_true', default=False,
+                        help='Enabling distributed evaluation')
+    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--pin_mem', action='store_true',
+                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
+    parser.set_defaults(pin_mem=True)
+
+    # distributed training parameters
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='number of distributed processes')
+    parser.add_argument('--local_rank', default=-1, type=int)
+    parser.add_argument('--dist_on_itp', action='store_true')
+    parser.add_argument('--dist_url', default='env://',
+                        help='url used to set up distributed training')
+
+    parser.add_argument('--dataset', default='TUAB', type=str,
+                        help='dataset: TUAB | TUEV | SEED-V')
+
+    known_args, _ = parser.parse_known_args()
+
+    return parser.parse_args()
+
+def get_models(args):
+    # load pretrained model
+    pretrained_model = EEGEncoder(args)
+    
+    # classification head
+    class ClassificationModel(nn.Module):
+        def __init__(self, encoder, num_classes):
+            super().__init__()
+            self.encoder = encoder
+
+            self.full_linear = nn.Linear(args.feature_size, args.feature_size)
+            self.full_gelu = nn.GELU()
+            self.channel_linear = nn.Linear(args.feature_size, args.feature_size)
+            self.channel_gelu = nn.GELU()
+            
+            self.classifier = nn.Sequential(
+                nn.Linear(args.feature_size, args.feature_size // 2),
+                nn.GELU(),
+                nn.Dropout(0.5),
+                nn.Linear(args.feature_size // 2, args.feature_size // 4),
+                nn.GELU(),
+                nn.Dropout(0.3),
+                nn.Linear(args.feature_size // 4, num_classes)
+            )
+            
+        def forward(self, x):
+            bs, ch, seq_len, feature_size = x.shape
+
+            features = self.encoder(x)
+            non_global_tokens = features[:, args.num_global_tokens:]  # [bs, ch*patch size, feature size]
+            non_global_tokens = non_global_tokens.reshape(-1, args.feature_size)
+            processed_features = self.full_linear(non_global_tokens)
+            processed_features = self.full_gelu(processed_features)
+
+            reshaped = processed_features.reshape(bs, ch, seq_len, args.feature_size)
+            
+            channel_pooled = torch.mean(reshaped, dim=1)  # [bs, seq_len, feature_size]
+
+            time_features = channel_pooled.reshape(-1, args.feature_size)  # [bs*seq_len, feature_size]
+            processed_features = self.channel_linear(time_features)  # [bs*seq_len, feature_size]
+            processed_features = self.channel_gelu(processed_features)
+            processed_features = processed_features.reshape(channel_pooled.size(0), seq_len, args.feature_size)  # [bs, seq_len, feature_size]
+            
+            time_pooled = torch.mean(processed_features, dim=1)  # [bs, feature_size]
+            
+            logits = self.classifier(time_pooled)
+            return logits
+    
+    model = ClassificationModel(pretrained_model, args.nb_classes)
+
+    return model
+
+
+def get_dataset(args):
+    if args.dataset == 'TUAB':
+        train_dataset, test_dataset, val_dataset = utils.prepare_TUAB_dataset("/path/to/dataset")
+        args.nb_classes = 1
+        metrics = ["pr_auc", "roc_auc", "accuracy", "balanced_accuracy"]
+    elif args.dataset == 'TUEV':
+        train_dataset, test_dataset, val_dataset = utils.prepare_TUEV_dataset("/path/to/dataset")
+        args.nb_classes = 6
+        metrics = ["accuracy", "balanced_accuracy", "cohen_kappa", "f1_weighted"]
+    elif args.dataset == 'SEED-V':
+        train_dataset, test_dataset, val_dataset = utils.prepare_SEEDV_dataset("/path/to/dataset")
+        args.nb_classes = 5
+        metrics = ["accuracy", "balanced_accuracy", "cohen_kappa", "f1_weighted"]
+    return train_dataset, test_dataset, val_dataset, metrics
+
+
+def main(args, ds_init):
+    if args.output_dir:
+        os.makedirs(args.output_dir, exist_ok=True)
+        current_script_path = os.path.abspath(__file__)
+        script_filename = os.path.basename(__file__)
+        target_script_path = os.path.join(args.output_dir, script_filename)
+        import shutil
+        shutil.copy2(current_script_path, target_script_path)
+        print(f"Copied current script to: {target_script_path}")
+
+    utils.init_distributed_mode(args)
+
+    if ds_init is not None:
+        utils.create_ds_config(args)
+
+    print(args)
+
+    device = torch.device(args.device)
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    cudnn.benchmark = True
+
+    dataset_train, dataset_test, dataset_val, metrics = get_dataset(args)
+
+    if True:  # args.distributed:
+        num_tasks = utils.get_world_size()
+        global_rank = utils.get_rank()
+        sampler_train = torch.utils.data.DistributedSampler(
+            dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+        )
+        print("Sampler_train = %s" % str(sampler_train))
+        if args.dist_eval:
+            if len(dataset_val) % num_tasks != 0:
+                print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
+                      'This will slightly alter validation results as extra duplicate entries are added to achieve '
+                      'equal num of samples per-process.')
+            sampler_val = torch.utils.data.DistributedSampler(
+                dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+            if type(dataset_test) == list:
+                sampler_test = [torch.utils.data.DistributedSampler(
+                    dataset, num_replicas=num_tasks, rank=global_rank, shuffle=False) for dataset in dataset_test]
+            else:
+                sampler_test = torch.utils.data.DistributedSampler(
+                    dataset_test, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+        else:
+            sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+            sampler_test = torch.utils.data.SequentialSampler(dataset_test)
+
+    if global_rank == 0 and args.log_dir is not None:
+        os.makedirs(args.log_dir, exist_ok=True)
+        log_writer = utils.TensorboardLogger(log_dir=args.log_dir)
+    else:
+        log_writer = None
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, sampler=sampler_train,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+    )
+
+    if dataset_val is not None:
+        data_loader_val = torch.utils.data.DataLoader(
+            dataset_val, sampler=sampler_val,
+            batch_size=int(1.5 * args.batch_size),
+            num_workers=args.num_workers,
+            pin_memory=args.pin_mem,
+            drop_last=False
+        )
+        if type(dataset_test) == list:
+            data_loader_test = [torch.utils.data.DataLoader(
+                dataset, sampler=sampler,
+                batch_size=int(1.5 * args.batch_size),
+                num_workers=args.num_workers,
+                pin_memory=args.pin_mem,
+                drop_last=False
+            ) for dataset, sampler in zip(dataset_test, sampler_test)]
+        else:
+            data_loader_test = torch.utils.data.DataLoader(
+                dataset_test, sampler=sampler_test,
+                batch_size=int(1.5 * args.batch_size),
+                num_workers=args.num_workers,
+                pin_memory=args.pin_mem,
+                drop_last=False
+            )
+    else:
+        data_loader_val = None
+        data_loader_test = None
+
+    model = get_models(args)
+
+    if args.finetune:
+        checkpoint = torch.load(args.finetune, map_location='cpu')
+
+        print("Load ckpt from %s" % args.finetune)
+        checkpoint_model = None
+        # for model_key in args.model_key.split('|'):
+        #     if model_key in checkpoint:
+        #         checkpoint_model = checkpoint[model_key]
+        #         print("Load state_dict by model_key = %s" % model_key)
+        #         break
+        if checkpoint_model is None:
+            checkpoint_model = checkpoint['state_dict']
+        if (checkpoint_model is not None):
+            all_keys = list(checkpoint_model.keys())
+            new_dict = OrderedDict()
+            for key in all_keys:
+                print(f"Processing key: {key}")
+                if key.startswith('module.student.'):
+                    new_key = 'encoder' + key[14:]
+                    print(f"Converting key {key} to {new_key}")
+                    new_dict[new_key] = checkpoint_model[key]
+            checkpoint_model = new_dict
+
+        state_dict = model.state_dict()
+        for k in ['head.weight', 'head.bias']:
+            if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
+                print(f"Removing key {k} from pretrained checkpoint")
+                del checkpoint_model[k]
+
+        all_keys = list(checkpoint_model.keys())
+        for key in all_keys:
+            if "relative_position_index" in key:
+                checkpoint_model.pop(key)
+
+        utils.load_state_dict(model, checkpoint_model, prefix=args.model_prefix)
+
+        if args.freeze_all_except_head:
+            print("Freezing all parameters except classification head...")
+            for name, param in model.named_parameters():
+                if 'classifier' not in name and 'channel_linear' not in name and 'full_linear' not in name and 'channel_gelu' not in name and 'full_gelu' not in name:
+                    param.requires_grad = False
+                else:
+                    print(f"Training parameter: {name}")
+                
+        total_params = sum(p.numel() for p in model.parameters())
+        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+        print(f"Number of trainable parameters: {trainable_params}")
+        print('Total number of params:', total_params)
+        print(f'Percentage of trainable parameters: {100 * trainable_params / total_params:.2f}%')
+        
+
+    model.to(device)
+
+    model_ema = None
+
+    model_without_ddp = model
+    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print("Model = %s" % str(model_without_ddp))
+    print('number of params:', n_parameters)
+    print(f'Percentage of trainable parameters: {100 * trainable_params / n_parameters:.2f}%')
+
+    total_batch_size = args.batch_size * args.update_freq * utils.get_world_size()
+    num_training_steps_per_epoch = len(dataset_train) // total_batch_size
+    print("LR = %.8f" % args.lr)
+    print("Batch size = %d" % total_batch_size)
+    print("Update frequent = %d" % args.update_freq)
+    print("Number of training examples = %d" % len(dataset_train))
+    print("Number of training training per epoch = %d" % num_training_steps_per_epoch)
+
+    # num_layers = model_without_ddp.get_num_layers()
+    num_layers = 12
+    if args.layer_decay < 1.0:
+        assigner = LayerDecayValueAssigner(list(args.layer_decay ** (num_layers + 1 - i) for i in range(num_layers + 2)))
+    else:
+        assigner = None
+
+    if assigner is not None:
+        print("Assigned values = %s" % str(assigner.values))
+
+    try:
+        skip_weight_decay_list = model.no_weight_decay()
+    except AttributeError:
+        skip_weight_decay_list = set()
+
+    if args.disable_weight_decay_on_rel_pos_bias:
+        for i in range(num_layers):
+            skip_weight_decay_list.add("blocks.%d.attn.relative_position_bias_table" % i)
+
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
+        model_without_ddp = model.module
+
+    optimizer = create_optimizer(
+        args, model_without_ddp, skip_list=skip_weight_decay_list,
+        get_num_layer=assigner.get_layer_id if assigner is not None else None, 
+        get_layer_scale=assigner.get_scale if assigner is not None else None)
+    loss_scaler = NativeScaler()
+
+    print("Use step level LR scheduler!")
+    lr_schedule_values = utils.cosine_scheduler(
+        args.lr, args.min_lr, args.epochs, num_training_steps_per_epoch,
+        warmup_epochs=args.warmup_epochs, warmup_steps=args.warmup_steps,
+    )
+    if args.weight_decay_end is None:
+        args.weight_decay_end = args.weight_decay
+    wd_schedule_values = utils.cosine_scheduler(
+        args.weight_decay, args.weight_decay_end, args.epochs, num_training_steps_per_epoch)
+    print("Max WD = %.7f, Min WD = %.7f" % (max(wd_schedule_values), min(wd_schedule_values)))
+
+    if args.nb_classes == 1:
+        criterion = torch.nn.BCEWithLogitsLoss()
+    elif args.smoothing > 0.:
+        criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
+    else:
+        criterion = torch.nn.CrossEntropyLoss()
+
+    print("criterion = %s" % str(criterion))
+
+    utils.auto_load_model(
+        args=args, model=model, model_without_ddp=model_without_ddp,
+        optimizer=optimizer, loss_scaler=loss_scaler, model_ema=model_ema)
+            
+    if args.eval:
+        balanced_accuracy = []
+        accuracy = []
+        for data_loader in data_loader_test:
+            test_stats = evaluate(data_loader, model, device, args.output_dir, header='Test:', metrics=metrics, is_binary=(args.nb_classes == 1), epoch=epoch)
+            accuracy.append(test_stats['accuracy'])
+            balanced_accuracy.append(test_stats['balanced_accuracy'])
+        print(f"======Accuracy: {np.mean(accuracy)} {np.std(accuracy)}, balanced accuracy: {np.mean(balanced_accuracy)} {np.std(balanced_accuracy)}")
+        exit(0)
+
+    print(f"Start training for {args.epochs} epochs")
+    start_time = time.time()
+    max_accuracy = 0.0
+    max_accuracy_test = 0.0
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            data_loader_train.sampler.set_epoch(epoch)
+        if log_writer is not None:
+            log_writer.set_step(epoch * num_training_steps_per_epoch * args.update_freq)
+        train_stats = train_one_epoch(
+            model, criterion, data_loader_train, optimizer,
+            device, epoch, loss_scaler, args.clip_grad, model_ema,
+            log_writer=log_writer, start_steps=epoch * num_training_steps_per_epoch,
+            lr_schedule_values=lr_schedule_values, wd_schedule_values=wd_schedule_values,
+            num_training_steps_per_epoch=num_training_steps_per_epoch, update_freq=args.update_freq, 
+            is_binary=args.nb_classes == 1
+        )
+        
+        if args.output_dir and args.save_ckpt:
+            utils.save_model(
+                args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+                loss_scaler=loss_scaler, epoch=epoch, model_ema=model_ema, save_ckpt_freq=args.save_ckpt_freq)
+            
+        if data_loader_val is not None:
+            val_stats = evaluate(data_loader_val, model, device, args.output_dir, header='Val:', 
+                               metrics=metrics, is_binary=args.nb_classes == 1, epoch=epoch)
+            print(f"Accuracy of the network on the {len(dataset_val)} val EEG: {val_stats['accuracy']:.2f}%")
+            test_stats = evaluate(data_loader_test, model, device, args.output_dir, header='Test:', 
+                               metrics=metrics, is_binary=args.nb_classes == 1, epoch=epoch)
+            print(f"Accuracy of the network on the {len(dataset_test)} test EEG: {test_stats['accuracy']:.2f}%")
+            
+            if max_accuracy < val_stats["accuracy"]:
+                max_accuracy = val_stats["accuracy"]
+                if args.output_dir and args.save_ckpt:
+                    utils.save_model(
+                        args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+                        loss_scaler=loss_scaler, epoch="best", model_ema=model_ema)
+                max_accuracy_test = test_stats["accuracy"]
+
+            print(f'Max accuracy val: {max_accuracy:.2f}%, max accuracy test: {max_accuracy_test:.2f}%')
+            if log_writer is not None:
+                for key, value in val_stats.items():
+                    if key == 'accuracy':
+                        log_writer.update(accuracy=value, head="val", step=epoch)
+                    elif key == 'balanced_accuracy':
+                        log_writer.update(balanced_accuracy=value, head="val", step=epoch)
+                    elif key == 'f1_weighted':
+                        log_writer.update(f1_weighted=value, head="val", step=epoch)
+                    elif key == 'pr_auc':
+                        log_writer.update(pr_auc=value, head="val", step=epoch)
+                    elif key == 'roc_auc':
+                        log_writer.update(roc_auc=value, head="val", step=epoch)
+                    elif key == 'cohen_kappa':
+                        log_writer.update(cohen_kappa=value, head="val", step=epoch)
+                    elif key == 'loss':
+                        log_writer.update(loss=value, head="val", step=epoch)
+                for key, value in test_stats.items():
+                    if key == 'accuracy':
+                        log_writer.update(accuracy=value, head="test", step=epoch)
+                    elif key == 'balanced_accuracy':
+                        log_writer.update(balanced_accuracy=value, head="test", step=epoch)
+                    elif key == 'f1_weighted':
+                        log_writer.update(f1_weighted=value, head="test", step=epoch)
+                    elif key == 'pr_auc':
+                        log_writer.update(pr_auc=value, head="test", step=epoch)
+                    elif key == 'roc_auc':
+                        log_writer.update(roc_auc=value, head="test", step=epoch)
+                    elif key == 'cohen_kappa':
+                        log_writer.update(cohen_kappa=value, head="test", step=epoch)
+                    elif key == 'loss':
+                        log_writer.update(loss=value, head="test", step=epoch)
+                
+            log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                         **{f'val_{k}': v for k, v in val_stats.items()},
+                         **{f'test_{k}': v for k, v in test_stats.items()},
+                         'epoch': epoch,
+                         'n_parameters': n_parameters}
+        else:
+            log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                         'epoch': epoch,
+                         'n_parameters': n_parameters}
+
+        if args.output_dir and utils.is_main_process():
+            if log_writer is not None:
+                log_writer.flush()
+            
+            args_dict = vars(args)
+            serializable_args = {
+                k: v if isinstance(v, (int, float, str, bool, type(None))) else str(v) 
+                for k, v in args_dict.items()
+            }
+            log_stats['args'] = serializable_args
+            
+            with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
+                f.write(json.dumps(log_stats) + "\n")
+
+        print(f"Epoch {epoch} confusion matrix:")
+        print(np.array(test_stats['confusion_matrix']))
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Training time {}'.format(total_time_str))
+
+
+if __name__ == '__main__':
+    opts = get_args()
+    ds_init = None
+    if opts.output_dir:
+        Path(opts.output_dir).mkdir(parents=True, exist_ok=True)
+    main(opts, ds_init)

utils.py

@@ -0,0 +1,804 @@
+# --------------------------------------------------------
+# EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation
+# Based on BEiT-v2, timm, DeiT, DINO v2, LaBraM and CBraMod code bases
+# 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/dinov2
+# https://github.com/935963004/LaBraM
+# https://github.com/wjq-learning/CBraMod
+# ---------------------------------------------------------
+
+import io
+import os
+import math
+import time
+import json
+import glob
+from collections import defaultdict, deque
+import datetime
+import numpy as np
+from timm.utils import get_state_dict
+
+from pathlib import Path
+import argparse
+
+import torch
+import torch.distributed as dist
+from torch import inf
+
+from tensorboardX import SummaryWriter
+import pickle
+from scipy.signal import resample
+from pyhealth.metrics import binary_metrics_fn, multiclass_metrics_fn
+
+def bool_flag(s):
+    """
+    Parse boolean arguments from the command line.
+    """
+    FALSY_STRINGS = {"off", "false", "0"}
+    TRUTHY_STRINGS = {"on", "true", "1"}
+    if s.lower() in FALSY_STRINGS:
+        return False
+    elif s.lower() in TRUTHY_STRINGS:
+        return True
+    else:
+        raise argparse.ArgumentTypeError("invalid value for a boolean flag")
+
+def get_model(model):
+    if isinstance(model, torch.nn.DataParallel) \
+      or isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        return model.module
+    else:
+        return model
+            
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}'
+        ]
+        if torch.cuda.is_available():
+            log_msg.append('max mem: {memory:.0f}')
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+class TensorboardLogger(object):
+    def __init__(self, log_dir):
+        self.writer = SummaryWriter(logdir=log_dir)
+        self.step = 0
+
+    def set_step(self, step=None):
+        if step is not None:
+            self.step = step
+        else:
+            self.step += 1
+
+    def update(self, head='scalar', step=None, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.writer.add_scalar(head + "/" + k, v, self.step if step is None else step)
+    
+    def update_image(self, head='images', step=None, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            self.writer.add_image(head + "/" + k, v, self.step if step is None else step)
+            
+    def flush(self):
+        self.writer.flush()
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+def all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=False):
+    world_size = get_world_size()
+
+    if world_size == 1:
+        return tensor
+    dist.all_reduce(tensor, op=op, async_op=async_op)
+
+    return tensor
+
+def all_gather_batch(tensors):
+    """
+    Performs all_gather operation on the provided tensors.
+    """
+    # Queue the gathered tensors
+    world_size = get_world_size()
+    # There is no need for reduction in the single-proc case
+    if world_size == 1:
+        return tensors
+    tensor_list = []
+    output_tensor = []
+    for tensor in tensors:
+        tensor_all = [torch.ones_like(tensor) for _ in range(world_size)]
+        dist.all_gather(
+            tensor_all,
+            tensor,
+            async_op=False  # performance opt
+        )
+
+        tensor_list.append(tensor_all)
+
+    for tensor_all in tensor_list:
+        output_tensor.append(torch.cat(tensor_all, dim=0))
+    return output_tensor
+
+class GatherLayer(torch.autograd.Function):
+    """
+    Gather tensors from all workers with support for backward propagation:
+    This implementation does not cut the gradients as torch.distributed.all_gather does.
+    """
+
+    @staticmethod
+    def forward(ctx, x):
+        output = [torch.zeros_like(x) for _ in range(dist.get_world_size())]
+        dist.all_gather(output, x)
+        return tuple(output)
+
+    @staticmethod
+    def backward(ctx, *grads):
+        all_gradients = torch.stack(grads)
+        dist.all_reduce(all_gradients)
+        return all_gradients[dist.get_rank()]
+
+
+def all_gather_batch_with_grad(tensors):
+    """
+    Performs all_gather operation on the provided tensors.
+    Graph remains connected for backward grad computation.
+    """
+    # Queue the gathered tensors
+    world_size = get_world_size()
+    # There is no need for reduction in the single-proc case
+    if world_size == 1:
+        return tensors
+    tensor_list = []
+    output_tensor = []
+
+    for tensor in tensors:
+        tensor_all = GatherLayer.apply(tensor)
+        tensor_list.append(tensor_all)
+
+    for tensor_all in tensor_list:
+        output_tensor.append(torch.cat(tensor_all, dim=0))
+    return output_tensor
+
+def _get_rank_env():
+    if "RANK" in os.environ:
+        return int(os.environ["RANK"])
+    else:
+        return int(os.environ['OMPI_COMM_WORLD_RANK'])
+
+
+def _get_local_rank_env():
+    if "LOCAL_RANK" in os.environ:
+        return int(os.environ["LOCAL_RANK"])
+    else:
+        return int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+
+
+def _get_world_size_env():
+    if "WORLD_SIZE" in os.environ:
+        return int(os.environ["WORLD_SIZE"])
+    else:
+        return int(os.environ['OMPI_COMM_WORLD_SIZE'])
+
+
+def init_distributed_mode(args):
+    if args.dist_on_itp:
+        args.rank = _get_rank_env()
+        args.world_size = _get_world_size_env()  # int(os.environ['OMPI_COMM_WORLD_SIZE'])
+        args.gpu = _get_local_rank_env()
+        args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
+        os.environ['LOCAL_RANK'] = str(args.gpu)
+        os.environ['RANK'] = str(args.rank)
+        os.environ['WORLD_SIZE'] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}, gpu {}'.format(
+        args.rank, args.dist_url, args.gpu), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+def load_state_dict(model, state_dict, prefix='', ignore_missing="relative_position_index"):
+    missing_keys = []
+    unexpected_keys = []
+    error_msgs = []
+    # copy state_dict so _load_from_state_dict can modify it
+    metadata = getattr(state_dict, '_metadata', None)
+    state_dict = state_dict.copy()
+    if metadata is not None:
+        state_dict._metadata = metadata
+
+    def load(module, prefix=''):
+        local_metadata = {} if metadata is None else metadata.get(
+            prefix[:-1], {})
+        module._load_from_state_dict(
+            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
+        for name, child in module._modules.items():
+            if child is not None:
+                load(child, prefix + name + '.')
+
+    load(model, prefix=prefix)
+
+    warn_missing_keys = []
+    ignore_missing_keys = []
+    for key in missing_keys:
+        keep_flag = True
+        for ignore_key in ignore_missing.split('|'):
+            if ignore_key in key:
+                keep_flag = False
+                break
+        if keep_flag:
+            warn_missing_keys.append(key)
+        else:
+            ignore_missing_keys.append(key)
+
+    missing_keys = warn_missing_keys
+
+    if len(missing_keys) > 0:
+        print("Weights of {} not initialized from pretrained model: {}".format(
+            model.__class__.__name__, missing_keys))
+    if len(unexpected_keys) > 0:
+        print("Weights from pretrained model not used in {}: {}".format(
+            model.__class__.__name__, unexpected_keys))
+    if len(ignore_missing_keys) > 0:
+        print("Ignored weights of {} not initialized from pretrained model: {}".format(
+            model.__class__.__name__, ignore_missing_keys))
+    if len(error_msgs) > 0:
+        print('\n'.join(error_msgs))
+
+def get_grad_norm(parameters, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+    total_norm = total_norm ** (1. / norm_type)
+    return total_norm
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True, layer_names=None):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters, layer_names=layer_names)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict): 
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0, layer_names=None) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    
+    parameters = [p for p in parameters if p.grad is not None]
+        
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.)
+    device = parameters[0].grad.device
+    
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        # total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
+        layer_norm = torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters])
+        total_norm = torch.norm(layer_norm, norm_type)
+        # print(layer_norm.max(dim=0))
+        
+        if layer_names is not None:
+            if torch.isnan(total_norm) or torch.isinf(total_norm) or total_norm > 1.0:
+                value_top, name_top = torch.topk(layer_norm, k=5)
+                print(f"Top norm value: {value_top}")
+                print(f"Top norm name: {[layer_names[i][7:] for i in name_top.tolist()]}")
+        
+    return total_norm
+
+
+def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0,
+                     start_warmup_value=0, warmup_steps=-1):
+    warmup_schedule = np.array([])
+    warmup_iters = warmup_epochs * niter_per_ep
+    if warmup_steps > 0:
+        warmup_iters = warmup_steps
+    print("Set warmup steps = %d" % warmup_iters)
+    if warmup_epochs > 0:
+        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+    iters = np.arange(epochs * niter_per_ep - warmup_iters)
+    schedule = np.array(
+        [final_value + 0.5 * (base_value - final_value) * (1 + math.cos(math.pi * i / (len(iters)))) for i in iters])
+
+    schedule = np.concatenate((warmup_schedule, schedule))
+
+    assert len(schedule) == epochs * niter_per_ep
+    return schedule
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler, model_ema=None, optimizer_disc=None, save_ckpt_freq=1):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+
+    if not getattr(args, 'enable_deepspeed', False):
+        checkpoint_paths = [output_dir / 'checkpoint.pth']
+        if epoch == 'best':
+            checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name),]
+        elif (epoch + 1) % save_ckpt_freq == 0:
+            checkpoint_paths.append(output_dir / ('checkpoint-%s.pth' % epoch_name))
+
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                # 'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+            if loss_scaler is not None:
+                to_save['scaler'] = loss_scaler.state_dict()
+
+            if model_ema is not None:
+                to_save['model_ema'] = get_state_dict(model_ema)
+                
+            if optimizer_disc is not None:
+                to_save['optimizer_disc'] = optimizer_disc.state_dict()
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        if model_ema is not None:
+            client_state['model_ema'] = get_state_dict(model_ema)
+        model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)           
+
+def auto_load_model(args, model, model_without_ddp, optimizer, loss_scaler, model_ema=None, optimizer_disc=None):
+    output_dir = Path(args.output_dir)
+    
+    if not getattr(args, 'enable_deepspeed', False):
+        # torch.amp
+        if args.auto_resume and len(args.resume) == 0:
+            all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint.pth'))
+            if len(all_checkpoints) > 0:
+                args.resume = os.path.join(output_dir, 'checkpoint.pth')
+            else:
+                all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*.pth'))
+                latest_ckpt = -1
+                for ckpt in all_checkpoints:
+                    t = ckpt.split('-')[-1].split('.')[0]
+                    if t.isdigit():
+                        latest_ckpt = max(int(t), latest_ckpt)
+                if latest_ckpt >= 0:
+                    args.resume = os.path.join(output_dir, 'checkpoint-%d.pth' % latest_ckpt)
+            print("Auto resume checkpoint: %s" % args.resume)
+
+        if args.resume:
+            if args.resume.startswith('https'):
+                checkpoint = torch.hub.load_state_dict_from_url(
+                    args.resume, map_location='cpu', check_hash=True)
+            else:
+                checkpoint = torch.load(args.resume, map_location='cpu', weights_only=False)
+            model_without_ddp.load_state_dict(checkpoint['model']) # strict: bool=True, , strict=False
+            print("Resume checkpoint %s" % args.resume)
+            if 'optimizer' in checkpoint and 'epoch' in checkpoint:
+                optimizer.load_state_dict(checkpoint['optimizer'])
+                print(f"Resume checkpoint at epoch {checkpoint['epoch']}")
+                args.start_epoch = 1#checkpoint['epoch'] + 1
+                if 'scaler' in checkpoint:
+                    loss_scaler.load_state_dict(checkpoint['scaler'])
+                print("With optim & sched!")
+            if 'optimizer_disc' in checkpoint:
+                optimizer_disc.load_state_dict(checkpoint['optimizer_disc'])
+    else:
+        # deepspeed, only support '--auto_resume'.
+        if args.auto_resume:
+            all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*'))
+            latest_ckpt = -1
+            for ckpt in all_checkpoints:
+                t = ckpt.split('-')[-1].split('.')[0]
+                if t.isdigit():
+                    latest_ckpt = max(int(t), latest_ckpt)
+            if latest_ckpt >= 0:
+                args.resume = os.path.join(output_dir, 'checkpoint-%d' % latest_ckpt)
+                print("Auto resume checkpoint: %d" % latest_ckpt)
+                _, client_states = model.load_checkpoint(args.output_dir, tag='checkpoint-%d' % latest_ckpt)
+                args.start_epoch = client_states['epoch'] + 1
+
+def create_ds_config(args):
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    with open(os.path.join(args.output_dir, "latest"), mode="w") as f:
+        pass
+
+    args.deepspeed_config = os.path.join(args.output_dir, "deepspeed_config.json")
+    with open(args.deepspeed_config, mode="w") as writer:
+        ds_config = {
+            "train_batch_size": args.batch_size * args.update_freq * get_world_size(),
+            "train_micro_batch_size_per_gpu": args.batch_size,
+            "steps_per_print": 1000,
+            "optimizer": {
+                "type": "Adam",
+                "adam_w_mode": True,
+                "params": {
+                    "lr": args.lr,
+                    "weight_decay": args.weight_decay,
+                    "bias_correction": True,
+                    "betas": [
+                        0.9,
+                        0.999
+                    ],
+                    "eps": 1e-8
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 0,
+                "initial_scale_power": 7,
+                "loss_scale_window": 128
+            }
+        }
+
+        writer.write(json.dumps(ds_config, indent=2))
+
+class TUABLoader(torch.utils.data.Dataset):
+    def __init__(self, root, files, sampling_rate=200):
+        self.root = root
+        self.files = files
+        self.default_rate = 200
+        self.sampling_rate = sampling_rate
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, index):
+        sample = pickle.load(open(os.path.join(self.root, self.files[index]), "rb"))
+        X = sample["X"]
+        if self.sampling_rate != self.default_rate:
+            X = resample(X, 10 * self.sampling_rate, axis=-1)
+        Y = sample["y"]
+        X = torch.FloatTensor(X)
+        return X, Y
+    
+
+class TUEVLoader(torch.utils.data.Dataset):
+    def __init__(self, root, files, sampling_rate=200):
+        self.root = root
+        self.files = files
+        self.default_rate = 200
+        self.sampling_rate = sampling_rate
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, index):
+        sample = pickle.load(open(os.path.join(self.root, self.files[index]), "rb"))
+        X = sample["signal"]
+        
+        if self.sampling_rate != self.default_rate:
+            X = resample(X, 5 * self.sampling_rate, axis=-1)
+        Y = int(sample["label"][0] - 1)
+        X = torch.FloatTensor(X)
+        return X, Y
+
+class SEEDVLoader(torch.utils.data.Dataset):
+    def __init__(self, root, files, sampling_rate=200):
+        self.root = root
+        self.files = files
+        self.default_rate = 200
+        self.sampling_rate = sampling_rate
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, index):
+        sample = pickle.load(open(os.path.join(self.root, self.files[index]), "rb"))
+        X = sample["X"]
+        if self.sampling_rate != self.default_rate:
+            X = resample(X, self.sampling_rate, axis=-1)
+        Y = int(sample["y"])
+        X = torch.FloatTensor(X)
+        return X, Y
+
+def prepare_TUEV_dataset(root):
+    # set random seed
+    seed = 8250
+    np.random.seed(seed)
+
+    train_files = os.listdir(os.path.join(root, "processed_train"))
+    val_files = os.listdir(os.path.join(root, "processed_eval"))
+    test_files = os.listdir(os.path.join(root, "processed_test"))
+
+    # prepare training and test data loader
+    train_dataset = TUEVLoader(
+        os.path.join(
+            root, "processed_train"), train_files
+    )
+    test_dataset = TUEVLoader(
+        os.path.join(
+            root, "processed_test"), test_files
+    )
+    val_dataset = TUEVLoader(
+        os.path.join(
+            root, "processed_eval"), val_files
+    )
+    print(len(train_files), len(val_files), len(test_files))
+    return train_dataset, test_dataset, val_dataset
+
+
+def prepare_TUAB_dataset(root):
+    # set random seed
+    seed = 12345
+    np.random.seed(seed)
+
+    train_files = os.listdir(os.path.join(root, "train"))
+    np.random.shuffle(train_files)
+    val_files = os.listdir(os.path.join(root, "val"))
+    test_files = os.listdir(os.path.join(root, "test"))
+
+    print(len(train_files), len(val_files), len(test_files))
+
+    # prepare training and test data loader
+    train_dataset = TUABLoader(os.path.join(root, "train"), train_files)
+    test_dataset = TUABLoader(os.path.join(root, "test"), test_files)
+    val_dataset = TUABLoader(os.path.join(root, "val"), val_files)
+    print(len(train_files), len(val_files), len(test_files))
+    return train_dataset, test_dataset, val_dataset
+
+def prepare_SEEDV_dataset(root):
+    # set random seed
+    seed = 8250
+    np.random.seed(seed)
+
+    train_files = os.listdir(os.path.join(root, "train"))
+    np.random.shuffle(train_files)
+    val_files = os.listdir(os.path.join(root, "val"))
+    test_files = os.listdir(os.path.join(root, "test"))
+
+    print(len(train_files), len(val_files), len(test_files))
+
+    # prepare training and test data loader
+    train_dataset = SEEDVLoader(os.path.join(root, "train"), train_files)
+    test_dataset = SEEDVLoader(os.path.join(root, "test"), test_files)
+    val_dataset = SEEDVLoader(os.path.join(root, "val"), val_files)
+    print(len(train_files), len(val_files), len(test_files))
+    return train_dataset, test_dataset, val_dataset
+
+def get_metrics(output, target, metrics, is_binary, threshold=0.5):
+    if is_binary:
+        if 'roc_auc' not in metrics or sum(target) * (len(target) - sum(target)) != 0:  # to prevent all 0 or all 1 and raise the AUROC error
+            results = binary_metrics_fn(
+                target,
+                output,
+                metrics=metrics,
+                threshold=threshold,
+            )
+        else:
+            results = {
+                "accuracy": 0.0,
+                "balanced_accuracy": 0.0,
+                "pr_auc": 0.0,
+                "roc_auc": 0.0,
+            }
+    else:
+        results = multiclass_metrics_fn(
+            target, output, metrics=metrics
+        )
+    return results