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