train_model.py

import os
import json
import time
import torch
import os.path as osp
from tqdm import tqdm
from datetime import datetime
from typing import Dict, List, Tuple, Any, Optional, Union 
from torch_geometric.nn import VGAE
from torch_geometric.loader import DataLoader
from torch_geometric.transforms import NormalizeFeatures
from utils import create_logger, Encoder, MLPDecoder, StandardGCN, IterativeGCN

# Set up variables
with open("config.json", "r") as file:
    configs = json.load(file)

DECODER = configs["decoders"][str(configs["decoder"])]
BATCH_SIZE = configs["batch_size"]
LEARNING_RATE = configs["learning_rate"]
EPOCHS = configs["epochs"]
LATENT_DIM = configs["latent_dim"]
GCN_HIDDEN_DIM = configs["gcn_hidden_dim"]
MLP_HIDDEN_DIM = configs["mlp_hidden_dim"]
ENCODER_HIDDEN_DIM = configs["encoder_hidden_dim"]
DROPOUT = configs["dropout"]
BETA = configs["beta"]
MLP_LAYERS = configs["mlp_layers"]
ITERATIONS = configs["iterations"]
DATA = configs["data"]

def make_dirs() -> None:
    """
    Create necessary directories for storing logs, models, and loss data.
    """
    dirs = [
        f"logs/{DATA.lower()}/{DECODER.lower()}",
        f"losses/{DATA.lower()}/epoch/{DECODER.lower()}",
        f"losses/{DATA.lower()}/batch/{DECODER.lower()}",
        f"models/{DATA.lower()}/{DECODER.lower()}"
    ]
    for dir in dirs:
        os.makedirs(dir, exist_ok=True)

def get_data_loader(batch_size: int = 32, shuffle: bool = True) -> DataLoader:
    """
    Create a DataLoader for the dataset from PyG data files (Dingo).
    
    Args:
        batch_size (int, optional): Number of samples per batch. Defaults to 32.
        shuffle (bool, optional): Whether to shuffle the dataset. Defaults to True.
        
    Returns:
        DataLoader: PyTorch Geometric DataLoader containing graph data.
    """
    # Create a data list from the pyg_data folder
    data_list = []
    for grid in os.scandir("pyg_data"):
        path_to_pt_file = osp.join(grid.path, "data.pt")
        data = torch.load(path_to_pt_file, weights_only=False)
        data_list.append(data)
    return DataLoader(data_list, batch_size=batch_size, shuffle=shuffle)

def get_graphpf_data_loader(batch_size: int = 32, shuffle: bool = True) -> DataLoader:
    """
    Create a DataLoader for the graphpf dataset.
    
    Args:
        batch_size (int, optional): Number of samples per batch. Defaults to 32.
        shuffle (bool, optional): Whether to shuffle the dataset. Defaults to True.
        
    Returns:
        DataLoader: PyTorch Geometric DataLoader containing graph data.
    """
    #TODO: ADD POWER DATA FOLDER IN THE REPO
    base_path = osp.join(os.getcwd(), "graphpf_data")
    list_dirs = [osp.join(base_path, dir, "train", "dataset.pt") for dir in os.listdir(base_path)]
    data_list = []
    for data_path in list_dirs:
        data = torch.load(data_path, weights_only=False)
        for d in data:
            d.x = torch.nan_to_num(d.x)
            d.edge_attr = torch.nan_to_num(d.edge_attr)
            assert not torch.isnan(d.edge_index).any(), f"Edge Index contains NaN value(s): {data_path}"
        data_list.extend(data)
    return DataLoader(data_list, batch_size=batch_size, shuffle=shuffle)

def train(
    model: VGAE, 
    optimizer: torch.optim.Optimizer, 
    loader: DataLoader, 
    device: torch.device, 
    epochs: int = 10, 
    beta: float = BETA
) -> Tuple[Dict[str, List[float]], Dict[str, List[float]]]:
    """
    Train the VGAE model.
    
    Args:
        model (VGAE): The variational graph autoencoder model.
        optimizer (torch.optim.Optimizer): Optimizer for training.
        loader (DataLoader): DataLoader containing training data.
        device (torch.device): Device to run training on (CPU or GPU).
        epochs (int, optional): Number of training epochs. Defaults to 10.
        beta (float, optional): Weight for KL divergence loss. Defaults to BETA.
        
    Returns:
        Tuple[Dict[str, List[float]], Dict[str, List[float]]]: 
            - Epoch-wise losses (reconstruction, KL divergence, total)
            - Batch-wise losses (reconstruction, KL divergence, total)
    """
    # Initialize loss tracking
    epoch_recon_losses = []
    epoch_kl_losses = []
    epoch_total_losses = []
    batch_recon_losses = []  
    batch_kl_losses = []     
    batch_total_losses = []
    
    # Training loop
    for epoch in range(1, epochs + 1):
        model.train()
        epoch_recon = 0.0
        epoch_kl = 0.0
        epoch_total = 0.0
        batch_count = 0
        
        with tqdm(loader, unit="batch", desc=f"Epoch {epoch}/{epochs}") as tepoch:
            for data in tepoch:
                # Prepare data
                data = data.to(device)
                if hasattr(data, "transform"):
                    data.transform = NormalizeFeatures()
                
                # Forward pass
                optimizer.zero_grad()
                z = model.encode(data.x, data.edge_index)
                recon_loss = model.recon_loss(z, data.edge_index)
                kl_loss = beta * model.kl_loss()
                total_loss = recon_loss + kl_loss
                
                # Backward pass
                total_loss.backward()
                optimizer.step()

                # Track batch losses
                batch_recon_losses.append(recon_loss.item())
                batch_kl_losses.append(kl_loss.item())
                batch_total_losses.append(total_loss.item())
                
                # Track epoch losses
                epoch_recon += recon_loss.item()
                epoch_kl += kl_loss.item()
                epoch_total += total_loss.item()
                batch_count += 1
                
                # Update progress bar
                tepoch.set_postfix({
                    'recon_loss': recon_loss.item(),
                    'kl_loss': kl_loss.item(),
                    'total_loss': total_loss.item()
                })
        
        # Calculate epoch averages
        avg_recon = epoch_recon / batch_count
        avg_kl = epoch_kl / batch_count
        avg_total = epoch_total / batch_count
        
        # Store epoch losses
        epoch_recon_losses.append(avg_recon)
        epoch_kl_losses.append(avg_kl)
        epoch_total_losses.append(avg_total)
        
        logger.info(f"\nEpoch {epoch} Summary:")
        logger.info(f"Recon Loss: {avg_recon:.4f} | KL Loss: {avg_kl:.4f} | Total Loss: {avg_total:.4f}")
    
    # Return all losses for visualization
    return {
        'recon_losses': epoch_recon_losses,
        'kl_losses': epoch_kl_losses,
        'total_losses': epoch_total_losses
    }, {
        'batch_recon_losses': batch_recon_losses,
        'batch_kl_losses': batch_kl_losses,
        'batch_total_losses': batch_total_losses
    }
    
def initialize_model(
    loader: DataLoader, 
    device: torch.device
) -> Tuple[VGAE, Encoder, Optional[Union[MLPDecoder, StandardGCN, IterativeGCN]]]:
    """
    Initialize the VGAE model with the specified encoder and decoder.
    
    Args:
        loader (DataLoader): DataLoader containing the dataset.
        device (torch.device): Device to place the model on.
        
    Returns:
        Tuple[VGAE, Encoder, Optional[Union[MLPDecoder, StandardGCN, IterativeGCN]]]: 
            - Initialized VGAE model
            - Encoder instance
            - Decoder instance (or None for default inner product decoder)
            
    Raises:
        ValueError: If an unknown decoder type is specified.
    """
    ENCODER_INPUT_FEATURES = loader.dataset[0].num_features
    encoder_hidden_dim = ENCODER_HIDDEN_DIM
    
    encoder = Encoder(
        input_features=ENCODER_INPUT_FEATURES, hidden_dim=encoder_hidden_dim, 
        latent_dim=LATENT_DIM, dropout=DROPOUT
    )
    # Initialize decoder based on type
    if DECODER.upper() == "DEFAULT":
        decoder = None
        logger.info("Using Inner Product decoder")
    elif DECODER.upper() == "MLP":
        decoder = MLPDecoder(latent_dim=LATENT_DIM, hidden_dim=MLP_HIDDEN_DIM,
                             num_layers=MLP_LAYERS, dropout=DROPOUT)
        logger.info(f"Using MLP decoder model")
    elif DECODER.upper() == "GCN":
        decoder = StandardGCN(
            latent_dim=LATENT_DIM, hidden_dim=GCN_HIDDEN_DIM, dropout=DROPOUT)
        logger.info(f"Using GCN decoder model")
    elif DECODER.upper() == "ITERATIVE_GCN":
        decoder = IterativeGCN(
            latent_dim=LATENT_DIM, hidden_dim=GCN_HIDDEN_DIM, dropout=DROPOUT,
            iterations=ITERATIONS)
        logger.info(f"Using Iterative GCN Decoder")
    else:
        raise ValueError(f"Unknown decoder type: {DECODER}")
    
    model = VGAE(encoder, decoder).to(device)
    
    return model, encoder, decoder

def save_model(model: VGAE) -> str:
    """
    Save the model parameters to a file.
    
    Args:
        model (VGAE): Trained VGAE model to save.
        
    Returns:
        str: Path where the model was saved.
    """
    save_path = f"models/{DATA.lower()}/{DECODER.lower()}/{TIMESTAMP}_{EPOCHS}.pt"
    torch.save({
        "state_dict": model.state_dict(),
        "config": {
            "decoder_type": DECODER,
            "encoder_config": {
                'hidden_dim': ENCODER_HIDDEN_DIM,
                'latent_dim': LATENT_DIM,
                'dropout': DROPOUT
            },
            "decoder_config": {
                'hidden_dim': MLP_HIDDEN_DIM if DECODER == "MLP" else GCN_HIDDEN_DIM if DECODER == "GCN" else None,
                'latent_dim': LATENT_DIM,
                'num_layers': MLP_LAYERS if DECODER == "MLP" else None,
                'dropout': DROPOUT
            },
        }
    }, save_path)
    return save_path

def calculate_training_time(elapsed_time: float) -> str:
    """
    Format elapsed training time into a human-readable string.
    
    Args:
        elapsed_time (float): Total training time in seconds.
        
    Returns:
        str: Formatted time string (e.g., "2h 30m 45s").
    """
    hours = int(elapsed_time // 3600)
    minutes = int((elapsed_time % 3600) // 60)
    seconds = int(elapsed_time % 60)
    return f"{hours}h {minutes}m {seconds}s"

if __name__ == "__main__":
    # Create dirs for saving files, if not already present
    make_dirs()
    
    TIMESTAMP = datetime.now().strftime("%Y%m%d_%H%M%S")
    # Set up the logger
    logger = create_logger(f"VGAE_{DECODER}", f"logs/{DATA}/{DECODER.lower()}/{TIMESTAMP}_{EPOCHS}.log")
    for key, val in configs.items():
        if key != "decoders":
            logger.info(f"{key}: {val}")

    # Set device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device: {device}")

    logger.info("Create data loader")
    if DATA.lower() == "dingo":
        loader = get_data_loader(batch_size=BATCH_SIZE)
    elif DATA.lower() == "graphpf":
        loader = get_graphpf_data_loader(batch_size=BATCH_SIZE)
    else:
        raise ValueError(f"Loader not created. Unknown data, please check config!")
    logger.info(f"Dataset used to create Data loader: {DATA.lower()}")
    logger.info(f"Data loader created with {len(loader.dataset)} graphs.")
    
    # Initialize model
    vgae_model, encoder, decoder = initialize_model(loader, device)
    logger.info(f"VGAE model\n{vgae_model}")
    
    
    # # Initialize optimizer
    logger.info(f"Using Adam optimizer with learning rate: {LEARNING_RATE}")
    optimizer = torch.optim.Adam(vgae_model.parameters(), lr=LEARNING_RATE)
    
    # Train the model
    logger.info("Training the model...")
    start_time = time.time()
    epoch_losses, batch_losses = train(vgae_model, optimizer, loader, device, epochs=EPOCHS, beta=BETA)
    elapsed_time = time.time() - start_time
    logger.info(f"Training completed in {calculate_training_time(elapsed_time)}")

    with open(f"losses/{DATA.lower()}/epoch/{DECODER.lower()}/{TIMESTAMP}_{EPOCHS}.json", "w") as f:
        json.dump(epoch_losses, f)
    with open(f"losses/{DATA.lower()}/batch/{DECODER.lower()}/{TIMESTAMP}_{EPOCHS}.json", "w") as f:
        json.dump(batch_losses, f)
    logger.info("Losses saved.")
    
    # Save the model
    save_path = save_model(vgae_model)
    logger.info(f"Model saved at {save_path}")