evaluate.py

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