pretrain.py

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()