SleepMaMi / pretrain_seq.py

Keondopark

SDB segmentation update to report Class-F1

6196ce4 2 days ago

19.6 kB

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

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