train.py

"""
MCVRP GNN Training Script
=========================

This script trains a Graph Neural Network (GNN) to predict the probability 
of two nodes (stops) belonging to the same route in the MCVRP.

Usage Examples:
---------------
1. Default training (Sparse kNN graph k=20, new codex data, 50 epochs):
   python train.py

2. Experiment with a fully connected graph (k=200):
   python train.py --knn_k 200 --epochs 10 --output_dir checkpoints_full

3. Baseline experiment on old data:
   python train.py --labels_dir nsga_labels --output_dir checkpoints_old_labels

Available Parameters:
---------------------
  --labels_dir    Directory containing input labels (inside data/ directory, or absolute path). Default: nsga_codex_labels
  --output_dir    Directory to save trained models and plots. Default: checkpoints_50epochs
  --knn_k         Number of nearest neighbors. Set to 200 for a fully connected graph. Default: 20
  --epochs        Number of training epochs. Default: 50
  --batch_size    Training batch size. Default: 8
  --lr            Learning rate for Adam optimizer. Default: 1e-3
"""
import os
import argparse
import torch
import torch.nn as nn
from torch.optim import Adam
from torch_geometric.loader import DataLoader
from sklearn.metrics import roc_auc_score, average_precision_score
import matplotlib.pyplot as plt
from pathlib import Path
from tqdm import tqdm

from dataset import VRPLabelDataset
from gnn import VRPEdgePredictor

def train_model(args):
    DATA_DIR = Path(__file__).resolve().parent.parent / "data"
    
    # Handle absolute vs relative paths for labels_dir
    labels_dir_path = Path(args.labels_dir)
    if not labels_dir_path.is_absolute():
        labels_dir_path = DATA_DIR / args.labels_dir
        
    # Handle absolute vs relative paths for output_dir
    checkpoint_dir = Path(args.output_dir)
    if not checkpoint_dir.is_absolute():
        checkpoint_dir = Path(__file__).resolve().parent / args.output_dir
        
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"Loading VRP Label Dataset from {labels_dir_path}...")
    full_dataset = VRPLabelDataset(
        graphs_dir=str(DATA_DIR / "xset_graphs"),
        labels_dir=str(labels_dir_path),
        max_nodes=200,
        param_filter="100_0p50_0p20",
        seed_filter=None,
        knn_k=args.knn_k
    )
    
    if len(full_dataset) == 0:
        print("Dataset is empty. Ensure paths and data are correct.")
        return
        
    print(f"Total labeled graphs found: {len(full_dataset)}")
    
    # Train / Val Split (90/10) split by INSTANCE, not by label!
    # This prevents the GNN from simply memorizing specific graphs.
    unique_instances = list(set([item["instance"] for item in full_dataset.data_items]))
    unique_instances.sort()  # Maintain determinism
    
    import random
    rng = random.Random(42)
    rng.shuffle(unique_instances)
    
    train_count = int(0.9 * len(unique_instances))
    train_instances = set(unique_instances[:train_count])
    val_instances = set(unique_instances[train_count:])
    
    # Save the exact graphs used for the splits so the user can check them!
    with open(checkpoint_dir / "train_instances.txt", "w") as f:
        f.write("\n".join(sorted(list(train_instances))))
    with open(checkpoint_dir / "val_instances.txt", "w") as f:
        f.write("\n".join(sorted(list(val_instances))))
    
    train_indices = []
    val_indices = []
    for i, item in enumerate(full_dataset.data_items):
        if item["instance"] in train_instances:
            train_indices.append(i)
        else:
            val_indices.append(i)
            
    train_dataset = torch.utils.data.Subset(full_dataset, train_indices)
    val_dataset = torch.utils.data.Subset(full_dataset, val_indices)
    
    print(f"Split completed: {len(train_indices)} labels (Train) | {len(val_indices)} labels (Val Test)")
    
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)

    # Calculate class imbalance strategy for positive weight
    # We sample a few batches to get a decent estimate of the ratio
    print("Estimating positive class weight...")
    pos_samples = 0
    neg_samples = 0
    for data in DataLoader(train_dataset, batch_size=32, shuffle=True):
        pos_samples += (data.y == 1).sum().item()
        neg_samples += (data.y == 0).sum().item()
        break # One batch is enough for an estimate in similar graphs
        
    pos_weight = torch.tensor([neg_samples / max(1, pos_samples)], dtype=torch.float)
    print(f"Estimated Pos/Neg Ratio -> pos_weight = {pos_weight.item():.2f}")

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    model = VRPEdgePredictor(node_in_dim=5, edge_in_dim=2, hidden_dim=64, num_layers=4).to(device)
    optimizer = Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight.to(device))

    best_val_ap = 0.0
    history = {'train_loss': [], 'val_loss': [], 'val_roc': [], 'val_ap': []}
    
    for epoch in range(args.epochs):
        model.train()
        train_loss = 0.0
        
        # Training loop
        for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs} [Train]"):
            batch = batch.to(device)
            optimizer.zero_grad()
            
            logits = model(batch.x, batch.edge_index, batch.edge_attr)
            loss = criterion(logits, batch.y)
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            
            train_loss += loss.item()
            
        train_loss /= len(train_loader)
        
        # Validation Loop
        model.eval()
        val_loss = 0.0
        all_preds = []
        all_targets = []
        
        with torch.no_grad():
            for batch in tqdm(val_loader, desc=f"Epoch {epoch+1}/{args.epochs} [Val]"):
                batch = batch.to(device)
                logits = model(batch.x, batch.edge_index, batch.edge_attr)
                loss = criterion(logits, batch.y)
                val_loss += loss.item()
                
                probs = torch.sigmoid(logits)
                all_preds.append(probs.cpu())
                all_targets.append(batch.y.cpu())
                
        val_loss /= len(val_loader)
        
        all_preds = torch.cat(all_preds).numpy()
        all_targets = torch.cat(all_targets).numpy()
        
        try:
            val_roc = roc_auc_score(all_targets, all_preds)
            val_ap = average_precision_score(all_targets, all_preds)
        except ValueError:
            val_roc, val_ap = 0.0, 0.0
            
        print(f"Epoch {epoch+1} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val ROC: {val_roc:.4f} | Val AP: {val_ap:.4f}")
        
        # Record history
        history['train_loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['val_roc'].append(val_roc)
        history['val_ap'].append(val_ap)
        
        # Save best model based on Average Precision (Precision-Recall is better for imbalanced clustering)
        if val_ap > best_val_ap:
            best_val_ap = val_ap
            save_path = checkpoint_dir / "vrp_edge_predictor_best.pt"
            torch.save(model.state_dict(), save_path)
            print(f"  -> Saved new best model to {save_path}")

    # Plot Convergence History
    print("\nGenerating convergence plot...")
    plt.figure(figsize=(12, 5))
    
    # Subplot 1: Loss
    plt.subplot(1, 2, 1)
    plt.plot(range(1, args.epochs + 1), history['train_loss'], label='Train Loss (BCE)', marker='o')
    plt.plot(range(1, args.epochs + 1), history['val_loss'], label='Val Loss (BCE)', marker='o')
    plt.title('Loss over Epochs')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.grid(True)
    plt.legend()
    
    # Subplot 2: Metrics
    plt.subplot(1, 2, 2)
    plt.plot(range(1, args.epochs + 1), history['val_roc'], label='Val ROC-AUC', marker='s')
    plt.plot(range(1, args.epochs + 1), history['val_ap'], label='Val Average Precision', marker='s')
    plt.title('Validation Metrics over Epochs')
    plt.xlabel('Epoch')
    plt.ylabel('Score')
    plt.grid(True)
    plt.legend()
    
    plot_path = checkpoint_dir / "convergence.png"
    plt.tight_layout()
    plt.savefig(plot_path, dpi=150)
    plt.close()
    print(f"Convergence plot successfully saved to {plot_path}!")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train GNN for MCVRP Edge Prediction")
    parser.add_argument("--labels_dir", type=str, default="nsga_codex_labels", 
                        help="Directory containing label files (relative to data/ or absolute path). Default: nsga_codex_labels")
    parser.add_argument("--output_dir", type=str, default="checkpoints_50epochs", 
                        help="Directory to save model checkpoints and plots. Default: checkpoints_50epochs")
    parser.add_argument("--knn_k", type=int, default=20, 
                        help="Number of nearest neighbors for the graph structure. Use 200 for fully connected. Default: 20")
    parser.add_argument("--epochs", type=int, default=50, 
                        help="Number of training epochs. Default: 50")
    parser.add_argument("--batch_size", type=int, default=8, 
                        help="Batch size for training. Default: 8")
    parser.add_argument("--lr", type=float, default=1e-3, 
                        help="Learning rate for Adam optimizer. Default: 1e-3")
    
    args = parser.parse_args()
    
    import multiprocessing
    # Optional performance fix for PyG on windows
    torch.multiprocessing.set_sharing_strategy('file_system')
    
    print(f"Starting training with arguments: {vars(args)}")
    train_model(args)