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