import pytorch_lightning as pl from pytorch_lightning.loggers import WandbLogger import torch import torch.nn as nn from torch.optim.lr_scheduler import _LRScheduler, 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 train_utils import get_datasets, WarmupCosineAnnealingLR import random import argparse from configs.global_cfg import TARGET_RESAMPLE, DATASET_PATH 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) 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=60 ): 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 def setup(self, stage=None): self.train_data_list = self.splits_set['train'] 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', pretrain_micro=True, num_classes=self.num_classes, modalities=self.modalities, sig_len_target=self.sig_len_target ) print("Loaders are ready") print("Length of train_loader: ", len(self.train_loader)) def train_dataloader(self): return self.train_loader class ConvMAE1d_Pretrain(pl.LightningModule): def __init__(self, cfg, ): super().__init__() self.save_hyperparameters() self.cfg = cfg self.num_classes = self.cfg["num_classes"] self.modalities = self.cfg["modalities"] self.num_experts = self.cfg["num_experts"] self.num_experts_per = self.cfg["num_experts_per"] self.sig_len = self.cfg['sig_len_target'] * TARGET_RESAMPLE self.use_shared = self.cfg["use_shared"] self.use_moe = self.cfg["use_moe"] self.use_revin = self.cfg["use_revin"] self.patch_size = self.cfg["patch_size"] self.loss_weights = self.cfg["loss_weights"] kwargs = {'mask_ratio': cfg.get('mask_ratio', 0.5)} self.model = models.__dict__[self.cfg["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, patch_size=self.patch_size, contrastive_training=(self.loss_weights['contrastive'] > 0.0), koleo_training=(self.loss_weights['koleo'] > 0.0), **kwargs ) if self.cfg["loss"] == 'MSE': self.criterion = nn.MSELoss() elif self.cfg["loss"] == 'MAE': self.criterion = nn.L1Loss() print(self.model) def forward(self, s) : return self.model(s) def training_step(self, batch): inputs = batch for key in inputs.keys(): inputs[key] = inputs[key].to(self.device, non_blocking=True) losses, moe_loss, contr_loss, koleo_loss, _, _, _ = self(inputs) self.log('learning_rate', self.trainer.optimizers[0].param_groups[0]['lr'], on_step=True, on_epoch=True, sync_dist=True) self.log('MoE_load_balancing_loss', moe_loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True) loss = 0.0 for m in inputs.keys(): self.log(f'train_loss_{m}', losses[m], on_step=True, on_epoch=True, sync_dist=True, prog_bar=True) loss += losses[m] loss = loss / len(inputs.keys()) * self.loss_weights['mae'] loss += self.loss_weights['moe'] * moe_loss self.log('Contrastive loss', contr_loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True) loss += self.loss_weights['contrastive'] * contr_loss self.log('KoLeo loss', koleo_loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True) loss += self.loss_weights['koleo'] *koleo_loss # koleo_loss is already weighted by self.koleo_weight in the model self.log('train_loss', loss, on_step=True, on_epoch=True, sync_dist=True, prog_bar=True) return loss def configure_optimizers(self): weight_decay = float(self.cfg["weight_decay"]) if hasattr(self.cfg, 'weight_decay') else 0.05 optimizer = getattr(torch.optim, self.cfg["optimizer_name"])( filter(lambda p: p.requires_grad, self.parameters()), lr=self.cfg["lr"], weight_decay=weight_decay # self.parameters(), lr=self.hparams.lr ) if self.cfg["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 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'])) # pd.DataFrame(splits_set['train']).to_csv('train_subjects.csv', index=False, header=False) all_configs['sig_len_target'] = all_configs.get('sig_len_target', 60) all_configs['patch_size'] = all_configs.get('patch_size', 10) all_configs['loss_weights'] = all_configs.get('loss_weights', {'mae': 1.0,'moe': 0.001,'contrastive': 0.0, 'koleo': 0.0}) all_configs['use_shared'] = all_configs.get('use_shared', True) all_configs['use_moe'] = all_configs.get('use_moe', True) all_configs['use_revin'] = all_configs.get('use_revin', False) all_configs['loss'] = all_configs.get('loss', 'MSE') 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 = all_configs['sig_len_target'] ) print("Data Loaded") print("Model Loading ...") if all_configs['pretrained']: model = ConvMAE1d_Pretrain.load_from_checkpoint( all_configs['pretrained'], strict=False, cfg=all_configs, # **kwargs ) else: model = ConvMAE1d_Pretrain( all_configs, # **kwargs ) print("Model Loaded") checkpoint_callback = ModelCheckpoint( dirpath=all_configs['ckpt_save_path'], filename=all_configs['experiment_name'] + '-{step:05d}-{train_loss:.2f}', every_n_train_steps=5000, # Save every 5000 steps save_top_k=-1, # Set to -1 to keep all checkpoints, or 1 for only the latest ) 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', callbacks=[checkpoint_callback, SaveLastEpochCheckpoint()], strategy= 'ddp_find_unused_parameters_true', accumulate_grad_batches=all_configs['accumulate_grad_batches'], # No sanity check num_sanity_val_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()