import pytorch_lightning as pl from pytorch_lightning.loggers import WandbLogger import torch 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 import random import argparse from configs.global_cfg import TARGET_RESAMPLE, DATASET_PATH from train_utils import get_datasets, 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) from torch.optim.lr_scheduler import LambdaLR BASE_DIR = os.path.dirname(os.path.abspath(__file__)) class PretrainDataModule(pl.LightningDataModule): def __init__(self, data_root, batch_size_t, transforms_t, num_workers, proportion_t, num_classes=5, splits_set=None, modalities=["EEG"], sig_len_target=30, seq_truncate=True, dual_views=False ): super().__init__() self.data_root = data_root self.transforms_t = transforms_t self.batch_size_t = batch_size_t self.num_workers = num_workers self.proportion_t = proportion_t self.num_classes = num_classes self.splits_set = splits_set self.modalities = modalities self.sig_len_target = sig_len_target self.seq_truncate = seq_truncate self.dual_views = dual_views 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'] print("_"*50) print("Number of training subjects: ", len(self.train_data_list)) is_distributed = self.trainer.num_devices > 1 and self.trainer.strategy.root_device.type == "cuda" sampler = None if is_distributed and not self.dual_views: sampler = 'distributed' self.train_loader = data_loader.build_loader( self.data_root, self.train_data_list, batch_size=self.batch_size_t, shuffle=False, transforms=self.transforms_t, sampler=sampler, 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, full_night=True, pretrain_macro=True, seq_truncate=self.seq_truncate, dual_views=self.dual_views) print("Train loader ready") print("Length of train_loader: ", len(self.train_loader)) def train_dataloader(self): return self.train_loader class ConvMAE1d_PretrainSeq(pl.LightningModule): def __init__(self, model_name='convmae1d_multi', num_classes=5, modalities=["EEG"], num_experts=4, num_experts_per=2, sig_len=6000, use_shared=True, use_moe=True, use_revin=False, full_finetune=False, optimizer_name='AdamW', lr=0.0005, weight_decay=0.05, scheduler_name='CosineAnnealingLR', num_epochs=10, loss='MSE', patch_size=10, 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.use_revin = use_revin self.patch_size = patch_size self.full_finetune = full_finetune self.simclr = cfg.get('simclr', False) if cfg is not None else False self.dino = cfg.get('dino', False) if cfg is not None else False kwargs = {'demo_keys': cfg.get('demo_keys', ['age','bmi','sex'])} if cfg is not None else {} kwargs['cycle_len'] = cfg.get('cycle_len', 180) kwargs['use_transformer'] = cfg.get('use_transformer', False) kwargs['bidi'] = False if kwargs['use_transformer'] else cfg.get('bidi', True) kwargs['simclr'] = self.simclr kwargs['dino'] = self.dino 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, pretrain=True, **kwargs) # wrap backbone with pretraining wrapper when using SimCLR or DINO if self.simclr or self.dino: from models.macro_pretraining_abl import PretrainingWrapper self.model = PretrainingWrapper(self.model, simclr=self.simclr, simclr_temp=cfg.get('simclr_temp', 0.1) if cfg is not None else 0.1, dino=self.dino, dino_student_temp=cfg.get('dino_student_temp', 0.1) if cfg is not None else 0.1, dino_teacher_temp=cfg.get('dino_teacher_temp', 0.04) if cfg is not None else 0.04, dino_center_momentum=cfg.get('dino_center_momentum', 0.9) if cfg is not None else 0.9, dino_momentum=cfg.get('dino_momentum', 0.996) if cfg is not None else 0.996) 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() def forward(self, s, mask=None, demo=None): return self.model(s, mask=mask, demo=demo) def training_step(self, batch, batch_idx): inputs, labels, mask, demo, pids = batch # signals = inputs["EEG"].to(self.device, non_blocking=True) if self.simclr or self.dino: inputs1, inputs2 = inputs for key in inputs1.keys(): inputs1[key] = inputs1[key].to(self.device, non_blocking=True) inputs2[key] = inputs2[key].to(self.device, non_blocking=True) inputs = (inputs1, inputs2) else: 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) loss = self(inputs, mask=mask, demo=demo) # print("loss", loss.item()) # print("Demographics: ", demo) if loss.abs() > 1e4: print("PIDs: ", pids) # exit(0) 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) # wandb.log({"train_loss_every_step": loss.item()}) return loss 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) # Initialize teacher weights for DINO if applicable if self.dino: self.model.init_teacher() print(f"Loaded pretrained model from {ckpt}") print(msg) def configure_optimizers(self): # Base learning rate and weight decay base_lr = self.hparams.lr weight_decay = float(self.hparams.weight_decay) if hasattr(self.hparams, 'weight_decay') else 0.05 layer_decay = float(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 # 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.cfg["num_epochs"] / len(self.cfg['gpus'])) + 1 print("max steps", max_steps, steps_per_epoch, self.hparams.num_epochs) 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 = int(steps_per_epoch * self.cfg["num_epochs"] / len(self.cfg['gpus'])) + 1 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 import yaml def get_all_configs_from_yaml(file_path): with open(file_path, 'r') as file: configs = yaml.safe_load(file) return configs def main(): # Parse command-line arguments parser = argparse.ArgumentParser(description="Pretrain singal convMAE 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']) print("Number of train subjects: ", len(splits_set['train'])) 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 = PretrainDataModule( data_root = all_configs['data_root'], transforms_t = None, batch_size_t = all_configs['batch_size_t'], num_workers = all_configs['num_workers'], proportion_t = all_configs['proportion_t'], num_classes = all_configs['num_classes'], splits_set = splits_set, modalities = all_configs['modalities'], sig_len_target = sig_len_target, seq_truncate = all_configs.get('seq_truncate', True), dual_views = all_configs.get('simclr', False) or all_configs.get('dino', False) ) print("Data Loaded") print("Model Loading ...") kwargs = { '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'], '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': all_configs.get('weight_decay', 0.05), 'scheduler_name': all_configs['scheduler_name'], 'num_epochs': all_configs['num_epochs'], 'use_revin': all_configs.get('use_revin', False), 'patch_size': all_configs.get('patch_size', 10), } model = ConvMAE1d_PretrainSeq( **kwargs, cfg=all_configs ) if all_configs['pretrain_checkpoint']: ckpt_file = os.path.join('experiments', all_configs['pretrain_checkpoint']) model.load_pretrained(ckpt_file) print("Model Loaded") checkpoint_callback = ModelCheckpoint( monitor='train_loss', mode='min', save_top_k=1, dirpath=all_configs['ckpt_save_path'], filename=all_configs['experiment_name'] + '-{epoch:02d}-{train_loss:.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}") class ReshuffleDatasetCallback(pl.Callback): def on_train_epoch_start(self, trainer, pl_module): # Reshuffle the dataset at the start of each epoch if hasattr(trainer.datamodule, 'train_loader') and hasattr(trainer.datamodule.train_loader.dataset, 'reshuffle_full_night_index'): trainer.datamodule.train_loader.dataset.reshuffle_full_night_index() print(f"Reshuffled full night index for epoch {trainer.current_epoch}") # 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', callbacks=[checkpoint_callback, SaveLastEpochCheckpoint(), ReshuffleDatasetCallback()], strategy= 'ddp_find_unused_parameters_true', accumulate_grad_batches=all_configs['accumulate_grad_batches'], num_sanity_val_steps=1, log_every_n_steps=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) wandb_logger.experiment.finish() if __name__ == '__main__': main()