pretrain_seq.py

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