train_gnn.py

import torch
import argparse
from models.gnn_model import SimpleGraphEncoder
from models.trainer import SimpleTrainer, VRPDataset
from data.generator import VRPDataGenerator

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--n_nodes', type=int, default=20)
    parser.add_argument('--capacity', type=float, default=30.0)
    parser.add_argument('--n_train', type=int, default=5000)
    parser.add_argument('--n_val', type=int, default=500)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--lr', type=float, default=1e-3)
    args = parser.parse_args()
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Device: {device}")
    
    print("Generating data...")
    train_gen = VRPDataGenerator(args.n_nodes, args.capacity, seed=1234)
    train_coords, train_demands, train_sols = train_gen.generate_with_solutions(args.n_train)
    
    val_gen = VRPDataGenerator(args.n_nodes, args.capacity, seed=5678)
    val_coords, val_demands, val_sols = val_gen.generate_with_solutions(args.n_val)
    
    train_dataset = VRPDataset(train_coords, train_demands, train_sols, args.capacity)
    val_dataset = VRPDataset(val_coords, val_demands, val_sols, args.capacity)
    
    print("Creating model...")
    model = SimpleGraphEncoder(embedding_dim=64, n_layers=5, n_heads=4)
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    print("Training...")
    trainer = SimpleTrainer(model, device, args.lr, args.capacity)
    trainer.train(train_dataset, val_dataset, args.epochs, args.batch_size)
    
    print("Done!")

if __name__ == '__main__':
    main()