from configs.global_cfg import DATASET_PATH import os import pickle import pandas as pd from torch.optim.lr_scheduler import _LRScheduler import math import torch from safetensors.torch import save_file from lifelines.utils import concordance_index from torchmetrics import Metric, MetricCollection from torchmetrics.utilities import dim_zero_cat def get_splits_from_file(ds, all_configs, data_root, split=None): sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects if split is None: sample_dir = os.path.join(data_root, DATASET_PATH[ds], sample_modality) else: sample_dir = os.path.join(data_root, DATASET_PATH[ds], split, 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) subjects = sorted([sub for sub in sig_info['sig_info'].keys()]) return subjects def get_datasets(datasets, all_configs, data_root, splits=['train']): splits_set = dict() for split in splits: splits_set[split] = [] for ds in datasets: # Use Physionet 2018 dataset if ds == 'shhs1': 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_val = get_shhs_lists("eval_list.txt") list_files_test = get_shhs_lists("test_list.txt") if 'train' in splits: train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train') missing_cases = ['shhs1-204709','shhs1-204806'] train_subjects = set(train_subjects) - set(list_files_val) - set(list_files_test) - set(missing_cases) train_subjects = sorted(list(train_subjects)) splits_set['train'].extend([('shhs1',f) for f in train_subjects]) if 'val' in splits: val_subjects = get_splits_from_file(ds, all_configs, data_root, split='val') val_subjects = set(list_files_val) & set(val_subjects) val_subjects = sorted(list(val_subjects)) splits_set['val'].extend([('shhs1',f) for f in val_subjects]) if 'test' in splits: test_subjects = get_splits_from_file(ds, all_configs, data_root, split='test') test_subjects = set(list_files_test) & set(test_subjects) test_subjects = sorted(list(test_subjects)) splits_set['test'].extend([('shhs1',f) for f in test_subjects]) elif ds == 'kiss' or 'cfs' in ds or 'isruc' in ds or 'sof' in ds: ds_main = 'kiss' if ds in 'kiss' else 'cfs' if 'cfs' in ds else 'isruc' if 'isruc' in ds else 'sof' if 'train' in splits: train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train') splits_set['train'].extend([(ds_main,f) for f in train_subjects]) if 'val' in splits: val_subjects = get_splits_from_file(ds, all_configs, data_root, split='val') splits_set['val'].extend([(ds_main,f) for f in val_subjects]) if 'test' in splits: test_subjects = get_splits_from_file(ds, all_configs, data_root, split='test') splits_set['test'].extend([(ds_main,f) for f in test_subjects]) elif ds in ['phy-tr', 'phy-te', 'mros1-val', 'mros1-te','mros2-val','mros2-te','mesa-val','mesa-te']: # sample_modality = all_configs.get('modalities', ['EEG'])[0] # sample_dir = os.path.join(data_root, DATASET_PATH[ds], 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) # all_subjects = sorted([sub for sub in sig_info['sig_info'].keys()]) if 'train' in splits: train_subjects = get_splits_from_file(ds, all_configs, data_root, split=None) splits_set['train'].extend([(ds, f) for f in train_subjects]) elif ds in ['shhs2', 'kiss','kvss','mros1','mros2','mesa']: if 'train' in splits: train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train') splits_set['train'].extend([(ds, f) for f in train_subjects]) # splits_set['train'].extend([(ds, f) for f in all_subjects]) return splits_set def export_inference_weights(ckpt_path, model, out_path, metadata=None): """ Strip a Lightning checkpoint down to the weights `model` loads and save them as safetensors (no optimizer / scheduler / callback state, no pickle). `model` is the inference backbone, e.g. `inference.build_micro_model()` or `inference.build_full_night_model()`. Keys are saved without the `model.` prefix; `pos_embed` is skipped since it is a fixed sin-cos table rebuilt at model construction. Keys absent from the checkpoint (e.g. `head` for a pretraining checkpoint) are not saved. from inference import build_micro_model export_inference_weights("experiments/micro.ckpt", build_micro_model(), "checkpoints/micro.safetensors", {"task": "micro"}) """ ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) src = ckpt["state_dict"] if "state_dict" in ckpt else ckpt weights = {} for k in model.state_dict(): if "pos_embed" in k: continue if f"model.{k}" in src: weights[k] = src[f"model.{k}"].detach().clone().contiguous() elif k in src: weights[k] = src[k].detach().clone().contiguous() meta = {"source_epoch": str(ckpt.get("epoch")), "source_global_step": str(ckpt.get("global_step"))} meta.update({k: str(v) for k, v in (metadata or {}).items()}) save_file(weights, out_path, metadata=meta) return weights def get_layer_decay_param_groups(model, base_lr, weight_decay, layer_decay): """ Create parameter groups with layer-wise learning rate decay. Args: model: The model whose parameters are grouped. base_lr: Base learning rate. weight_decay: Weight decay for regularization. layer_decay: Layer decay factor (e.g., 0.75). Returns: A list of parameter groups for the optimizer. """ param_groups = [] layers = list(model.children()) num_layers = len(layers) # We use named_parameters to filter by name (e.g., A_log, dt, bias) # But we need to map them back to their layer index for LR decay. for i, layer in enumerate(layers): layer_scale = layer_decay ** (num_layers - i - 1) curr_lr = base_lr * layer_scale # We split parameters in each layer into 'decay' and 'no_decay' params_to_decay = [] params_no_decay = [] for name, param in layer.named_parameters(): if not param.requires_grad: continue # Check if it's a Mamba-specific parameter or a standard "no-decay" candidate # Common Mamba names: A_log, D, dt_proj, x_proj, in_proj, conv1d # Standard no-decay: bias, norm weights is_ssm_param = any(ssm_name in name for ssm_name in ['A_log', 'D', 'dt', 'B', 'C']) is_bias_or_norm = name.endswith("bias") or "norm" in name if is_ssm_param or is_bias_or_norm: params_no_decay.append(param) else: params_to_decay.append(param) # Add groups for this specific layer if params_to_decay: param_groups.append({ "params": params_to_decay, "lr": curr_lr, "weight_decay": weight_decay, }) if params_no_decay: param_groups.append({ "params": params_no_decay, "lr": curr_lr, "weight_decay": 0.0, }) return param_groups class WarmupCosineAnnealingLR(_LRScheduler): def __init__(self, optimizer, warmup_epochs, max_epochs, eta_min=1e-8, last_epoch=-1): """ Step-based Warmup Cosine Annealing Learning Rate Scheduler. Args: optimizer: The optimizer for which to schedule the learning rate. warmup_epochs: Number of steps for the warmup phase. max_epochs: Total number of steps for the scheduler. eta_min: Minimum learning rate at the end of the schedule. last_epoch: The index of the last step (default: -1). """ self.warmup_epochs = warmup_epochs self.max_epochs = max_epochs self.eta_min = eta_min super(WarmupCosineAnnealingLR, self).__init__(optimizer, last_epoch) def get_lr(self): if self.last_epoch < self.warmup_epochs: # Linear warmup return [base_lr * (self.last_epoch + 1) / self.warmup_epochs for base_lr in self.base_lrs] else: # Cosine annealing progress = (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs) return [self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * progress)) / 2 for base_lr in self.base_lrs] def cox_ph_loss(hazards, event_times, is_event): # Sort event times and get corresponding indices for sorting other tensors event_times, sorted_idx = event_times.sort(dim=0, descending=True) # Make gathering robust: if sorted_idx is 1D and hazards/is_event are 2D, # use advanced indexing; otherwise fall back to gather which requires # the index tensor to have the same number of dimensions as input. if sorted_idx.dim() == 1 and hazards.dim() == 2: hazards = hazards[sorted_idx] is_event = is_event[sorted_idx] else: hazards = hazards.gather(0, sorted_idx) is_event = is_event.gather(0, sorted_idx) log_cumulative_hazard = torch.logcumsumexp(hazards.float(), dim=0) # Calculate losses for all labels simultaneously losses = (hazards - log_cumulative_hazard) * is_event losses = -losses # Negative for maximization # Average loss per label label_loss = losses.sum(dim=0) / (is_event.sum(dim=0) + 1e-9) # Avoid division by zero # Average across labels total_loss = label_loss.mean() return total_loss def c_index_torch(scores, event_times, is_event): return float(concordance_index(event_times, -scores, is_event)) class CIndex(Metric): def __init__(self, num_classes=6, **kwargs): super().__init__(**kwargs) self.add_state("scores", default=[], dist_reduce_fx="cat") self.add_state("event_time", default=[], dist_reduce_fx="cat") self.add_state("is_event", default=[], dist_reduce_fx="cat") self.num_classes = num_classes def update(self, score, t, e) -> None: self.scores.append(score) self.event_time.append(t) self.is_event.append(e) def compute(self): # parse inputs # Convert to float32 before numpy() since numpy doesn't support BFloat16 scores = dim_zero_cat(self.scores).detach().cpu().float().numpy() event_time = dim_zero_cat(self.event_time).detach().cpu().float().numpy() is_event = dim_zero_cat(self.is_event).detach().cpu().float().numpy() from lifelines.utils import concordance_index cindex = [] for i in range(self.num_classes): try: c = concordance_index(event_time[:,i], -scores[:,i], is_event[:,i]) cindex.append(c) except (ZeroDivisionError, ValueError): # Not enough admissible pairs (common during sanity check or with small batches) # Return 0.5 (random predictor) as default cindex.append(0.5) return torch.tensor(cindex) class AUROC(Metric): def __init__(self, num_classes=6, **kwargs): super().__init__(**kwargs) self.add_state("scores", default=[], dist_reduce_fx="cat") self.add_state("event_time", default=[], dist_reduce_fx="cat") self.add_state("is_event", default=[], dist_reduce_fx="cat") self.num_classes = num_classes self.measure_date = 365 * 6 # at 6 years def update(self, score, t, e) -> None: self.scores.append(score) self.event_time.append(t) self.is_event.append(e) def compute(self): # parse inputs # Convert to float32 before numpy() since numpy doesn't support BFloat16 scores = dim_zero_cat(self.scores).detach().cpu().float().numpy() event_time = dim_zero_cat(self.event_time).detach().cpu().float().numpy() is_event = dim_zero_cat(self.is_event).detach().cpu().float().numpy() from sklearn.metrics import roc_auc_score, roc_curve # auroc = [] # # For each disease, calculate AUROC # for i in range(self.num_classes): # y_true = (event_time[:,i] <= self.measure_date).astype(int) # y_scores = scores[:,i] # try: # c = roc_auc_score(y_true, y_scores) # auroc.append(c) # except (ZeroDivisionError, ValueError): # # Not enough admissible pairs (common during sanity check or with small batches) # # Return 0.5 (random predictor) as default # auroc.append(0.5) # return torch.tensor(auroc) auroc = [] # For each disease, calculate AUROC at 6-year follow-up (time-dependent) for i in range(self.num_classes): # Filter out zero times (invalid entries) valid_mask = event_time[:, i] > 0 if valid_mask.sum() < 2: auroc.append(0.5) continue valid_scores = scores[valid_mask, i] valid_times = event_time[valid_mask, i] valid_events = is_event[valid_mask, i] # Time-dependent AUROC at measure_date: # - Samples with event before measure_date: y=1 # - Samples with no event and observed until after measure_date: y=0 # - Samples with censoring before measure_date: excluded (potentially informative censoring) # Consider only samples observed until measure_date observed_mask = valid_times >= self.measure_date # Or event occurred before measure_date event_before_mask = valid_events == 1 # Use all samples with valid time info for AUROC # y_true: 1 if event before measure_date, 0 otherwise y_true = (valid_events == 1) & (valid_times <= self.measure_date) y_true = y_true.astype(int) y_scores = valid_scores try: # Check if there's variation in labels if len(np.unique(y_true)) < 2: auroc.append(0.5) else: c = roc_auc_score(y_true, y_scores) auroc.append(c) except (ZeroDivisionError, ValueError) as e: # Not enough admissible pairs auroc.append(0.5) return torch.tensor(auroc)