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) # 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 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): # Split data into train and validation sets # all_data_list = os.listdir(self.video_path) 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) # Survival analysis 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) # Classification tasks 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() # Metrics 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) # Classification tasks elif self.ft_task in REGRESSION_TASKS: if loss == 'mse': self.criterion = nn.MSELoss() elif loss == 'mae': self.criterion = nn.L1Loss() # Metrics 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): # Load the pretrained model checkpoint 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']: # print("Matching key: ", k) new_state_dict[k] = checkpoint['state_dict'][k] else: # print(f"Missing key (Only in the current model), tries to find model.{k}") if f"model.{k}" in checkpoint['state_dict']: # print("Matching key, starts with 'model.' in the checkpoint, delete 'model. prefix': ", k) new_state_dict[k] = checkpoint['state_dict'][f"model.{k}"] # Remove 'model.' prefix 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. """ # Set the train mode for the whole model (including head and encoder) super().train(mode) # Now, *force* the encoder back to eval 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': # Log metrics 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 # 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) return loss def forward_epoch_level(self, batch, mode='train'): 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) 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 # signals = inputs["EEG"].to(self.device, non_blocking=True) 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.permute(0,2,1)[mask] pred = pred[mask] labels = labels[mask] # labels /= 100 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 learning rate and weight decay 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 # Create parameter groups with layer-wise learning rate decay 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) # 10% of total steps for warmup scheduler = WarmupCosineAnnealingLR( optimizer, warmup_epochs=warmup_steps, # Number of 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 # 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) 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 = os.path.join(BASE_DIR, all_configs['data_root']) datasets = all_configs.get('datasets', ['phy']) # Load data 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() # Setup the data module to prepare the datasets and dataloaders 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': # For full night training, labels is a list of arrays 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:]) # Combine 1 and 2 # class_weights = np.sum(class_count_adj)/class_count_adj # class_weights = np.maximum(np.emath.logn(2, np.sum(class_count_adj)/class_count_adj), 1.0) 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), # Default to False if not specified 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): # 마지막 epoch의 체크포인트 저장 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}") # if previous_run_id is not None, load the best model from the previous run 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 # 이전 실험의 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', # gradient_clip_val=1.0, 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()