train_model.py

import numpy as np
import wandb
import os
import torch
import random
import warnings

import hydra
from omegaconf import DictConfig, OmegaConf

# Suppress nested tensor prototype warnings
warnings.filterwarnings(
    "ignore", message="The PyTorch API of nested tensors is in prototype stage"
)

import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from dataset.utils import (
    DynGraphDataset,
    pad_collate_fn,
    agcrn_collate_fn,
    create_weighted_sampler,
)
from models.models import DynGraphModel
from torch.utils.data import DataLoader, Subset
from tqdm import tqdm
from sklearn.decomposition import PCA
from einops import rearrange
from functools import partial

from utils.config_utils import is_sweep

# Profile the code
import cProfile
import pstats

pr = cProfile.Profile()
pr.enable()


def run_training(cfg: DictConfig, run_id: str = None) -> str:
    """Core training logic. Can be called directly or via Hydra.

    Args:
        cfg: Configuration dictionary
        run_id: Optional wandb run ID to resume. If None, creates new run.

    Returns:
        run_id: The wandb run ID used
    """
    encoder_type = cfg.model.encoder_type
    dataset_name = cfg.dataset.name

    # Init wandb - skip if already running (e.g., called from run_experiment.py)
    if wandb.run is None:
        if run_id:
            wandb.init(
                project=cfg.wandb.project,
                id=run_id,
                resume="allow",
                config=OmegaConf.to_container(cfg, resolve=True),
            )
            wandb.config.update(
                OmegaConf.to_container(cfg, resolve=True), allow_val_change=True
            )
        else:
            wandb.init(
                project=cfg.wandb.project,
                config=OmegaConf.to_container(cfg, resolve=True),
            )
    run_id = wandb.run.id

    # Select one GPU if more are available
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    verbose = cfg.verbose
    plot = cfg.plot

    if verbose:
        print("Configuration:", OmegaConf.to_yaml(cfg))

    seed = cfg.seed
    random.seed(seed)
    torch.manual_seed(seed)
    np.random.seed(seed)

    # Define dataset
    dataset = DynGraphDataset(**cfg.dataset)

    # Split it into train and validation sets
    indices = list(range(len(dataset)))
    train_idx, val_idx = train_test_split(indices, test_size=0.2, random_state=seed)

    train_dataset = Subset(dataset, train_idx)
    val_dataset = Subset(dataset, val_idx)

    # Choose the appropriate collate function based on encoder type
    if encoder_type == "agcrn":
        collate_fn = partial(
            agcrn_collate_fn, max_nodes=cfg.dataset.max_nodes, pad_value=-1
        )
    else:
        # For other encoders, pad_nodes is always False since encoder_type != 'agcrn'
        collate_fn = partial(
            pad_collate_fn,
            pad_nodes=False,
            max_nodes=cfg.dataset.max_nodes,
            pad_times=cfg.dataset.pad_times,
        )

    balance_labels = cfg.dataset.get("balance_labels", False)
    if balance_labels:
        train_labels = dataset.graph_labels[train_idx]
        train_sampler = create_weighted_sampler(train_labels)
        shuffle = False
    else:
        train_sampler = None
        shuffle = True

    train_loader = DataLoader(
        train_dataset,
        batch_size=cfg.training.batch_size,
        shuffle=shuffle,
        sampler=train_sampler,
        collate_fn=collate_fn,
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=cfg.training.batch_size,
        shuffle=False,
        collate_fn=collate_fn,
    )

    # Define model
    input_size = dataset[0].x.shape[-1]
    output_size = 1 if dataset.num_classes == 2 else dataset.num_classes
    # Pass model config + dataset params needed by model
    model_params = OmegaConf.to_container(cfg.model, resolve=True)
    model_params.update(
        {
            "input_size": input_size,
            "output_size": output_size,
            "max_nodes": cfg.dataset.max_nodes,
            "emb_size": cfg.get("emb_size", 16),
            "evolve_variant": cfg.get("evolve_variant", "H"),
            "temporal_kernel_size": cfg.model.get("temporal_kernel_size", 2),
            "dilation": cfg.model.get("dilation", 1),
        }
    )
    model = DynGraphModel(**model_params).to(device)

    # Watch model with wandb for gradient and parameter logging
    if verbose:
        wandb.watch(model, log="all", log_freq=20)

    # Define loss function and optimizer
    is_multiclass = output_size > 1
    criterion_pred = (
        torch.nn.CrossEntropyLoss() if is_multiclass else torch.nn.BCEWithLogitsLoss()
    )
    criterion_rec = (
        torch.nn.CrossEntropyLoss() if is_multiclass else torch.nn.BCEWithLogitsLoss()
    )
    criterion_obs = torch.nn.MSELoss()
    optimizer = torch.optim.Adam(
        model.parameters(), lr=cfg.training.lr, weight_decay=cfg.training.weight_decay
    )
    # Define scheduler
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer, step_size=cfg.training.step_size, gamma=cfg.training.gamma
    )

    # Set the model to training mode
    model.train()
    # Check if we need to freeze K parameters
    if cfg.training.alpha == 0 and cfg.training.beta == 0:
        # Freeze the K parameters in the encoder if both beta and alpha are zero
        model.encoder.K.requires_grad = False

    # Train the model
    num_epochs = cfg.training.max_epochs
    best_acc = -1.0
    best_model_path = None
    best_epoch = -1
    patience = cfg.training.patience
    counter = 0

    for epoch in tqdm(range(num_epochs), desc="Training", position=0, leave=True):
        for batch in tqdm(train_loader, position=1, leave=False):

            # Move the inputs and labels to the device
            if encoder_type == "agcrn":
                # For AGCRN: batch.x is already (B, T, N, F)
                input = batch.x.to(device)
                edge_index = None
                edge_weight = None
            else:
                # For other models: add batch dimension to TSL format
                input = batch.x.unsqueeze(0).to(device)
                edge_index = batch.edge_index
                edge_weight = (
                    batch.edge_weight if hasattr(batch, "edge_weight") else None
                )

            label = batch.y.squeeze().to(device)

            if isinstance(edge_index, list):
                edge_index = torch.nested.nested_tensor(edge_index).to(device)
            elif isinstance(edge_index, torch.Tensor):
                edge_index = edge_index.to(device)
            else:
                edge_index = None

            if isinstance(edge_weight, list):
                edge_weight = torch.nested.nested_tensor(edge_weight).to(device)
            elif isinstance(edge_weight, torch.Tensor):
                edge_weight = edge_weight.to(device)
            else:
                edge_weight = None

            # Get sequence lengths for this batch before padding
            if hasattr(batch, "last_timestep_mask"):
                last_timestep_mask = batch.last_timestep_mask.to(device)
            else:
                last_timestep_mask = None

            # Get node mask for this batch
            if hasattr(batch, "real_node_mask"):
                real_node_mask = batch.real_node_mask.to(device)
            else:
                real_node_mask = None

            # Zero the gradients
            optimizer.zero_grad()

            # Forward pass
            x, h, x_rec, h_rec = model(
                input,
                edge_index,
                edge_weight,
                batch=batch.batch.to(device),
                last_timestep_mask=last_timestep_mask,
                real_node_mask=real_node_mask,
            )

            # Compute the loss
            loss_pred = criterion_pred(x.squeeze(), label)
            if cfg.training.beta or cfg.training.alpha > 0:
                loss_rec = criterion_rec(x_rec.squeeze(), label)
                l2_reg = cfg.training.weight_decay * torch.sum(
                    torch.pow(model.encoder.K, 2)
                )
                loss_obs = criterion_obs(h_rec, h)
                loss_ridge = loss_obs + l2_reg
                loss_sum = (
                    loss_pred
                    + cfg.training.alpha * loss_rec
                    + cfg.training.beta * loss_ridge
                )
            else:
                loss_sum = loss_pred

            # Backward pass and optimization
            loss_sum.backward()
            optimizer.step()

        # Step the scheduler
        scheduler.step()
        wandb.log({"epoch": epoch, "lr": scheduler.get_last_lr()[0]})

        # Validation
        total_loss = 0
        total_loss_pred, total_loss_rec, total_loss_ridge, total_loss_obs = 0, 0, 0, 0
        predictions, labels_val, hs = [], [], []
        with torch.no_grad():
            for batch in tqdm(val_loader, desc="Validation", position=1, leave=False):

                # Move the inputs and labels to the device
                if encoder_type == "agcrn":
                    # For AGCRN: batch.x is already (B, T, N, F)
                    input = batch.x.to(device)
                    edge_index = None
                else:
                    # For other models: add batch dimension to TSL format
                    input = batch.x.unsqueeze(0).to(device)
                    edge_index = batch.edge_index

                label = batch.y.squeeze().to(device)
                if isinstance(edge_index, list):
                    edge_index = torch.nested.nested_tensor(edge_index).to(device)
                elif isinstance(edge_index, torch.Tensor):
                    edge_index = edge_index.to(device)
                else:
                    edge_index = None

                label = batch.y.squeeze().to(device)

                # Get sequence lengths for this batch before padding
                if hasattr(batch, "last_timestep_mask"):
                    last_timestep_mask = batch.last_timestep_mask.to(device)
                else:
                    last_timestep_mask = None

                # Get node mask for this batch
                if hasattr(batch, "real_node_mask"):
                    real_node_mask = batch.real_node_mask.to(device)
                else:
                    real_node_mask = None

                # Forward pass
                x, h, x_rec, h_rec = model(
                    input,
                    edge_index,
                    edge_weight=None,
                    batch=batch.batch.to(device),
                    last_timestep_mask=last_timestep_mask,
                    real_node_mask=real_node_mask,
                )
                predictions.append(x)
                labels_val.append(label)

                # Collect hidden states for plot:
                # sum over nodes for each batch and timestep
                if plot:
                    if batch.batch is not None and encoder_type != "agcrn":
                        batch_size = batch.batch.max().item() + 1
                        batch_hidden_states = [
                            h[0, :, batch.batch == n, :].sum(-2)
                            for n in range(batch_size)
                        ]
                        hs.extend(batch_hidden_states)
                    elif batch.batch is not None and encoder_type == "agcrn":
                        # For AGCRN: h has shape (B, T, N, F), sum over real nodes for each batch sample
                        real_node_mask = (
                            batch.real_node_mask.unsqueeze(1).unsqueeze(-1).to(device)
                        )  # (B, 1, N, 1)
                        batch_hidden_states = (h * real_node_mask).sum(
                            dim=2
                        )  # (B, T, F)
                        hs.extend(batch_hidden_states)
                    else:
                        hs.append(h[0].sum(-2))

                # Compute the loss
                loss_pred = criterion_pred(x.squeeze(), label)
                if cfg.training.beta or cfg.training.alpha > 0:
                    loss_rec = criterion_rec(x_rec.squeeze(), label)
                    l2_reg = cfg.training.weight_decay * torch.sum(
                        torch.pow(model.encoder.K, 2)
                    )
                    loss_obs = criterion_obs(h_rec, h)
                    loss_ridge = loss_obs + l2_reg
                    loss_sum = (
                        loss_pred
                        + cfg.training.alpha * loss_rec
                        + cfg.training.beta * loss_ridge
                    )
                else:
                    loss_rec = torch.tensor(0.0, device=device)
                    loss_ridge = torch.tensor(0.0, device=device)
                    loss_obs = torch.tensor(0.0, device=device)
                    loss_sum = loss_pred

                # Accumulate the total loss
                total_loss += loss_sum.item()
                total_loss_pred += loss_pred.item()
                total_loss_rec += loss_rec.item()
                total_loss_ridge += loss_ridge.item()
                total_loss_obs += loss_obs.item()

        # Calculate the average validation loss
        avg_loss = total_loss / len(val_loader)
        avg_loss_pred = total_loss_pred / len(val_loader)
        avg_loss_rec = total_loss_rec / len(val_loader)
        avg_loss_ridge = total_loss_ridge / len(val_loader)
        avg_loss_obs = total_loss_obs / len(val_loader)

        # Calculate the average validation accuracy
        predictions = torch.cat(predictions)
        labels_val = [lbl if lbl.dim() > 0 else lbl.unsqueeze(0) for lbl in labels_val]
        labels_val = torch.cat(labels_val)
        if dataset.num_classes > 2:
            # Handle multi-class classification case
            predictions = torch.argmax(predictions, dim=1)
        else:
            # Handle binary classification case
            predictions = torch.sigmoid(predictions) > 0.5
            predictions = predictions.squeeze()
        accuracy = (predictions == labels_val).sum().item() / len(labels_val)

        # Log the average validation loss and accuracy
        wandb.log(
            {
                "epoch": epoch,
                "val_loss": avg_loss,
                "val_loss_pred": avg_loss_pred,
                "val_loss_rec": avg_loss_rec,
                "val_loss_ridge": avg_loss_ridge,
                "val_loss_obs": avg_loss_obs,
                "val_acc": accuracy,
            }
        )
        if verbose:
            print("Validation Loss: {:.6f}".format(avg_loss))

        # Check for NaN loss and trigger early stopping
        if np.isnan(avg_loss):
            if verbose:
                print(f"NaN loss detected at epoch {epoch}. Triggering early stopping.")
            break

        # Check if the current accuracy is the best so far
        if accuracy > best_acc:
            best_acc = accuracy
            counter = 0
            # Save the best model to a temporary file (overwrite previous best)
            best_model_path = (
                f"models/saved/temp_best_{encoder_type}_{dataset_name}_{run_id}.pt"
            )
            torch.save(model.state_dict(), best_model_path)
            best_epoch = epoch
            if verbose:
                print(f"New best model at epoch {epoch}")
        else:
            counter += 1

        # Check if early stopping criteria is met
        if counter >= patience:
            if verbose:
                print("Early stopping at epoch", epoch)
            break

    # Load the best model for evaluation

    if best_model_path is not None:
        # Load the best model state
        model.load_state_dict(torch.load(best_model_path))
        # Clean up temporary file
        os.remove(best_model_path)
        if verbose:
            print(f"Loaded best model from epoch {best_epoch}")
    else:
        if verbose:
            print("No best model found, using the last model for evaluation")

    # Set the model to evaluation mode
    model.eval()

    # Validation
    predictions, hs_val, labels_val = [], [], []
    with torch.no_grad():
        for batch in tqdm(val_loader, desc="Validation", position=1, leave=False):

            # Move the inputs and labels to the device
            if encoder_type == "agcrn":
                # For AGCRN: batch.x is already (B, T, N, F)
                input = batch.x.to(device)
                edge_index = None
            else:
                # For other models: add batch dimension to TSL format
                input = batch.x.unsqueeze(0).to(device)
                edge_index = batch.edge_index

            label = batch.y.squeeze().to(device)
            if isinstance(edge_index, list):
                edge_index = torch.nested.nested_tensor(edge_index).to(device)
            elif isinstance(edge_index, torch.Tensor):
                edge_index = edge_index.to(device)
            else:
                edge_index = None

            label = batch.y.squeeze().to(device)

            # Get sequence lengths for this batch before padding
            if hasattr(batch, "last_timestep_mask"):
                last_timestep_mask = batch.last_timestep_mask.to(device)
            else:
                last_timestep_mask = None

            # Get node mask for this batch
            if hasattr(batch, "real_node_mask"):
                real_node_mask = batch.real_node_mask.to(device)
            else:
                real_node_mask = None

            # Forward pass
            x, h, x_rec, h_rec = model(
                input,
                edge_index,
                edge_weight=None,
                batch=batch.batch.to(device),
                last_timestep_mask=last_timestep_mask,
                real_node_mask=real_node_mask,
            )
            predictions.append(x)

            # Collect hidden states for plots:
            # sum over nodes for each batch and timestep
            if plot:
                if batch.batch is not None and encoder_type != "agcrn":
                    batch_size = batch.batch.max().item() + 1
                    batch_hidden_states = [
                        h[0, :, batch.batch == n, :].sum(-2) for n in range(batch_size)
                    ]
                    hs_val.extend(batch_hidden_states)
                elif batch.batch is not None and encoder_type == "agcrn":
                    # For AGCRN: h has shape (B, T, N, F), sum over real nodes for each batch sample
                    real_node_mask = (
                        batch.real_node_mask.unsqueeze(1).unsqueeze(-1).to(device)
                    )  # (B, 1, N, 1)
                    batch_hidden_states = (h * real_node_mask).sum(dim=2)  # (B, T, F)
                else:
                    hs_val.append(h[0].sum(-2))

            labels_val.append(label)

    # Compute classification accuracy
    predictions = torch.cat(predictions)
    labels_val = [lbl if lbl.dim() > 0 else lbl.unsqueeze(0) for lbl in labels_val]
    labels_val = torch.cat(labels_val)
    if dataset.num_classes > 2:
        # Handle multi-class classification case
        predictions = torch.argmax(predictions, dim=1)
    else:
        # Handle binary classification case
        predictions = torch.sigmoid(predictions) > 0.5
        predictions = predictions.squeeze()
    accuracy = (predictions == labels_val).sum().item() / len(labels_val)
    wandb.log({"acc": accuracy})

    if verbose:
        print("Accuracy: {:.4f}".format(accuracy))

    # Save the final model state with config and accuracy
    final_model_path = f"models/saved/{encoder_type}_{dataset_name}_{run_id}_.pt"
    final_config_path = f"models/saved/{encoder_type}_{dataset_name}_{run_id}_.yaml"
    # Save config as YAML
    OmegaConf.save(cfg, final_config_path)
    # Save model state dict, accuracy, and run_id
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "accuracy": accuracy,
            "run_id": run_id,
        },
        final_model_path,
    )
    if verbose:
        print("Final model saved to {}".format(final_model_path))

    if plot:
        # Perform PCA on the hidden states
        # Train states
        hs = torch.stack(hs)  # shape [batch, time, hidden_size]
        hs = hs.cpu().numpy()
        # Validation states
        hs_val = torch.stack(hs_val)  # shape [batch, time, hidden_size]
        hs_val = hs_val.cpu().numpy()

        # Dimensionality reduction
        dim_red = cfg.dim_red_plot
        pca = PCA(n_components=dim_red)

        if np.isnan(hs_val).any():
            warnings.warn("NaN values in the hidden states")
        elif labels_val.dim() > 1:
            warnings.warn("Histogram not implemented for more than two classes")
        else:
            hs_val_red = pca.fit_transform(hs_val.reshape(-1, hs_val.shape[-1]))
            hs_val_red = rearrange(
                hs_val_red,
                "(b t) f -> b t f",
                b=hs_val.shape[0],
                t=hs_val.shape[1],
                f=dim_red,
            )

            # Plots

            # Plot covariance matrix of reduced states
            fig, ax = plt.subplots()
            cov = ax.imshow(
                hs_val_red.reshape(-1, hs_val_red.shape[-1]).T
                @ hs_val_red.reshape(-1, hs_val_red.shape[-1]),
                cmap="viridis",
            )
            plt.colorbar(cov, ax=ax)
            wandb.log({"cov_img": wandb.Image(fig)})
            plt.close(fig)

            # Plot state distribution of the first 2 PCA components
            idx0 = labels_val.cpu().numpy() == 0
            idx1 = labels_val.cpu().numpy() == 1
            label_0 = hs_val_red[idx0, -1, :2]
            label_1 = hs_val_red[idx1, -1, :2]

            # Create a figure with two subplots
            fig, axs = plt.subplots(1, 2, figsize=(10, 4))
            # Plot the first histogram in the first subplot
            hist0 = axs[0].hist2d(
                label_0[:, 0], label_0[:, 1], bins=10, cmap="Blues", alpha=0.6
            )
            axs[0].set_xlabel("PC 0")
            axs[0].set_ylabel("PC 1")
            axs[0].set_title("2D Histogram - Label 0")
            plt.colorbar(hist0[3], ax=axs[0])
            # Plot the second histogram in the second subplot
            hist1 = axs[1].hist2d(
                label_1[:, 0], label_1[:, 1], bins=10, cmap="Reds", alpha=0.6
            )
            axs[1].set_xlabel("PC 0")
            axs[1].set_ylabel("PC 1")
            axs[1].set_title("2D Histogram - Label 1")
            plt.colorbar(hist1[3], ax=axs[1])
            # Adjust the spacing between subplots
            plt.tight_layout()

            wandb.log({"hist_PC_img": wandb.Image(fig)})
            plt.close(fig)

    # End profiling
    pr.disable()
    if verbose:
        with open("profile_results.prof", "w") as f:
            ps = pstats.Stats(pr, stream=f).sort_stats("cumulative")
            ps.dump_stats("profile_results.prof")

    return run_id


@hydra.main(version_base=None, config_path="configs", config_name="train")
def train_model(cfg: DictConfig) -> None:
    """Hydra entry point for training."""
    encoder_type = cfg.model.encoder_type
    dataset_name = cfg.dataset.name

    # If not a sweep, try to load saved config and merge CLI overrides
    if not is_sweep():
        import glob
        import re
        files = glob.glob(f"models/saved/{encoder_type}_{dataset_name}_*.yaml")
        files = [f for f in files if not re.search(r"_\.yaml$", f)]
        if len(files) == 1:
            config_file = files[0]
            # 1. Start from defaults (cfg)
            # 2. Override with saved config
            config_cfg = OmegaConf.load(config_file)
            cfg = OmegaConf.merge(cfg, config_cfg)

            # 3. Override only with CLI arguments
            from hydra.core.hydra_config import HydraConfig
            overrides = HydraConfig.get().overrides.task
            filtered_overrides = [
                o
                for o in overrides
                if "=" in o and o.split("=")[0] not in ["dataset", "model"]
            ]
            override_cfg = OmegaConf.from_dotlist(filtered_overrides)
            cfg = OmegaConf.merge(cfg, override_cfg)

            print(f"Loaded configs from: {config_file}")

    # Get run_id from config if provided (e.g., from run_experiment.py)
    run_id = (
        cfg.wandb.run_id if hasattr(cfg.wandb, "run_id") and cfg.wandb.run_id else None
    )

    run_training(cfg, run_id)


if __name__ == "__main__":
    train_model()