SleepMaMi / pretrain.py

Keondopark

SDB segmentation update to report Class-F1

6196ce4 2 days ago

13.2 kB

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

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