main_pipelines/main_pretrain.py

from pprint import pprint
import os
from argparse import ArgumentParser, Namespace
import datetime
from dateutil import tz
import random
import numpy as np
import torch
import warnings
from datetime import timedelta
from pytorch_lightning import seed_everything, Trainer
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor, Callback
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.strategies import DDPStrategy


class DenseStepCheckpoint(Callback):
    """Save checkpoints at specific training steps."""
    
    def __init__(self, dirpath: str, save_steps: list = None):
        super().__init__()
        self.dirpath = dirpath
        self.save_steps = set(save_steps) if save_steps else {1, 10, 100, 1000, 10000, 100000}
        self.saved_steps = set()
        
    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        global_step = trainer.global_step
        if global_step in self.save_steps and global_step not in self.saved_steps:
            ckpt_path = os.path.join(self.dirpath, f"step={global_step}.ckpt")
            trainer.save_checkpoint(ckpt_path)
            self.saved_steps.add(global_step)
            if trainer.is_global_zero:
                print(f"[DenseStepCheckpoint] Saved checkpoint at step {global_step}: {ckpt_path}")

from osf.datasets.pretrain_datamodule import SleepDataModule
from osf.models.dino_model_cls import DINOCLSModel
from config import *
from train_config import *

warnings.filterwarnings("ignore")
os.environ["TOKENIZERS_PARALLELISM"] = "false"
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision('high')


torch._dynamo.config.cache_size_limit = 128
torch._dynamo.config.optimize_ddp = False



def param_stats(model: torch.nn.Module, verbose: bool = False):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    if verbose:
        print(f"{'Name':40s} {'Shape':20s} {'#Params':>10s} {'Train?':>6s}")
        print("-" * 80)
        for name, p in model.named_parameters():
            print(f"{name:40s} {str(list(p.shape)):20s} {p.numel():10d} {str(p.requires_grad):>6s}")
        print("-" * 80)
    print(f"Total parameters:     {total / 1e6:.3f} M ({total})")
    print(f"  Trainable params:   {trainable / 1e6:.3f} M ({trainable})")
    print(f"  Frozen params:      {(total-trainable) / 1e6:.3f} M ({total-trainable})")
def main(hparams: Namespace):

    now = datetime.datetime.now(tz.tzlocal())
    extension = now.strftime("%Y_%m_%d_%H_%M_%S")
    extension = f"final_sleep_unimodal_{hparams.model_name}_{hparams.psg_encoder_name}_bz{hparams.batch_size}_{extension}"
    ckpt_dir = os.path.join(
        CKPT_PATH, f"logs/sleepuni/ckpts/{extension}")
    os.makedirs(ckpt_dir, exist_ok=True)
    if hparams.model_name in MODEL_LIST:
        callbacks = [
            LearningRateMonitor(logging_interval="step"),
            ModelCheckpoint(monitor="val/loss", dirpath=ckpt_dir,
                            save_last=True, every_n_epochs=2, mode="min", save_top_k=-1,
                            save_on_train_epoch_end=False, auto_insert_metric_name=True),
        ]
        if hparams.dense_ckpt:
            dense_ckpt_dir = os.path.join(ckpt_dir, "dense_steps")
            os.makedirs(dense_ckpt_dir, exist_ok=True)
            callbacks.append(DenseStepCheckpoint(
                dirpath=dense_ckpt_dir,
                save_steps=hparams.dense_ckpt_steps
            ))
    else:
        raise NotImplementedError
    logger_dir = os.path.join(CKPT_PATH, "logs/sleepuni")
    os.makedirs(logger_dir, exist_ok=True)
    print("wandb logger dir: ", logger_dir)
    wandb_logger = WandbLogger(
        project=hparams.wandb_proj_name + f'final_{hparams.model_name}_{hparams.psg_encoder_name}_bz{hparams.batch_size}', save_dir=logger_dir, name=extension)

    strategy = DDPStrategy(
        find_unused_parameters=True,
        static_graph=False,
        timeout=timedelta(minutes=15),
    )

    trainer = Trainer(
        max_epochs=hparams.max_epochs,
        accelerator="gpu",
        accumulate_grad_batches=hparams.accumulate_grad_batches,
        devices=hparams.num_devices,
        num_nodes=hparams.num_nodes,
        precision=hparams.precision,
        gradient_clip_val=3.0,
        gradient_clip_algorithm="norm",
        strategy=strategy,
        callbacks=callbacks,
        logger=wandb_logger,
        log_every_n_steps=10,
    )

    hparams.exp_log_dir = os.path.join(
        CKPT_PATH, f"data/{extension}/exp_logs")
    train_edf_cols = MONITOR_TYPE_MAP.get(hparams.monitor_type, TRAIN_EDF_COLS_UNI_ENC)
    hparams.num_leads = len(train_edf_cols)
    
    dm = SleepDataModule(
            is_pretrain    = 1,
            csv_dir        = SPLIT_DATA_FOLDER,
            train_edf_cols = train_edf_cols,
            batch_size     = hparams.batch_size,
            num_workers    = hparams.num_workers,
            data_pct       = hparams.train_data_pct,
            window_size = 30,
            sample_rate = 64,
            val_dataset_list = hparams.val_dataset_list,
            data_source = hparams.data_source,
            include_datasets = hparams.include_datasets,
        )

    hparams.simclr_augmentation = AUGMENTATION_MAP.get(hparams.model_name, "none")

    # Create DINO model
    model = DINOCLSModel(**vars(hparams))
    model.training_steps_per_epoch = len(dm.train_dataloader()) // hparams.accumulate_grad_batches // hparams.num_devices
    model.teacher_temp_warmup_iters = model.training_steps_per_epoch * 0.1 * hparams.max_epochs
    print(f"[INFO] DINO teacher warmup steps: {model.teacher_temp_warmup_iters}")
    pprint(vars(hparams))

    if hparams.ckpt_path:
        trainer.fit(model, datamodule = dm, ckpt_path=hparams.ckpt_path)
    else:
        trainer.fit(model, datamodule = dm)


if __name__ == '__main__':
    parser = ArgumentParser(description="Pretraining DINO model for sleep PSG data.")
    parser.add_argument("--model_name", type=str, default="dino_ours",
                        choices=MODEL_LIST)
    
    parser.add_argument("--psg_encoder_name", type=str, default="vit_base")
    parser.add_argument("--val_dataset_list", default=PRETRAIN_VAL_DATASET_LIST)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--train_data_pct", type=float, default=1.)
    parser.add_argument("--data_source", type=str, default="auto",
                        choices=["auto", "pretrain", "downstream", "both"])
    parser.add_argument("--include_datasets", type=str, nargs="*", default=None)
    parser.add_argument("--monitor_type", type=str, default="main",
                        choices=["main", "type3", "type4"],
                        help="Channel configuration: main (12ch), type3 (5ch), type4 (3ch)")
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--patch_size_time", type=int, default=4)
    parser.add_argument("--patch_size_ch", type=int, default=4)
    parser.add_argument("--use_2d_pos_embed", type=bool, default=True)
    parser.add_argument("--sample_rate", type=int, default=64)
    parser.add_argument("--num_workers", type=int, default=64)
    parser.add_argument("--num_devices", type=int, default=4)
    parser.add_argument("--num_nodes", type=int, default=1)
    parser.add_argument("--max_epochs", type=int, default=30)
    parser.add_argument("--accumulate_grad_batches", type=int, default=1)
    parser.add_argument("--precision", type=str, default="32-true")
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--text_encoder_name", type=str, default="google/flan-t5-base")
    parser.add_argument("--lead_wise", type=int, default=0)
    parser.add_argument("--use_lead_embedding", type=int, default=1)
    # DINO-specific args
    parser.add_argument("--koleo_lambda", type=float, default=0.0)
    parser.add_argument("--ibot_lambda", type=float, default=0.0)
    parser.add_argument("--dino_out_dim", type=int, default=2048)
    parser.add_argument("--dino_patch_out_dim", type=int, default=2048)
    parser.add_argument("--dino_hidden_dim", type=int, default=2048)
    parser.add_argument("--dino_bottleneck_dim", type=int, default=256)
    parser.add_argument("--wandb_proj_name", type=str, default="sleepuni")
    parser.add_argument("--ckpt_path", type=str, default=None)
    parser.add_argument("--dense_ckpt", action="store_true")
    parser.add_argument("--dense_ckpt_steps", type=int, nargs="+", default=[10, 100, 200, 400, 500, 800, 1000, 1600, 2500, 3200, 6400, 10000, 12500, 12800, 25600, 51200, 62500, 100000])
    


    hparams = parser.parse_args()
    
    seed_everything(hparams.seed)
    main(hparams)