train_from_checkpoint.py

#!/usr/bin/env python
import os
import wandb
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
import torch
import yaml
from omegaconf import OmegaConf, DictConfig
import hydra

from models.hetero_gnn.simba_gnn import HeteroGNN
from models.hetero_gnn.lightning_module import HeteroGNNTrainer
from models.hetero_gnn.data_module import HeteroGraphDataModule

def load_checkpoint(checkpoint_path, model, stage):
    """Load a checkpoint and return the model state dict."""
    # Add safe globals for omegaconf
    torch.serialization.add_safe_globals(['omegaconf.dictconfig.DictConfig'])
    
    try:
        # First try loading with weights_only=True (new default in PyTorch 2.6)
        checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True)
    except Exception as e:
        print(f"Failed to load with weights_only=True, trying with weights_only=False: {str(e)}")
        # If that fails, try loading with weights_only=False (old behavior)
        checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
    
    # If the checkpoint is a Lightning checkpoint, extract the model state dict
    if 'state_dict' in checkpoint:
        state_dict = checkpoint['state_dict']
        # Remove 'model.' prefix from keys if it exists
        new_state_dict = {}
        for k, v in state_dict.items():
            if k.startswith('model.'):
                new_state_dict[k[6:]] = v
            else:
                new_state_dict[k] = v
        model.load_state_dict(new_state_dict)
    else:
        # If it's a direct model state dict
        model.load_state_dict(checkpoint)
    
    return model

def train_from_checkpoint(checkpoint_path, target_stage):
    """Train a model from a checkpoint for a specific stage."""
    # Load checkpoint
    checkpoint_data = torch.load(checkpoint_path, map_location='cpu', weights_only=False)

    # Create a comprehensive default config
    default_config = {
        'general': {
            'seed': 42,
            'num_gpus': 1,
            'precision': 'bf16-mixed',
            'compile': False
        },
        'data': {
            'data_dir': '/home/aac/Mohammad/GNN/MetaBiome/data',
            'num_workers': 4,
            'train_ratio': 0.8,
            'val_ratio': 0.1,
            'test_ratio': 0.1,
            'shuffle': True,
            'use_test_data': True,
            'sample_fraction': 1.0,
            'similarity_threshold': 0.85
        },
        'model': {
            'encoder': {
                'd_model': 768,
                'num_layers': 3,
                'num_heads': 12, # Will be adjusted based on d_model later
                'dropout': 0.2
            },
            'edge_drop': 0.2,
            'feature_mask': 0.1
        },
        'training': {
            'gradient_clip_val': 1.0,
            'gradient_clip_algorithm': 'norm',
            'loss': { # General loss definitions
                'presence': {'name': 'BCEWithLogitsLoss', 'weight': 1.0},
                'abundance': {'name': 'HuberLoss', 'delta': 1.0, 'weight': 1.0},
                'rank': {'name': 'MarginRankingLoss', 'margin': 0.1, 'weight': 0.1},
                'flux': {'name': 'BCEWithLogitsLoss', 'weight': 1.0}, # Changed from MSELoss to BCEWithLogitsLoss as per lightning_module
                'contrastive': {'weight': 1.0} # For SSL, temperature is used as weight
            },
            'stages': {
                'ssl': {
                    'epochs': 50,
                    'batch_size': 256,
                    'accumulate_grad_batches': 1,
                    'optimizer': {'lr': 1.0e-3, 'weight_decay': 1e-4, 'betas': [0.9, 0.98]},
                    'scheduler': {'name': 'OneCycleLR', 'max_lr': 2e-3, 'pct_start': 0.1, 'div_factor': 5, 'final_div_factor': 100},
                    'loss': { # Stage-specific SSL loss weights
                        'contrastive': {'weight': 1.0},
                        'presence': {'weight': 1.0},
                        'abundance': {'weight': 1.0},
                        'flux': {'weight': 0.0},
                        'rank': {'weight': 0.1}
                    }
                },
                'supervised': {
                    'epochs': 40,
                    'batch_size': 4,
                    'accumulate_grad_batches': 1,
                    'optimizer': {'lr': 5e-4, 'weight_decay': 1e-4, 'betas': [0.9, 0.98]},
                    'scheduler': {'name': 'OneCycleLR', 'max_lr': 1e-3, 'pct_start': 0.1, 'div_factor': 5, 'final_div_factor': 100},
                    'loss': { # Stage-specific supervised loss weights
                        'contrastive': {'weight': 0.0},
                        'presence': {'weight': 1.0},
                        'abundance': {'weight': 1.0},
                        'flux': {'weight': 1.0},
                        'rank': {'weight': 0.2}
                    }
                },
                'finetune': {
                    'epochs': 40,
                    'batch_size': 1,
                    'accumulate_grad_batches': 1,
                    'optimizer': {'lr': 1e-4, 'weight_decay': 1e-4, 'betas': [0.9, 0.98]},
                    'scheduler': {'name': 'ConstantLR', 'factor': 1.0, 'total_iters': 1},
                    'loss': { # Stage-specific finetune loss weights
                        'contrastive': {'weight': 0.0},
                        'presence': {'weight': 1.0},
                        'abundance': {'weight': 1.0},
                        'flux': {'weight': 0.0},
                        'rank': {'weight': 0.1} # Default, can be swept
                    }
                }
            }
        },
        'checkpointing': {
            'dirpath': '/home/aac/Mohammad/GNN/MetaBiome/checkpoints', # Base path, will be appended with run_name/stage
            'save_top_k': 1,
            'save_last': True,
            # filename, monitor, mode will be set dynamically
        },
        'early_stopping': {
            'patience': 15,
            'mode': 'max',
            'min_delta': 0.01,
            # monitor will be set dynamically
        }
    }
    conf = OmegaConf.create(default_config)

    # Extract config from checkpoint if available and merge
    if 'hyper_parameters' in checkpoint_data and checkpoint_data['hyper_parameters']:
        checkpoint_conf_dict = checkpoint_data['hyper_parameters']
        # Ensure that the loaded hyper_parameters are treated as a dict
        if isinstance(checkpoint_conf_dict, DictConfig):
             checkpoint_conf_dict = OmegaConf.to_container(checkpoint_conf_dict, resolve=True)
        
        # Before merging, we need to handle potential dot-separated keys from older wandb configs
        flat_checkpoint_conf = {}
        if isinstance(checkpoint_conf_dict, dict):
            for k, v in checkpoint_conf_dict.items():
                if '.' in k: # Indicates a flat key e.g., "model.encoder.d_model"
                    # Convert to nested dict structure for OmegaConf merging
                    keys = k.split('.')
                    d = flat_checkpoint_conf
                    for key_part in keys[:-1]:
                        d = d.setdefault(key_part, {})
                    d[keys[-1]] = v
                else:
                    # if not dot-separated, it might be a top-level key like 'model' itself
                    # or a non-config item. We'll try to merge it directly.
                    # If it's a dict, OmegaConf will merge it.
                    flat_checkpoint_conf[k] = v
            
            # If flat_checkpoint_conf was populated, use it. Otherwise, use original checkpoint_conf_dict
            if flat_checkpoint_conf:
                 conf.merge_with(OmegaConf.create(flat_checkpoint_conf))
            else: # If no dot-separated keys, merge directly
                 conf.merge_with(OmegaConf.create(checkpoint_conf_dict))

    # Set seeds for reproducibility
    pl.seed_everything(conf.general.seed)

    # Adjust num_heads based on d_model
    if conf.model.encoder.d_model <= 512:
        conf.model.encoder.num_heads = 8
    else:
        conf.model.encoder.num_heads = 12
    
    # Create a temporary data module to get node features
    temp_data_module = HeteroGraphDataModule(
        data_dir=conf.data.data_dir,
        batch_size=1,
        num_workers=1,
        stage=target_stage, # Or a default like "ssl"
        sample_fraction=conf.data.sample_fraction,
        similarity_threshold=conf.data.similarity_threshold
    )
    temp_data_module.setup()
    node_features = temp_data_module.get_node_features()

    # Create model
    model = HeteroGNN(
        node_features=node_features,
        d_model=conf.model.encoder.d_model,
        num_layers=conf.model.encoder.num_layers,
        num_heads=conf.model.encoder.num_heads,
        dropout=conf.model.encoder.dropout,
        edge_dropout=conf.model.edge_drop,
        feature_mask=conf.model.feature_mask
    )

    # Load model weights from checkpoint
    if 'state_dict' in checkpoint_data:
        state_dict = checkpoint_data['state_dict']
        new_state_dict = {}
        for k, v in state_dict.items():
            if k.startswith('model.'):
                new_state_dict[k[len('model.'):]] = v
            else:
                new_state_dict[k] = v
        model.load_state_dict(new_state_dict, strict=False) # Use strict=False if some keys might mismatch
    else:
        # If no 'state_dict', assume checkpoint_data itself is the state_dict (less common for Lightning)
        model.load_state_dict(checkpoint_data, strict=False)

    # Convert OmegaConf to dict for wandb
    wandb_config_dict = OmegaConf.to_container(conf, resolve=True)
    
    run_name_suffix = f"{conf.model.encoder.d_model}D-{conf.model.encoder.num_heads}H-{conf.model.edge_drop}ED-{conf.model.feature_mask}FM-{conf.data.similarity_threshold}ST"
    wandb_run_name = f"{target_stage}_from_ckpt_{run_name_suffix}"


    # Initialize wandb
    wandb_run = wandb.init(
        project="MetaBiomeX_Resumed", # Or your project name
        entity="mypersonalteam",  # Your wandb entity
        config=wandb_config_dict,
        name=wandb_run_name,
        resume="allow", # Allow resuming if run ID exists
        id=wandb.util.generate_id() # Generate a new ID for each new "from_checkpoint" run
    )

    # Create WandB logger
    wandb_logger = WandbLogger(
        experiment=wandb_run,
        log_model=False # Or True if you want to log the model
    )

    # Create data module for target stage
    data_module = HeteroGraphDataModule(
        data_dir=conf.data.data_dir,
        batch_size=conf.training.stages[target_stage].batch_size,
        num_workers=conf.data.num_workers,
        stage=target_stage,
        edge_drop=conf.model.edge_drop, # Pass these if DataModule uses them
        feature_mask=conf.model.feature_mask,
        sample_fraction=conf.data.sample_fraction,
        similarity_threshold=conf.data.similarity_threshold
    )
    data_module.setup()

    # Create lightning module
    lightning_module = HeteroGNNTrainer(
        model=model,
        config=conf, # Pass the OmegaConf object here
        stage=target_stage,
        edge_drop=conf.model.edge_drop,
        feature_mask=conf.model.feature_mask,
        log_every_n_steps=5 # Or from config
    )

    # Checkpointing setup
    checkpoint_dir = os.path.join(conf.checkpointing.dirpath, wandb_run.name, target_stage)
    os.makedirs(checkpoint_dir, exist_ok=True)
    
    # Define monitor metric based on stage
    monitor_metric_map = {
        "ssl": f"{target_stage}_val_contrastive_loss",
        "supervised": f"{target_stage}_val_abundance_r2_all",
        "finetune": f"{target_stage}_val_abundance_r2_all"
    }
    monitor_metric = monitor_metric_map.get(target_stage, f"{target_stage}_val_loss")

    checkpoint_callback = ModelCheckpoint(
        dirpath=checkpoint_dir,
        filename=f"{wandb_run.name}-{target_stage}-{{epoch}}-{{{monitor_metric}:.2f}}_resumed",
        monitor=monitor_metric,
        mode=conf.checkpointing.get('mode', 'max' if 'r2' in monitor_metric else 'min'),
        save_top_k=conf.checkpointing.save_top_k,
        save_last=conf.checkpointing.save_last
    )

    callbacks = [checkpoint_callback]
    if conf.early_stopping.get('patience', 0) > 0:
        early_stopping = EarlyStopping(
            monitor=monitor_metric,
            mode='max',
            patience=conf.early_stopping.patience,
            verbose=True
        )
        callbacks.append(early_stopping)

    # Create trainer
    trainer = pl.Trainer(
        max_epochs=conf.training.stages[target_stage].epochs,
        accelerator="gpu" if torch.cuda.is_available() else "cpu",
        devices=conf.general.num_gpus,
        precision=conf.general.precision,
        logger=wandb_logger,
        callbacks=callbacks,
        gradient_clip_val=conf.training.gradient_clip_val,
        accumulate_grad_batches=conf.training.stages[target_stage].accumulate_grad_batches,
        # accumulate_grad_batches=conf.training.accumulate_grad_batches, # Check if this is global or per stage
        gradient_clip_algorithm=conf.training.gradient_clip_algorithm,
        deterministic=True, # As per original script
        # progress_bar_refresh_rate=1 # Deprecated, use enable_progress_bar
        enable_progress_bar=True
    )

    # Train
    trainer.fit(lightning_module, datamodule=data_module)

    # Close wandb
    wandb_run.finish()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint file")
    parser.add_argument("--stage", type=str, required=True, choices=["supervised", "finetune"], help="Target stage to train")
    args = parser.parse_args()
    
    train_from_checkpoint(args.checkpoint, args.stage)