| from collections import defaultdict |
| import pytorch_lightning as pl |
| from pytorch_lightning.loggers import WandbLogger |
| from torchmetrics import ConfusionMatrix, CohenKappa, MeanAbsoluteError, MeanSquaredError |
| from torchmetrics.classification import Accuracy, Precision, Recall, F1Score |
| import torch |
| import torch.nn as nn |
|
|
| from torch.optim.lr_scheduler import CosineAnnealingLR |
| import os |
|
|
| from pytorch_lightning.callbacks import ModelCheckpoint, Callback |
| import numpy as np |
|
|
| import data.data_loader as data_loader |
| import yaml |
|
|
| import models |
| from torchmetrics import MetricCollection |
|
|
| import random |
| import argparse |
| from configs.global_cfg import TARGET_RESAMPLE, DATASET_PATH, CLASSIFICATION_TASKS, REGRESSION_TASKS |
| import pickle |
| from einops import rearrange |
| import pandas as pd |
| from train_utils import get_datasets, cox_ph_loss, CIndex, AUROC, get_layer_decay_param_groups, WarmupCosineAnnealingLR |
|
|
| 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) |
|
|
| |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
|
|
|
|
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) |
|
|
| class FinetuneDataModule(pl.LightningDataModule): |
| def __init__(self, |
| data_root, |
| batch_size_t, |
| batch_size_v, |
| transforms_t, |
| transforms_v, |
| num_workers, |
| proportion_t, |
| 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_t = transforms_t |
| self.transforms_v = transforms_v |
| self.batch_size_t = batch_size_t |
| self.batch_size_v = batch_size_v |
| self.num_workers = num_workers |
| self.proportion_t = proportion_t |
| 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.train_data_list = self.splits_set['train'] |
| self.val_data_list = self.splits_set['val'] |
| self.test_data_list = self.splits_set['test'] |
| |
| print("_"*50) |
| print("Number of training subjects: ", len(self.train_data_list)) |
| |
| self.train_loader = data_loader.build_loader( |
| self.data_root, self.train_data_list, |
| batch_size=self.batch_size_t, shuffle=True, |
| transforms=self.transforms_t, sampler=None, |
| num_workers=self.num_workers, proportion=self.proportion_t, |
| split='train', |
| 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("Train loader ready") |
| print("Length of train_loader: ", len(self.train_loader)) |
|
|
| self.val_loader = data_loader.build_loader( |
| self.data_root, self.val_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='val', |
| 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("Val loader ready") |
| print("Length of val_loader: ", len(self.val_loader)) |
|
|
| 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 train_dataloader(self): |
| return self.train_loader |
|
|
| def val_dataloader(self): |
| return self.val_loader |
|
|
| def test_dataloader(self): |
| return self.test_loader |
|
|
| |
|
|
| class ConvMAE1d_Finetune(pl.LightningModule): |
| def __init__(self, |
| model_name='classification_multi', |
| num_classes=5, |
| modalities=["EEG"], |
| num_experts=4, |
| num_experts_per=2, |
| loss='ce', |
| class_weights=np.ones(5), |
| sig_len=6000, |
| use_shared=True, |
| use_moe=True, |
| use_revin=False, |
| full_finetune=False, |
| optimizer_name='AdamW', |
| lr=0.0005, |
| weight_decay=1e-6, |
| scheduler_name='CosineAnnealingLR', |
| num_epochs=10, |
| ft_task='sleep_staging', |
| patch_size=10, |
| from_seq_pretrain=False, |
| full_night=False, |
| cfg=None |
| ): |
| super().__init__() |
| self.cfg = cfg |
| self.save_hyperparameters() |
| single_channel = True |
| self.num_classes = num_classes |
| self.modalities = modalities |
| self.num_experts = num_experts |
| self.num_experts_per = num_experts_per |
| self.sig_len = sig_len |
| self.use_shared = use_shared |
| self.use_moe = use_moe |
| self.full_finetune = full_finetune |
| self.use_revin = use_revin |
| self.ft_task = ft_task |
| self.patch_size = patch_size |
| self.from_seq_pretrain = from_seq_pretrain |
| self.full_night = full_night |
| self.f1_class_index = cfg.get('f1_class_index') if cfg is not None else None |
| if self.f1_class_index is not None and not 0 <= self.f1_class_index < num_classes: |
| raise ValueError(f"f1_class_index must be between 0 and {num_classes - 1}") |
| |
| kwargs = {} |
| if model_name == 'classification_multi_full_night': |
| kwargs = {'use_micro': cfg.get('use_micro', False), |
| 'use_lstm': cfg.get('use_lstm', False), |
| 'use_transformer': cfg.get('use_transformer', False), |
| 'bidi': False if cfg.get('use_transformer', False) else cfg.get('bidi', True)} |
|
|
| self.model = models.__dict__[model_name](modalities=self.modalities, |
| num_experts=self.num_experts, |
| num_experts_per=self.num_experts_per, |
| sig_len=self.sig_len, |
| use_shared=self.use_shared, |
| use_moe=self.use_moe, |
| use_revin=self.use_revin, |
| ft_task=ft_task, |
| patch_size=self.patch_size, |
| num_classes=self.num_classes, |
| **kwargs) |
| |
| if self.ft_task == 'diagnosis': |
| assert loss == 'coxph', "Diagnosis task requires Cox Proportional Hazards loss." |
| self.criterion = cox_ph_loss |
| self.train_cindex = CIndex() |
| self.val_cindex = CIndex() |
| self.test_cindex = CIndex() |
| self.train_auroc = AUROC(num_classes=num_classes) |
| self.val_auroc = AUROC(num_classes=num_classes) |
| self.test_auroc = AUROC(num_classes=num_classes) |
| |
| |
| elif self.ft_task in CLASSIFICATION_TASKS: |
| if loss == 'ce': |
| self.criterion = nn.CrossEntropyLoss() |
| |
| elif loss == 'wce': |
| class_weights = torch.tensor(class_weights, dtype=torch.float32) |
| self.criterion = nn.CrossEntropyLoss(weight=class_weights) |
| elif loss == 'mse': |
| self.criterion = nn.MSELoss() |
| |
| |
| |
| self.train_metrics = MetricCollection( |
| { |
| "acc": Accuracy(task="multiclass", num_classes=num_classes), |
| "f1": F1Score(task="multiclass", num_classes=num_classes), |
| }, |
| prefix="train_", |
| ) |
|
|
| self.val_metrics = MetricCollection( |
| { |
| "acc": Accuracy(task="multiclass", num_classes=num_classes), |
| "f1": F1Score(task="multiclass", num_classes=num_classes, average='macro'), |
| "kappa": CohenKappa(task="multiclass", num_classes=num_classes), |
| "precision": Precision(task="multiclass", num_classes=num_classes, average='macro'), |
| "recall": Recall(task="multiclass", num_classes=num_classes, average='macro'), |
| }, |
| prefix="val_", |
| ) |
| if self.f1_class_index is not None: |
| self.val_f1_class = F1Score( |
| task="multiclass", num_classes=num_classes, average='none' |
| ) |
|
|
| self.test_metrics = MetricCollection( |
| { |
| "acc": Accuracy(task="multiclass", num_classes=num_classes), |
| "f1": F1Score(task="multiclass", num_classes=num_classes, average='macro'), |
| "kappa": CohenKappa(task="multiclass", num_classes=num_classes), |
| "precision": Precision(task="multiclass", num_classes=num_classes, average='macro'), |
| "recall": Recall(task="multiclass", num_classes=num_classes, average='macro'), |
| }, |
| prefix="test_", |
| ) |
| if self.f1_class_index is not None: |
| self.test_f1_class = F1Score( |
| task="multiclass", num_classes=num_classes, average='none' |
| ) |
| self.val_conf_matrix = ConfusionMatrix(task="multiclass", num_classes=num_classes) |
| self.test_conf_matrix = ConfusionMatrix(task="multiclass", num_classes=num_classes) |
| |
| elif self.ft_task in REGRESSION_TASKS: |
| if loss == 'mse': |
| self.criterion = nn.MSELoss() |
| elif loss == 'mae': |
| self.criterion = nn.L1Loss() |
| |
| |
| |
| self.train_metrics = MetricCollection( |
| { |
| "mae": MeanAbsoluteError(), |
| "mse": MeanSquaredError() |
| }, |
| prefix="train_", |
| ) |
|
|
| self.val_metrics = MetricCollection( |
| { |
| "mae": MeanAbsoluteError(), |
| "mse": MeanSquaredError() |
| }, |
| prefix="val_", |
| ) |
|
|
| self.test_metrics = MetricCollection( |
| { |
| "mae": MeanAbsoluteError(), |
| "mse": MeanSquaredError() |
| }, |
| prefix="test_", |
| ) |
| |
| |
| def load_pretrained(self, ckpt): |
| |
| checkpoint = torch.load(ckpt, map_location=self.device, weights_only=False) |
| new_state_dict = {} |
| for k in self.model.state_dict().keys(): |
| if 'pos_embed' in k: |
| print("Skipping key (pos_embed): ", k) |
| continue |
| if k in checkpoint['state_dict']: |
| |
| new_state_dict[k] = checkpoint['state_dict'][k] |
| else: |
| |
| if f"model.{k}" in checkpoint['state_dict']: |
| |
| new_state_dict[k] = checkpoint['state_dict'][f"model.{k}"] |
| else: |
| print("Missing key (Only in the current model): ", k) |
| |
| |
| for k in checkpoint['state_dict'].keys(): |
| if not (k.startswith('model.') and k[6:] in self.model.state_dict()) and not k in self.model.state_dict(): |
| print("Missing key (Only in the Checkpoint): ", k) |
| |
| msg = self.model.load_state_dict(new_state_dict, strict=False) |
| print(f"Loaded pretrained model from {ckpt}") |
| print(msg) |
| |
| def on_train_epoch_start(self): |
| if self.ft_task != 'diagnosis': |
| self.train_metrics.reset() |
|
|
| def on_validation_epoch_end(self): |
| if self.ft_task == 'diagnosis': |
| cindex = self.val_cindex.compute() |
| auroc = self.val_auroc.compute() |
| self.log('val_cindex', torch.mean(cindex), on_step=False, on_epoch=True, prog_bar=True) |
| self.log('val_auroc', torch.mean(auroc), on_step=False, on_epoch=True, prog_bar=True) |
| print("Validation C-Index:\n", cindex) |
| print("Validation AUROC:\n", auroc) |
| |
| self.val_cindex.reset() |
| self.val_auroc.reset() |
| else: |
| val_metrics = self.val_metrics.compute() |
| self.log_dict(val_metrics, on_epoch=True, sync_dist=True) |
| if self.ft_task in CLASSIFICATION_TASKS and self.f1_class_index is not None: |
| self.log( |
| f'val_f1_class_{self.f1_class_index}', |
| self.val_f1_class.compute()[self.f1_class_index], |
| on_step=False, on_epoch=True, sync_dist=True, prog_bar=True |
| ) |
| self.val_metrics.reset() |
| if self.ft_task in CLASSIFICATION_TASKS: |
| if self.f1_class_index is not None: |
| self.val_f1_class.reset() |
| print("Validation Confusion Matrix:\n", self.val_conf_matrix.compute()) |
|
|
| def on_test_epoch_end(self): |
| 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: |
| test_metrics = self.test_metrics.compute() |
| self.log_dict(test_metrics, on_epoch=True, sync_dist=True) |
| if self.ft_task in CLASSIFICATION_TASKS and self.f1_class_index is not None: |
| self.log( |
| f'test_f1_class_{self.f1_class_index}', |
| self.test_f1_class.compute()[self.f1_class_index], |
| on_step=False, on_epoch=True, sync_dist=True |
| ) |
| self.test_f1_class.reset() |
| print("Test Confusion Matrix:\n", self.test_conf_matrix.compute()) |
|
|
| |
|
|
| def forward(self, s, mask=None, demo=None): |
| if self.full_night: |
| return self.model(s, mask=mask, demo=demo) |
| else: |
| return self.model(s) |
| |
|
|
| def train(self, mode: bool = True): |
| """ |
| Override the default .train() method to keep the encoder in eval mode. |
| """ |
| |
| super().train(mode) |
| |
| |
| if not self.full_finetune: |
| self.model.encoder.eval() |
| self.model.fuser.eval() |
| if self.from_seq_pretrain and not self.cfg.get('use_micro', False): |
| self.model.proj.eval() |
| self.model.sequence_encoder.eval() |
|
|
| def log_metrics(self, loss, pred_cls, labels, mode='train'): |
| |
| if mode == 'train': |
| batch_value = self.train_metrics(pred_cls, labels) |
| self.log_dict(batch_value, on_step=True, on_epoch=True, sync_dist=True) |
| self.log('learning_rate', self.trainer.optimizers[0].param_groups[0]['lr'], on_step=True, on_epoch=True, sync_dist=True) |
| self.log('train_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True) |
| |
| elif mode == 'val': |
| self.val_metrics.update(pred_cls, labels) |
| if self.ft_task in CLASSIFICATION_TASKS and self.f1_class_index is not None: |
| self.val_f1_class.update(pred_cls, labels) |
| self.val_conf_matrix.update(pred_cls, labels) |
| self.log('val_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True, batch_size=self.cfg['batch_size_v']) |
| |
| |
| elif mode == 'test': |
| |
| self.test_metrics.update(pred_cls, labels) |
| if self.ft_task in CLASSIFICATION_TASKS and self.f1_class_index is not None: |
| self.test_f1_class.update(pred_cls, labels) |
| self.test_conf_matrix.update(pred_cls, labels) |
| self.log('test_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True, batch_size=self.cfg['batch_size_v']) |
|
|
| def forward_diagnosis(self, batch, mode='train'): |
| inputs, is_event, event_times, mask, demo, _ = batch |
|
|
| |
| 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) |
|
|
| return loss |
| |
| |
|
|
| def forward_epoch_level(self, batch, mode='train'): |
| inputs, labels, _, _, pids = batch |
|
|
| |
| 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) |
| |
|
|
| self.log_metrics(loss, pred_cls, labels, mode=mode) |
|
|
| return loss |
|
|
| def forward_full_night(self, batch, mode='train'): |
| 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) |
| |
| pred = self(inputs, mask=mask) |
| |
| if self.ft_task == 'sleep_staging' and (mask is not None) and (mask.numel() > 0): |
| |
| pred = pred[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) |
| |
| self.log_metrics(loss, pred, labels, mode=mode) |
|
|
| return loss |
| |
| def training_step(self, batch, batch_idx): |
| if self.ft_task == 'diagnosis': |
| loss = self.forward_diagnosis(batch, mode='train') |
| elif not self.full_night: |
| loss = self.forward_epoch_level(batch, mode='train') |
| else: |
| loss = self.forward_full_night(batch, mode='train') |
|
|
| return loss |
|
|
| def validation_step(self, batch, batch_idx): |
| if self.ft_task == 'diagnosis': |
| loss = self.forward_diagnosis(batch, mode='val') |
| elif not self.full_night: |
| loss = self.forward_epoch_level(batch, mode='val') |
| else: |
| loss = self.forward_full_night(batch, mode='val') |
|
|
| return loss |
|
|
| def test_step(self, batch): |
| if self.ft_task == 'diagnosis': |
| loss = self.forward_diagnosis(batch, mode='test') |
| elif not self.full_night: |
| loss = self.forward_epoch_level(batch, mode='test') |
| else: |
| loss = self.forward_full_night(batch, mode='test') |
|
|
| return loss |
| |
| def configure_optimizers(self): |
| |
| base_lr = self.hparams.lr |
| weight_decay = self.hparams.weight_decay if hasattr(self.hparams, 'weight_decay') else 0.05 |
| layer_decay = self.hparams.layer_decay if hasattr(self.hparams, 'layer_decay') else 1.0 |
| if not self.full_finetune: |
| for param in self.model.encoder.parameters(): |
| param.requires_grad = False |
| if self.use_revin: |
| for param in self.model.revin.parameters(): |
| param.requires_grad = False |
| if self.use_shared: |
| for param in self.model.fuser.parameters(): |
| param.requires_grad = False |
| |
| if self.from_seq_pretrain and not self.cfg.get('use_micro', False): |
| for param in self.model.sequence_encoder.parameters(): |
| param.requires_grad = False |
| for param in self.model.proj.parameters(): |
| param.requires_grad = False |
|
|
| |
| param_groups = get_layer_decay_param_groups(self.model, base_lr, weight_decay, layer_decay) |
|
|
| optimizer = getattr(torch.optim, self.hparams.optimizer_name)(param_groups) |
| |
| if self.hparams.scheduler_name == 'CosineAnnealingLR': |
| steps_per_epoch = len(self.trainer.datamodule.train_dataloader()) |
| max_steps = int(steps_per_epoch * self.hparams.num_epochs / len(self.cfg['gpus'])) + 1 |
| return{ |
| "optimizer": optimizer, |
| "lr_scheduler": { |
| "scheduler": CosineAnnealingLR(optimizer, T_max=max_steps, eta_min=1e-8), |
| "interval": "step", |
| "frequency": 1 |
| } |
| } |
| elif self.hparams.scheduler_name == 'WarmupCosineAnnealingLR': |
| steps_per_epoch = len(self.trainer.datamodule.train_dataloader()) |
| max_steps = steps_per_epoch * self.hparams.num_epochs |
| warmup_steps = int(max_steps * 0.1) |
| |
| scheduler = WarmupCosineAnnealingLR( |
| optimizer, |
| warmup_epochs=warmup_steps, |
| max_epochs=max_steps, |
| eta_min=1e-8 |
| ) |
|
|
| return { |
| "optimizer": optimizer, |
| "lr_scheduler": { |
| "scheduler": scheduler, |
| "interval": "step", |
| "frequency": 1 |
| } |
| } |
|
|
| else: |
| return optimizer |
|
|
|
|
| def get_all_configs_from_yaml(file_path): |
| with open(file_path, 'r') as file: |
| configs = yaml.safe_load(file) |
| return configs |
| |
|
|
| def main(): |
| |
| |
| 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) |
|
|
| 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) |
|
|
| |
| data_root = os.path.join(BASE_DIR, all_configs['data_root']) |
| datasets = all_configs.get('datasets', ['phy']) |
|
|
|
|
| |
| print("Loading Data") |
| splits_set = get_datasets(datasets, all_configs, data_root, splits=['train', 'val', 'test']) |
|
|
| print("Number of train subjects: ", len(splits_set['train'])) |
| print("Number of valid subjects: ", len(splits_set['val'])) |
| print("Number of test subjects: ", len(splits_set['test'])) |
|
|
| if all_configs['pretrain_checkpoint']: |
| pretrain_exp_name = all_configs['pretrain_checkpoint'].split('.')[0].split('-')[0] |
| pretrain_config_path = os.path.join("experiments", pretrain_exp_name, "config.yaml") |
| if os.path.exists(pretrain_config_path): |
| pretrain_config = get_all_configs_from_yaml(os.path.join("experiments", pretrain_exp_name, "config.yaml")) |
| else: |
| pretrain_config = {} |
| print(f"Warning: Pretrain config file {pretrain_config_path} not found. Using default values for use_shared and use_moe.") |
| 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 = FinetuneDataModule( |
| data_root = all_configs['data_root'], |
| transforms_t = None, |
| transforms_v = None, |
| batch_size_t = all_configs['batch_size_t'], |
| batch_size_v = all_configs['batch_size_v'], |
| num_workers = all_configs['num_workers'], |
| proportion_t = all_configs['proportion_t'], |
| 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() |
|
|
| if all_configs['loss'] == 'wce': |
| train_loader = data_module.train_dataloader() |
| labels = train_loader.dataset.get_all_labels() |
| class_count = np.zeros(all_configs['num_classes'], dtype=np.int64) |
| if all_configs.get('full_night', False) or all_configs.get('ft_task', 'sleep_staging') == 'apnea_seg': |
| |
| labels = np.concatenate(labels, axis=0) |
| unique, counts = np.unique(labels, return_counts=True) |
| for i, cnt in zip(unique, counts): |
| class_count[i] += cnt |
| print("Class counts:", class_count) |
| |
| ft_task = all_configs.get('ft_task', 'sleep_staging') |
| if ft_task == 'sleep_staging': |
| class_weights = np.emath.logn(all_configs['num_classes'], np.sum(class_count)/class_count) |
| class_weights = np.maximum(1.0, class_weights) |
| print("Class weights:", class_weights) |
| |
| elif ft_task == 'apnea_seg': |
| class_count_adj = np.zeros(2, dtype=np.int64) |
| class_count_adj[0] = class_count[0] |
| class_count_adj[1] = np.sum(class_count[1:]) |
| |
| |
| |
| class_weights = np.emath.logn(2, np.sum(class_count_adj)/class_count_adj) + 1.0 |
| print("Class weights:", class_weights) |
|
|
| |
| else: |
| class_weights = np.ones(all_configs['num_classes']) |
|
|
| print("Data Loaded") |
| |
| |
| |
| print("Model Loading ...") |
| model = ConvMAE1d_Finetune( |
| 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'], |
| class_weights = class_weights, |
| sig_len = sig_len_target * TARGET_RESAMPLE, |
| use_shared = use_shared, |
| use_moe = use_moe, |
| full_finetune=all_configs.get('full_finetune', False), |
| optimizer_name = all_configs['optimizer_name'], |
| lr = all_configs['lr'], |
| weight_decay = float(all_configs.get('weight_decay', 0.05)), |
| scheduler_name = all_configs['scheduler_name'], |
| num_epochs = all_configs['num_epochs'], |
| 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) |
| |
| ft_task = all_configs.get('ft_task', 'sleep_staging') |
| if ft_task == 'diagnosis': |
| main_metric = 'val_cindex' |
| elif ft_task in CLASSIFICATION_TASKS: |
| if 'f1_class_index' in all_configs: |
| main_metric = f"val_f1_class_{all_configs['f1_class_index']}" |
| else: |
| main_metric = 'val_f1' |
| elif ft_task in REGRESSION_TASKS: |
| main_metric = 'val_mae' |
| |
| checkpoint_callback = ModelCheckpoint( |
| monitor=main_metric, |
| mode='max', |
| save_top_k=1, |
| dirpath=all_configs['ckpt_save_path'], |
| filename=all_configs['experiment_name'] + f'-{{epoch:02d}}-{{{main_metric}:.2f}}', |
| verbose=True |
| ) |
| class SaveLastEpochCheckpoint(Callback): |
| def on_train_end(self, trainer, pl_module): |
| |
| final_checkpoint_path = f"{all_configs['ckpt_save_path']}/{all_configs['experiment_name']}-final.ckpt" |
| trainer.save_checkpoint(final_checkpoint_path) |
| print(f"Final checkpoint saved at {final_checkpoint_path}") |
|
|
| |
| previous_run_id = all_configs['previous_run_id'] |
| if args.no_wandb: |
| wandb_logger = None |
| elif previous_run_id: |
| wandb_logger = WandbLogger( |
| project=all_configs['project_name'], |
| entity='gundo0102', |
| name=all_configs['experiment_name'], |
| config=all_configs, |
| resume='must', |
| id=previous_run_id |
| ) |
| else: |
| wandb_logger = WandbLogger( |
| project=all_configs['project_name'], |
| entity='gundo0102', |
| name=all_configs['experiment_name'], |
| config=all_configs |
| ) |
|
|
| trainer = pl.Trainer( |
| devices=all_configs['gpus'], |
| max_epochs=all_configs['num_epochs'], |
| accelerator='gpu', |
| logger=wandb_logger, |
| precision = 'bf16-mixed', |
| |
| callbacks=[checkpoint_callback, SaveLastEpochCheckpoint()], |
| strategy='ddp_find_unused_parameters_true', |
| accumulate_grad_batches=all_configs['accumulate_grad_batches'], |
| num_sanity_val_steps=1, |
| check_val_every_n_epoch=1, |
| ) |
| |
| if all_configs['previous_checkpoint']: |
| trainer.fit(model, datamodule=data_module, ckpt_path=all_configs['previous_checkpoint']) |
| else: |
| trainer.fit(model, datamodule=data_module) |
| |
| trainer.test(ckpt_path="best", datamodule=data_module) |
| |
| wandb_logger.experiment.finish() |
| |
| |
| if __name__ == '__main__': |
| main() |
|
|
|
|
|
|
|
|
