SleepMaMi / evaluate.py

KeondoPark

Initial commit for public release

35f60b9 5 months ago

21.4 kB

from collections import defaultdict
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from torchmetrics import ConfusionMatrix, CohenKappa
from torchmetrics.classification import Accuracy, Precision, Recall, F1Score 
import torch
import torch.nn as nn
import logging
import sys



import os

from pytorch_lightning.callbacks import ModelCheckpoint
import numpy as np

import data.data_loader as data_loader
import yaml

import models
from models.utils import cox_ph_loss, CIndex, AUROC
from torchmetrics import MetricCollection

import random
import importlib
import argparse
from configs.global_cfg import TARGET_RESAMPLE
from configs.global_cfg import TARGET_RESAMPLE, DATASET_PATH, CLASSIFICATION_TASKS, REGRESSION_TASKS
import pickle
from einops import rearrange
import pandas as pd

torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.cuda.manual_seed_all(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(42)
random.seed(42)

# In H100 GPUs
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False


BASE_DIR = os.path.dirname(os.path.abspath(__file__))

class EvalDataModule(pl.LightningDataModule):
    def __init__(self,                 
                data_root,                                               
                batch_size_v,
                transforms_v,
                num_workers,                 
                proportion_v,                 
                num_classes=5, 
                splits_set=None,
                modalities=["EEG"],
                sig_len_target=30,                
                ft_task='sleep_staging',
                full_night=False,
                seq_truncate=False
                ):
        super().__init__()                

        self.data_root = data_root        
        self.transforms_v = transforms_v                
        self.batch_size_v = batch_size_v
        self.num_workers = num_workers        
        self.proportion_v = proportion_v
        self.num_classes = num_classes
        self.splits_set = splits_set        
        self.modalities = modalities
        self.sig_len_target = sig_len_target
        self.ft_task = ft_task
        self.full_night = full_night
        self.seq_truncate = seq_truncate

        
    def setup(self, stage=None):       
        
        self.test_data_list = self.splits_set['test']        
        

        self.test_loader = data_loader.build_loader(
            self.data_root, self.test_data_list,
            batch_size=self.batch_size_v, shuffle=False,
            transforms=self.transforms_v, sampler=None,
            num_workers=self.num_workers, proportion=self.proportion_v,
            split='test',
            num_classes=self.num_classes, modalities=self.modalities,
            sig_len_target=self.sig_len_target,
            ft_task=self.ft_task, full_night=self.full_night, seq_truncate=self.seq_truncate)

        print("Test loader ready")
        print("Length of test_loader: ", len(self.test_loader))

    
    def test_dataloader(self):
        return self.test_loader
    
from finetune import ConvMAE1d_Finetune

class ConvMAE1d_Evaluate(ConvMAE1d_Finetune):
    def __init__(self, *args, **kwargs):        
        super().__init__(*args, **kwargs)
        
        self.save_results = self.cfg.get('save_results', False)
        if self.save_results:
            self.save_results_dir = "ft_results"
            os.makedirs(self.save_results_dir, exist_ok=True)
                # self.preds = defaultdict(list)
                # self.gt = defaultdict(list)            
            self.results = defaultdict(dict)  # Initialize nested dict for each PID
        
        self.save_emb = self.cfg.get('save_emb', False)
        if self.save_emb:
            self.save_emb_dir = "ft_results/embeddings"
            os.makedirs(self.save_emb_dir, exist_ok=True)
            self.embeddings = defaultdict(dict)

        self._peak_gpu_mem_bytes = None

    def on_test_epoch_start(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.reset_peak_memory_stats()
            self._peak_gpu_mem_bytes = None
            
    
    def forward_epoch_level(self, batch, mode='test'):
        inputs, labels, _, _, pids = batch

        # signals = inputs["EEG"].to(self.device, non_blocking=True)
        for key in inputs.keys():                            
            inputs[key] = inputs[key].to(self.device, non_blocking=True) 

        pred = self(inputs)     
        if self.ft_task == 'apnea_seg':            
            pred = rearrange(pred, 'n l c -> n c l')      
                                
        loss = self.criterion(pred, labels)   
        pred_cls = pred.argmax(dim=1)  
        
        if self.save_results:
            pred_cls = pred_cls.detach().cpu().numpy()
            labels = labels.detach().cpu().numpy()
            for idx, pid in enumerate(pids):
                if pid not in self.results:
                    self.results[pid] = {'pred': [], 'gt': [],}
                self.results[pid]['pred'].append(pred_cls[idx])
                self.results[pid]['gt'].append(labels[idx])

        self.log_metrics(loss, pred_cls, labels, mode=mode)           

            
        return loss

    def forward_full_night(self, batch, mode='test'):
        inputs, labels, mask, demo, pids = batch
        
        if self.ft_task in ['age','ahi']:
            labels = demo[self.ft_task].to(self.device, dtype=torch.float32, non_blocking=True)
        elif self.ft_task in ['sex']:
            labels = demo[self.ft_task].to(self.device, dtype=torch.long, non_blocking=True) - 1
        
        for key in inputs.keys():                            
            inputs[key] = inputs[key].to(self.device, non_blocking=True) 

        if (mask is not None) and (mask.numel() > 0):
            mask = mask.to(self.device, non_blocking=True)                    
            
        if self.save_emb:
            self.model.eval()
            with torch.no_grad():                
                (feats_macro_f, feats_macro_r), feats_micro = self.model.forward_features(inputs, mask=mask)                
                from models.utils import flip_padded_sequences
                feats_macro_r = flip_padded_sequences(feats_macro_r, mask)
                feats_macro = torch.cat([feats_macro_f, feats_macro_r], dim=-1)
                
                for idx, pid in enumerate(pids):                    
                    self.embeddings[pid] = {
                        'macro': feats_macro[idx].detach().cpu().numpy(),
                        'micro': feats_micro[idx].detach().cpu().numpy(),
                        'mask': mask[idx].detach().cpu().numpy() if mask is not None and mask.numel() > 0 else None,
                        'age': demo['age'][idx].detach().cpu().numpy(),
                        'sex': demo['sex'][idx].detach().cpu().numpy(),
                        'bmi': demo['bmi'][idx].detach().cpu().numpy(),
                        'ahi': demo['ahi'][idx].detach().cpu().numpy()
                    }
            return  # Skip further processing if only saving embeddings
                
        pred = self(inputs, mask=mask)
           
        if self.ft_task == 'sleep_staging' and (mask is not None) and (mask.numel() > 0):            
            pred = pred.permute(0,2,1)[mask]            
            labels = labels[mask]                
        
        if self.ft_task in REGRESSION_TASKS:
            pred = pred.squeeze()        
        loss = self.criterion(pred, labels)        
        
        if self.ft_task in CLASSIFICATION_TASKS:            
            pred = pred.argmax(dim=1)

        if self.save_results:           
            self.update_results(pred, labels, pids, demo)
        self.log_metrics(loss, pred, labels, mode=mode)     
    
    def forward_diagnosis(self, batch, mode='train'):
        inputs, is_event, event_times, mask, demo, pids = batch

        # signals = inputs["EEG"].to(self.device, non_blocking=True)
        for key in inputs.keys():                            
            inputs[key] = inputs[key].to(self.device, non_blocking=True) 

        for key in demo.keys():
            demo[key] = demo[key].to(self.device, non_blocking=True)

        if (mask is not None) and (mask.numel() > 0):
            mask = mask.to(self.device, non_blocking=True)                    
            
        pred = self(inputs, mask=mask, demo=demo)        
            
        loss = self.criterion(pred, event_times, is_event)
        if mode == 'train':
            self.log('train_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True)        
            self.train_cindex.update(pred.squeeze(), event_times, is_event)
            self.train_auroc.update(pred.squeeze(), event_times, is_event)
        elif mode == 'val':
            self.log('val_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True, batch_size=self.cfg['batch_size_v'])        
            self.val_cindex.update(pred.squeeze(), event_times, is_event)
            self.val_auroc.update(pred.squeeze(), event_times, is_event)
        elif mode == 'test':
            self.log('test_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True, batch_size=self.cfg['batch_size_v'])        
            self.test_cindex.update(pred.squeeze(), event_times, is_event)
            self.test_auroc.update(pred.squeeze(), event_times, is_event)
            
        if self.save_results:
            labels = torch.stack([event_times, is_event], dim=1)
            self.update_results(pred, labels, pids, demo)

        return loss
    
    
    def update_results(self, pred, labels, pids, demo):
        def _to_numpy(x):
            x = x.detach()
            if x.dtype in (torch.bfloat16, torch.float16, torch.float32, torch.float64):
                x = x.float()
            return x.cpu().numpy()

        pred_np = _to_numpy(pred)
        labels_np = _to_numpy(labels)
        for idx, pid in enumerate(pids):
            if pid not in self.results:
                self.results[pid] = {'pred': [], 
                                     'gt': [], 
                                     'age': _to_numpy(demo['age'][idx]),
                                     'sex': _to_numpy(demo['sex'][idx]),
                                     'bmi': _to_numpy(demo['bmi'][idx]),
                                     'ahi': _to_numpy(demo['ahi'][idx])
                                     }
            self.results[pid]['pred'].append(pred_np[idx])
            self.results[pid]['gt'].append(labels_np[idx])
            
    
    
    def on_test_epoch_end(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            self._peak_gpu_mem_bytes = torch.cuda.max_memory_allocated()
            peak_gb = self._peak_gpu_mem_bytes / (1024 ** 3)
            self.log('test_peak_gpu_mem_gb', peak_gb, on_step=False, on_epoch=True, prog_bar=False)
            print(f"Peak GPU memory allocated: {peak_gb:.3f} GB")

        if self.ft_task == 'diagnosis':
            cindex = self.test_cindex.compute()
            auroc = self.test_auroc.compute()            
            for i in range(len(cindex)):
                self.log(f'test_cindex_class_{i}', cindex[i], on_step=False, on_epoch=True, prog_bar=False)
                self.log(f'test_auroc_class_{i}', auroc[i], on_step=False, on_epoch=True, prog_bar=False)
            self.log('test_cindex', torch.mean(cindex), on_step=False, on_epoch=True, prog_bar=True)
            self.log('test_auroc', torch.mean(auroc), on_step=False, on_epoch=True, prog_bar=True)
            print("Test C-Index:\n", cindex)
            print("Test AUROC:\n", auroc)
            
            self.test_cindex.reset()
            self.test_auroc.reset()
        else:
            self.log_dict(self.test_metrics.compute(), on_epoch=True, sync_dist=True)            
            print("Test Confusion Matrix:\n", self.test_conf_matrix.compute())
        
        if self.save_results:            
            # with open(os.path.join(self.save_results_dir, f'preds_{self.cfg["experiment_name"]}.pickle'), 'wb') as f:
            #     pickle.dump(self.preds, f)

            # with open(os.path.join(self.save_results_dir, f'gt_{self.cfg["experiment_name"]}.pickle'), 'wb') as f:
            #     pickle.dump(self.gt, f)
            
            with open(os.path.join(self.save_results_dir, f'results_{self.cfg["experiment_name"]}.pickle'), 'wb') as f:
                pickle.dump(self.results, f)
                
        if self.save_emb:
            for pid, emb in self.embeddings.items():
                with open(os.path.join(self.save_emb_dir, f'embeddings_{self.cfg["experiment_name"]}_{pid}.pickle'), 'wb') as f:
                    pickle.dump(emb, f)
    
def get_all_configs_from_yaml(file_path):
    with open(file_path, 'r') as file:
        configs = yaml.safe_load(file)
    return configs
    # return Config(configs)

def main():
    
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="Train Perceiver Model")
    parser.add_argument('--config', type=str, required=True, help="Name of the configuration file (without .py)")
    parser.add_argument('--no_wandb', action='store_true', default=False, help="Disable Weights and Biases logging")
    args = parser.parse_args()   

    all_configs = get_all_configs_from_yaml(os.path.join("configs", args.config + ".yaml"))
    print(all_configs)    

    # Add transforms dynamically if needed
    # args.transforms = get_transforms()
    # print(args)

    print("Starting Experiment: ", all_configs['experiment_name'])    
        
    config_save_path = os.path.join(all_configs['config_save_path'], all_configs['experiment_name'])
     
    os.makedirs(config_save_path, exist_ok=True)
    with open(os.path.join(config_save_path, 'config.yaml'), 'w') as file:
        yaml.dump(all_configs, file)    

    # Set data paths 
    # data_root("datasets") directory should be in the same directory as this script
    data_root = os.path.join(BASE_DIR, all_configs['data_root'])
    datasets = all_configs.get('datasets', ['phy'])


    # Load data
    print("Loading Data")        
    splits_set = dict()
    splits_set['train'] = []    
    splits_set['val'] = []
    splits_set['test'] = []

    for ds in datasets:
        # Use Physionet 2018 dataset        
        if ds == 'kiss':

            sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects

            sample_dir = os.path.join(data_root, DATASET_PATH[ds], 'test', sample_modality)
            chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
            info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
            with open(info_path, 'rb') as f:
                sig_info = pickle.load(f)

            test_split = sorted([sub for sub in sig_info['sig_info'].keys()])

            if all_configs.get('save_emb',False):
                test_split = test_split[:500]  # For saving embeddings, limit to first 500 subjects only to save space

            splits_set['test'].extend([('kiss',f) for f in test_split])

        elif ds == 'shhs1':

            sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects
            sample_dir = os.path.join(data_root, DATASET_PATH[ds], 'test', sample_modality)
            chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
            info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
            with open(info_path, 'rb') as f:
                sig_info = pickle.load(f)
            test_subjects = sorted([sub for sub in sig_info['sig_info'].keys()])

            def get_shhs_lists(shhs_split_file):
                ids = pd.read_csv(os.path.join("data","shhs_meta_data", shhs_split_file), header=None, names=['a'])
                ids = ids['a'].values.tolist()    
                list_files = ["shhs1-20"+id[11:15] for id in ids]
                return list_files            

            list_files_test = get_shhs_lists("test_list.txt")
            test_subjects = set(list_files_test) & set(test_subjects)
            test_subjects = sorted(list(test_subjects))            
            
            if all_configs.get('save_emb',False):
                test_subjects = test_subjects[:500]  # For saving embeddings, limit to first 500 subjects only to save space
            splits_set['test'].extend([('shhs1',f) for f in test_subjects])


        elif 'isruc' in ds:

            sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects    

            sample_dir = os.path.join(data_root, DATASET_PATH[ds], 'test', sample_modality)
            info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
            with open(info_path, 'rb') as f:
                sig_info = pickle.load(f)
            test_subjects = sorted([sub for sub in sig_info['sig_info'].keys()])

            splits_set['test'].extend([(ds,f) for f in test_subjects])

    print("Number of test subjects: ", len(splits_set['test']))         

    if all_configs['pretrain_checkpoint']:
        pretrain_exp_name = all_configs['pretrain_checkpoint'].split('-')[0]
        pretrain_config = get_all_configs_from_yaml(os.path.join("experiments", pretrain_exp_name, "config.yaml"))        
        use_shared = pretrain_config.get('use_shared', all_configs.get('use_shared', True))
        use_moe = pretrain_config.get('use_moe', all_configs.get('use_moe', True))
    else:        
        use_shared = all_configs.get('use_shared', True)
        use_moe = all_configs.get('use_moe', True)
        

    sig_len_target = all_configs.get('sig_len_target', 60)
    data_module = EvalDataModule(   
        data_root = all_configs['data_root'],                
        transforms_v = None,              
        batch_size_v = all_configs['batch_size_v'],        
        num_workers = all_configs['num_workers'],        
        proportion_v = all_configs['proportion_v'],        
        num_classes = all_configs['num_classes'],
        splits_set = splits_set,
        modalities = all_configs['modalities'],
        sig_len_target=sig_len_target,
        ft_task=all_configs.get('ft_task', 'sleep_staging'),
        full_night=all_configs.get('full_night', False),
        seq_truncate=all_configs.get('seq_truncate', False)
        )
    data_module.setup()  # Setup the data module to prepare the datasets and dataloaders    
    
    
    print("Model Loading ...")
    model = ConvMAE1d_Evaluate(
        model_name=all_configs['model_name'],                
        num_classes = all_configs['num_classes'],                
        modalities = all_configs['modalities'],
        num_experts = all_configs['num_experts'],
        num_experts_per = all_configs['num_experts_per'],
        loss = all_configs['loss'],        
        sig_len = sig_len_target * TARGET_RESAMPLE,
        use_shared = use_shared,
        use_moe = use_moe,
        full_finetune=all_configs.get('full_finetune', False),  # Default to False if not specified                
        use_revin = all_configs['use_revin'],
        ft_task = all_configs.get('ft_task', 'sleep_staging'),
        patch_size=all_configs.get('patch_size', 10),
        from_seq_pretrain=all_configs.get('from_seq_pretrain', False),
        full_night=all_configs.get('full_night', False),
        cfg=all_configs 
        )
    
    print("Model Loaded")        

    if all_configs['pretrain_checkpoint']:
        ckpt_file = os.path.join('experiments', all_configs['pretrain_checkpoint'])
        model.load_pretrained(ckpt_file)    
    
    # If a pure pretraining checkpoint is specified (not a lightning finetune checkpoint),
    # keep the existing `load_pretrained` behavior so encoder weights are loaded.
    if all_configs.get('pretrain_checkpoint') and not all_configs.get('previous_checkpoint') and not all_configs.get('test_checkpoint'):
        ckpt_file = os.path.join('experiments', all_configs['pretrain_checkpoint'])
        try:
            model.load_pretrained(ckpt_file)
        except Exception as e:
            print(f"Warning: failed to load pretrained weights via load_pretrained(): {e}")

    from pytorch_lightning.callbacks import DeviceStatsMonitor

    monitor = DeviceStatsMonitor()    

    # Build a trainer for evaluation only
    trainer = pl.Trainer(
        devices=all_configs['gpus'],
        accelerator='gpu',
        logger=None,
        precision='bf16-mixed',
        callbacks=[monitor],
    )

    # Run test only (do not run training/finetuning)
    # print(f"Running test with checkpoint: {ckpt_path}")    
    import time
    start = time.time()
    trainer.test(model=model, datamodule=data_module)
    end = time.time()
    print(f"Elapsed time for {len(data_module.test_loader.dataset)}:{end - start}")


if __name__ == '__main__':
    # Add optional CLI arg for specifying checkpoint path
    if '--ckpt' not in sys.argv:
        # no-op; argparse will ignore unknown args here because we've already parsed
        pass
    main()

	from collections import defaultdict
	import pytorch_lightning as pl
	from pytorch_lightning.loggers import WandbLogger
	from torchmetrics import ConfusionMatrix, CohenKappa
	from torchmetrics.classification import Accuracy, Precision, Recall, F1Score
	import torch
	import torch.nn as nn
	import logging
	import sys



	import os

	from pytorch_lightning.callbacks import ModelCheckpoint
	import numpy as np

	import data.data_loader as data_loader
	import yaml

	import models
	from models.utils import cox_ph_loss, CIndex, AUROC
	from torchmetrics import MetricCollection

	import random
	import importlib
	import argparse
	from configs.global_cfg import TARGET_RESAMPLE
	from configs.global_cfg import TARGET_RESAMPLE, DATASET_PATH, CLASSIFICATION_TASKS, REGRESSION_TASKS
	import pickle
	from einops import rearrange
	import pandas as pd

	torch.manual_seed(42)
	torch.cuda.manual_seed(42)
	torch.cuda.manual_seed_all(42)
	torch.backends.cudnn.deterministic = True
	torch.backends.cudnn.benchmark = False
	np.random.seed(42)
	random.seed(42)

	# In H100 GPUs
	torch.backends.cuda.matmul.allow_tf32 = False
	torch.backends.cudnn.allow_tf32 = False


	BASE_DIR = os.path.dirname(os.path.abspath(__file__))

	class EvalDataModule(pl.LightningDataModule):
	def __init__(self,
	data_root,
	batch_size_v,
	transforms_v,
	num_workers,
	proportion_v,
	num_classes=5,
	splits_set=None,
	modalities=["EEG"],
	sig_len_target=30,
	ft_task='sleep_staging',
	full_night=False,
	seq_truncate=False
	):
	super().__init__()

	self.data_root = data_root
	self.transforms_v = transforms_v
	self.batch_size_v = batch_size_v
	self.num_workers = num_workers
	self.proportion_v = proportion_v
	self.num_classes = num_classes
	self.splits_set = splits_set
	self.modalities = modalities
	self.sig_len_target = sig_len_target
	self.ft_task = ft_task
	self.full_night = full_night
	self.seq_truncate = seq_truncate


	def setup(self, stage=None):

	self.test_data_list = self.splits_set['test']


	self.test_loader = data_loader.build_loader(
	self.data_root, self.test_data_list,
	batch_size=self.batch_size_v, shuffle=False,
	transforms=self.transforms_v, sampler=None,
	num_workers=self.num_workers, proportion=self.proportion_v,
	split='test',
	num_classes=self.num_classes, modalities=self.modalities,
	sig_len_target=self.sig_len_target,
	ft_task=self.ft_task, full_night=self.full_night, seq_truncate=self.seq_truncate)

	print("Test loader ready")
	print("Length of test_loader: ", len(self.test_loader))


	def test_dataloader(self):
	return self.test_loader

	from finetune import ConvMAE1d_Finetune

	class ConvMAE1d_Evaluate(ConvMAE1d_Finetune):
	def __init__(self, args, *kwargs):
	super().__init__(args, *kwargs)

	self.save_results = self.cfg.get('save_results', False)
	if self.save_results:
	self.save_results_dir = "ft_results"
	os.makedirs(self.save_results_dir, exist_ok=True)
	# self.preds = defaultdict(list)
	# self.gt = defaultdict(list)
	self.results = defaultdict(dict) # Initialize nested dict for each PID

	self.save_emb = self.cfg.get('save_emb', False)
	if self.save_emb:
	self.save_emb_dir = "ft_results/embeddings"
	os.makedirs(self.save_emb_dir, exist_ok=True)
	self.embeddings = defaultdict(dict)

	self._peak_gpu_mem_bytes = None

	def on_test_epoch_start(self):
	if torch.cuda.is_available():
	torch.cuda.synchronize()
	torch.cuda.reset_peak_memory_stats()
	self._peak_gpu_mem_bytes = None


	def forward_epoch_level(self, batch, mode='test'):
	inputs, labels, _, _, pids = batch

	# signals = inputs["EEG"].to(self.device, non_blocking=True)
	for key in inputs.keys():
	inputs[key] = inputs[key].to(self.device, non_blocking=True)

	pred = self(inputs)
	if self.ft_task == 'apnea_seg':
	pred = rearrange(pred, 'n l c -> n c l')

	loss = self.criterion(pred, labels)
	pred_cls = pred.argmax(dim=1)

	if self.save_results:
	pred_cls = pred_cls.detach().cpu().numpy()
	labels = labels.detach().cpu().numpy()
	for idx, pid in enumerate(pids):
	if pid not in self.results:
	self.results[pid] = {'pred': [], 'gt': [],}
	self.results[pid]['pred'].append(pred_cls[idx])
	self.results[pid]['gt'].append(labels[idx])

	self.log_metrics(loss, pred_cls, labels, mode=mode)


	return loss

	def forward_full_night(self, batch, mode='test'):
	inputs, labels, mask, demo, pids = batch

	if self.ft_task in ['age','ahi']:
	labels = demo[self.ft_task].to(self.device, dtype=torch.float32, non_blocking=True)
	elif self.ft_task in ['sex']:
	labels = demo[self.ft_task].to(self.device, dtype=torch.long, non_blocking=True) - 1

	for key in inputs.keys():
	inputs[key] = inputs[key].to(self.device, non_blocking=True)

	if (mask is not None) and (mask.numel() > 0):
	mask = mask.to(self.device, non_blocking=True)

	if self.save_emb:
	self.model.eval()
	with torch.no_grad():
	(feats_macro_f, feats_macro_r), feats_micro = self.model.forward_features(inputs, mask=mask)
	from models.utils import flip_padded_sequences
	feats_macro_r = flip_padded_sequences(feats_macro_r, mask)
	feats_macro = torch.cat([feats_macro_f, feats_macro_r], dim=-1)

	for idx, pid in enumerate(pids):
	self.embeddings[pid] = {
	'macro': feats_macro[idx].detach().cpu().numpy(),
	'micro': feats_micro[idx].detach().cpu().numpy(),
	'mask': mask[idx].detach().cpu().numpy() if mask is not None and mask.numel() > 0 else None,
	'age': demo['age'][idx].detach().cpu().numpy(),
	'sex': demo['sex'][idx].detach().cpu().numpy(),
	'bmi': demo['bmi'][idx].detach().cpu().numpy(),
	'ahi': demo['ahi'][idx].detach().cpu().numpy()
	}
	return # Skip further processing if only saving embeddings

	pred = self(inputs, mask=mask)

	if self.ft_task == 'sleep_staging' and (mask is not None) and (mask.numel() > 0):
	pred = pred.permute(0,2,1)[mask]
	labels = labels[mask]

	if self.ft_task in REGRESSION_TASKS:
	pred = pred.squeeze()
	loss = self.criterion(pred, labels)

	if self.ft_task in CLASSIFICATION_TASKS:
	pred = pred.argmax(dim=1)

	if self.save_results:
	self.update_results(pred, labels, pids, demo)
	self.log_metrics(loss, pred, labels, mode=mode)

	def forward_diagnosis(self, batch, mode='train'):
	inputs, is_event, event_times, mask, demo, pids = batch

	# signals = inputs["EEG"].to(self.device, non_blocking=True)
	for key in inputs.keys():
	inputs[key] = inputs[key].to(self.device, non_blocking=True)

	for key in demo.keys():
	demo[key] = demo[key].to(self.device, non_blocking=True)

	if (mask is not None) and (mask.numel() > 0):
	mask = mask.to(self.device, non_blocking=True)

	pred = self(inputs, mask=mask, demo=demo)

	loss = self.criterion(pred, event_times, is_event)
	if mode == 'train':
	self.log('train_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True)
	self.train_cindex.update(pred.squeeze(), event_times, is_event)
	self.train_auroc.update(pred.squeeze(), event_times, is_event)
	elif mode == 'val':
	self.log('val_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True, batch_size=self.cfg['batch_size_v'])
	self.val_cindex.update(pred.squeeze(), event_times, is_event)
	self.val_auroc.update(pred.squeeze(), event_times, is_event)
	elif mode == 'test':
	self.log('test_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True, batch_size=self.cfg['batch_size_v'])
	self.test_cindex.update(pred.squeeze(), event_times, is_event)
	self.test_auroc.update(pred.squeeze(), event_times, is_event)

	if self.save_results:
	labels = torch.stack([event_times, is_event], dim=1)
	self.update_results(pred, labels, pids, demo)

	return loss


	def update_results(self, pred, labels, pids, demo):
	def _to_numpy(x):
	x = x.detach()
	if x.dtype in (torch.bfloat16, torch.float16, torch.float32, torch.float64):
	x = x.float()
	return x.cpu().numpy()

	pred_np = _to_numpy(pred)
	labels_np = _to_numpy(labels)
	for idx, pid in enumerate(pids):
	if pid not in self.results:
	self.results[pid] = {'pred': [],
	'gt': [],
	'age': _to_numpy(demo['age'][idx]),
	'sex': _to_numpy(demo['sex'][idx]),
	'bmi': _to_numpy(demo['bmi'][idx]),
	'ahi': _to_numpy(demo['ahi'][idx])
	}
	self.results[pid]['pred'].append(pred_np[idx])
	self.results[pid]['gt'].append(labels_np[idx])



	def on_test_epoch_end(self):
	if torch.cuda.is_available():
	torch.cuda.synchronize()
	self._peak_gpu_mem_bytes = torch.cuda.max_memory_allocated()
	peak_gb = self._peak_gpu_mem_bytes / (1024 ** 3)
	self.log('test_peak_gpu_mem_gb', peak_gb, on_step=False, on_epoch=True, prog_bar=False)
	print(f"Peak GPU memory allocated: {peak_gb:.3f} GB")

	if self.ft_task == 'diagnosis':
	cindex = self.test_cindex.compute()
	auroc = self.test_auroc.compute()
	for i in range(len(cindex)):
	self.log(f'test_cindex_class_{i}', cindex[i], on_step=False, on_epoch=True, prog_bar=False)
	self.log(f'test_auroc_class_{i}', auroc[i], on_step=False, on_epoch=True, prog_bar=False)
	self.log('test_cindex', torch.mean(cindex), on_step=False, on_epoch=True, prog_bar=True)
	self.log('test_auroc', torch.mean(auroc), on_step=False, on_epoch=True, prog_bar=True)
	print("Test C-Index:\n", cindex)
	print("Test AUROC:\n", auroc)

	self.test_cindex.reset()
	self.test_auroc.reset()
	else:
	self.log_dict(self.test_metrics.compute(), on_epoch=True, sync_dist=True)
	print("Test Confusion Matrix:\n", self.test_conf_matrix.compute())

	if self.save_results:
	# with open(os.path.join(self.save_results_dir, f'preds_{self.cfg["experiment_name"]}.pickle'), 'wb') as f:
	# pickle.dump(self.preds, f)

	# with open(os.path.join(self.save_results_dir, f'gt_{self.cfg["experiment_name"]}.pickle'), 'wb') as f:
	# pickle.dump(self.gt, f)

	with open(os.path.join(self.save_results_dir, f'results_{self.cfg["experiment_name"]}.pickle'), 'wb') as f:
	pickle.dump(self.results, f)

	if self.save_emb:
	for pid, emb in self.embeddings.items():
	with open(os.path.join(self.save_emb_dir, f'embeddings_{self.cfg["experiment_name"]}_{pid}.pickle'), 'wb') as f:
	pickle.dump(emb, f)

	def get_all_configs_from_yaml(file_path):
	with open(file_path, 'r') as file:
	configs = yaml.safe_load(file)
	return configs
	# return Config(configs)

	def main():

	# Parse command-line arguments
	parser = argparse.ArgumentParser(description="Train Perceiver Model")
	parser.add_argument('--config', type=str, required=True, help="Name of the configuration file (without .py)")
	parser.add_argument('--no_wandb', action='store_true', default=False, help="Disable Weights and Biases logging")
	args = parser.parse_args()

	all_configs = get_all_configs_from_yaml(os.path.join("configs", args.config + ".yaml"))
	print(all_configs)

	# Add transforms dynamically if needed
	# args.transforms = get_transforms()
	# print(args)

	print("Starting Experiment: ", all_configs['experiment_name'])

	config_save_path = os.path.join(all_configs['config_save_path'], all_configs['experiment_name'])

	os.makedirs(config_save_path, exist_ok=True)
	with open(os.path.join(config_save_path, 'config.yaml'), 'w') as file:
	yaml.dump(all_configs, file)

	# Set data paths
	# data_root("datasets") directory should be in the same directory as this script
	data_root = os.path.join(BASE_DIR, all_configs['data_root'])
	datasets = all_configs.get('datasets', ['phy'])


	# Load data
	print("Loading Data")
	splits_set = dict()
	splits_set['train'] = []
	splits_set['val'] = []
	splits_set['test'] = []

	for ds in datasets:
	# Use Physionet 2018 dataset
	if ds == 'kiss':

	sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects

	sample_dir = os.path.join(data_root, DATASET_PATH[ds], 'test', sample_modality)
	chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
	info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
	with open(info_path, 'rb') as f:
	sig_info = pickle.load(f)

	test_split = sorted([sub for sub in sig_info['sig_info'].keys()])

	if all_configs.get('save_emb',False):
	test_split = test_split[:500] # For saving embeddings, limit to first 500 subjects only to save space

	splits_set['test'].extend([('kiss',f) for f in test_split])

	elif ds == 'shhs1':

	sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects
	sample_dir = os.path.join(data_root, DATASET_PATH[ds], 'test', sample_modality)
	chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
	info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
	with open(info_path, 'rb') as f:
	sig_info = pickle.load(f)
	test_subjects = sorted([sub for sub in sig_info['sig_info'].keys()])

	def get_shhs_lists(shhs_split_file):
	ids = pd.read_csv(os.path.join("data","shhs_meta_data", shhs_split_file), header=None, names=['a'])
	ids = ids['a'].values.tolist()
	list_files = ["shhs1-20"+id[11:15] for id in ids]
	return list_files

	list_files_test = get_shhs_lists("test_list.txt")
	test_subjects = set(list_files_test) & set(test_subjects)
	test_subjects = sorted(list(test_subjects))

	if all_configs.get('save_emb',False):
	test_subjects = test_subjects[:500] # For saving embeddings, limit to first 500 subjects only to save space
	splits_set['test'].extend([('shhs1',f) for f in test_subjects])


	elif 'isruc' in ds:

	sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects

	sample_dir = os.path.join(data_root, DATASET_PATH[ds], 'test', sample_modality)
	info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
	with open(info_path, 'rb') as f:
	sig_info = pickle.load(f)
	test_subjects = sorted([sub for sub in sig_info['sig_info'].keys()])

	splits_set['test'].extend([(ds,f) for f in test_subjects])

	print("Number of test subjects: ", len(splits_set['test']))

	if all_configs['pretrain_checkpoint']:
	pretrain_exp_name = all_configs['pretrain_checkpoint'].split('-')[0]
	pretrain_config = get_all_configs_from_yaml(os.path.join("experiments", pretrain_exp_name, "config.yaml"))
	use_shared = pretrain_config.get('use_shared', all_configs.get('use_shared', True))
	use_moe = pretrain_config.get('use_moe', all_configs.get('use_moe', True))
	else:
	use_shared = all_configs.get('use_shared', True)
	use_moe = all_configs.get('use_moe', True)


	sig_len_target = all_configs.get('sig_len_target', 60)
	data_module = EvalDataModule(
	data_root = all_configs['data_root'],
	transforms_v = None,
	batch_size_v = all_configs['batch_size_v'],
	num_workers = all_configs['num_workers'],
	proportion_v = all_configs['proportion_v'],
	num_classes = all_configs['num_classes'],
	splits_set = splits_set,
	modalities = all_configs['modalities'],
	sig_len_target=sig_len_target,
	ft_task=all_configs.get('ft_task', 'sleep_staging'),
	full_night=all_configs.get('full_night', False),
	seq_truncate=all_configs.get('seq_truncate', False)
	)
	data_module.setup() # Setup the data module to prepare the datasets and dataloaders


	print("Model Loading ...")
	model = ConvMAE1d_Evaluate(
	model_name=all_configs['model_name'],
	num_classes = all_configs['num_classes'],
	modalities = all_configs['modalities'],
	num_experts = all_configs['num_experts'],
	num_experts_per = all_configs['num_experts_per'],
	loss = all_configs['loss'],
	sig_len = sig_len_target * TARGET_RESAMPLE,
	use_shared = use_shared,
	use_moe = use_moe,
	full_finetune=all_configs.get('full_finetune', False), # Default to False if not specified
	use_revin = all_configs['use_revin'],
	ft_task = all_configs.get('ft_task', 'sleep_staging'),
	patch_size=all_configs.get('patch_size', 10),
	from_seq_pretrain=all_configs.get('from_seq_pretrain', False),
	full_night=all_configs.get('full_night', False),
	cfg=all_configs
	)

	print("Model Loaded")

	if all_configs['pretrain_checkpoint']:
	ckpt_file = os.path.join('experiments', all_configs['pretrain_checkpoint'])
	model.load_pretrained(ckpt_file)

	# If a pure pretraining checkpoint is specified (not a lightning finetune checkpoint),
	# keep the existing `load_pretrained` behavior so encoder weights are loaded.
	if all_configs.get('pretrain_checkpoint') and not all_configs.get('previous_checkpoint') and not all_configs.get('test_checkpoint'):
	ckpt_file = os.path.join('experiments', all_configs['pretrain_checkpoint'])
	try:
	model.load_pretrained(ckpt_file)
	except Exception as e:
	print(f"Warning: failed to load pretrained weights via load_pretrained(): {e}")

	from pytorch_lightning.callbacks import DeviceStatsMonitor

	monitor = DeviceStatsMonitor()

	# Build a trainer for evaluation only
	trainer = pl.Trainer(
	devices=all_configs['gpus'],
	accelerator='gpu',
	logger=None,
	precision='bf16-mixed',
	callbacks=[monitor],
	)

	# Run test only (do not run training/finetuning)
	# print(f"Running test with checkpoint: {ckpt_path}")
	import time
	start = time.time()
	trainer.test(model=model, datamodule=data_module)
	end = time.time()
	print(f"Elapsed time for {len(data_module.test_loader.dataset)}:{end - start}")


	if __name__ == '__main__':
	# Add optional CLI arg for specifying checkpoint path
	if '--ckpt' not in sys.argv:
	# no-op; argparse will ignore unknown args here because we've already parsed
	pass
	main()