train.py

from __future__ import annotations
import argparse
from collections import defaultdict
from typing import Dict, Any
from pathlib import Path
import datetime
import os 

import torch
import torch.nn.functional as F
from torch_geometric.loader import DataLoader as GeoLoader
from torch_geometric.data import Dataset
import matplotlib.pyplot as plt


from Utils import plot_training_curves
from GraphMatching import BruteForceMatcher, BruteForceSampleMatcher, PyGMHungarianMatcher, permute_batch, permute_complete_dataset, permute_graphs_over_n
from Models import VGAE, build_encoder, PairNodeSkeletonDecoder, TransformerDecoder  # expects graph-level latent implemented in this module



def make_matcher(kind: str):
    k = kind.lower()
    if k == "optimal":
        return BruteForceMatcher()
    elif k.startswith("top"):
        return BruteForceSampleMatcher(int(k[3:]))
    else:
        raise ValueError("Unknown matcher kind. Use 'optimal', 'top10', 'top50', 'top100'.")


def train_epoch(model, loader, opt, matcher, beta_kl: float, device: torch.device, permutations):
    model.train()
    running = defaultdict(float)
    for data in loader:
        data = data.to(device)
        data=permute_batch(data,permutations,device)
        
        opt.zero_grad()
        data_rec, mu, logstd, _, pred_num = model(data)
        #num_nodes = 0.5 * (torch.sqrt(8 * torch.bincount(data.batch) + 1) - 1)
        
        #num_loss = 0.1 * F.mse_loss(pred_num, num_nodes)
        rec_batch, mu, logstd, _, _ = model(data)
        matched_loss, _ = matcher.match(data, data_rec)
        kl_l = model.kl_loss(mu, logstd)
        loss = matched_loss + beta_kl * kl_l #+ num_loss

        loss.backward()
        opt.step()
        
        #running["num_loss"] += num_loss.item()
        running["reconstruction_matched"] += matched_loss.item()
        running["kl_loss"] += kl_l.item()
        running["loss"] += loss.item()
    for k in running:
        running[k] /= max(1, len(loader))
    return running


@torch.no_grad()
def val_epoch(model, loader, matcher, beta_kl: float, device: torch.device):
    model.eval()
    running = defaultdict(float)
    for data in loader:
        data = data.to(device)
        data_rec, mu, logstd, _, pred_num = model(data)
        #num_nodes = 0.5 * (torch.sqrt(8 * torch.bincount(data.batch) + 1) - 1)
        
        #num_loss = 0.1 * F.mse_loss(pred_num, num_nodes)
        matched_loss, _ = matcher.match(data, data_rec)
        kl_l = model.kl_loss(mu, logstd)
        loss = matched_loss + beta_kl * kl_l
        
        #running["num_loss"] += num_loss.item()
        running["reconstruction_matched"] += matched_loss.item()
        running["kl_loss"] += kl_l.item()
        running["loss"] += loss.item()
    for k in running:
        running[k] /= max(1, len(loader))
    return running


def parse_args():
    p = argparse.ArgumentParser(description="Train VGAE (graph-level latent) with graph-matching recon loss")
    # Data
    p.add_argument("--root", type=str, default="data/edge_node_qm9")
    p.add_argument("--qm9-root", type=str, default="data/qm9")
    p.add_argument("--limit", type=int, default=None)
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--num-workers", type=int, default=0)
    # Encoder
    p.add_argument("--encoder-type", type=str, default="gin", choices=["gin","gine","gcn","graph","gat"])
    p.add_argument("--in-channels", type=int, default=6)
    p.add_argument("--encoder-layers", type=int, default=3)
    p.add_argument("--encoder-hidden", type=str, default="64")   # int or comma list
    p.add_argument("--encoder-out", type=int, default=128)
    p.add_argument("--norm", type=str, default="graph")
    p.add_argument("--dropout", type=float, default=0.1)
    p.add_argument("--act", type=str, default="leakyrelu")
    # VAE
    p.add_argument("--latent-dim", type=int, default=50)
    p.add_argument("--beta-kl", type=float, default=1.0)
    # Train
    p.add_argument("--epochs", type=int, default=50)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--matcher", type=str, default="top10")
    return p.parse_args()

from torch.utils.data import Subset

def subset_max_atoms(ds, max_atoms=9):
    keep = []
    dropped = 0
    for i in range(len(ds)):
        d = ds[i]
        n_atoms = int((d.x[:, 0] == 1).sum())  # indicator +1 for atoms
        if n_atoms <= max_atoms:
            keep.append(i)
        else:
            dropped += 1
    print(f"[filter] keeping {len(keep)}/{len(ds)} (dropped {dropped} with >{max_atoms} atoms)")
    return Subset(ds, keep)

class SavedGraphDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
        self.files = sorted(
            [f for f in os.listdir(root) if f.startswith("data_") and f.endswith(".pt")],
            key=lambda f: int(f.split("_")[1].split(".")[0])
        )

    def len(self):
        return len(self.files)

    def get(self, idx):
        return torch.load(os.path.join(self.root, self.files[idx]))

def main():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Using device:", device)

    ds = SavedGraphDataset("./data")
    split     = (20_000, 30_000)                   # standard QM9 split
    train, val, test = ds[:split[0]], ds[split[0]:split[1]], ds[split[1]:]
    
    train_ds = subset_max_atoms(train, 9)
    val_ds   = subset_max_atoms(val,   9)
    
    train_loader = GeoLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=4)
    
    permutations = [permute_graphs_over_n(i).to(device) for i in range(1,10)]
    
    val_ds = permute_complete_dataset(val_ds, permutations , device)
    val_loader   = GeoLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=4)

    # Encoder
    hidden = args.encoder_hidden
    if isinstance(hidden, str):
        hidden = [int(x) for x in hidden.split(",")] if "," in hidden else int(hidden)
    cfg: Dict[str, Any] = {
        "encoder_type": args.encoder_type,
        "in_channels": args.in_channels,
        "encoder_layers": args.encoder_layers,
        "encoder_hidden_dims": hidden,
        "encoder_output_dim": args.encoder_out,
        "encoder_norm": (None if args.norm.lower() == "none" else args.norm),
        "encoder_norm_kwargs": {},
        "act": args.act,
        "dropout": args.dropout,
        "heads": 4 if args.encoder_type == "gat" else 1,                                                   
        "v2": False,
    }
    encoder = build_encoder(cfg, device)
    decoder = TransformerDecoder(args.latent_dim, 128 ,8 , 5, "layer").to(device)
    #decoder = PairNodeSkeletonDecoder(latent_dim=args.latent_dim, hidden=256, max_atoms=9)
    model = VGAE(encoder=encoder, decoder=decoder, latent_level='graph', latent_dim=args.latent_dim, property_out_dim=1).to(device)


    opt = torch.optim.Adam(model.parameters(), lr=args.lr)
    matcher = make_matcher(args.matcher)
    #matcher = PyGMHungarianMatcher()
    ckpt_dir = Path("checkpoints") / datetime.datetime.now().strftime("%Y-%m-%d_%H-%M")
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    
    best_val = float("inf")
    for epoch in range(1, args.epochs + 1):
        tr = train_epoch(model, train_loader, opt, matcher, args.beta_kl, device, permutations)
        va = val_epoch(model, val_loader, matcher, args.beta_kl, device)
        print(f"Epoch {epoch:03d}  "
              f"train: loss {tr['loss']:.4f}  recon {tr['reconstruction_matched']:.4f}  KL {tr['kl_loss']:.4f} | "
              f"val: loss {va['loss']:.4f}  recon {va['reconstruction_matched']:.4f}  KL {va['kl_loss']:.4f}")
        if va["loss"] < best_val:
            best_val = va["loss"]
            torch.save(model.state_dict(), ckpt_dir / "best.pt")

    print("Training finished; best val loss:", best)
    

    curves = (
        "loss", "val_loss",
        "kl_loss", "val_kl_loss",
        "reconstruction_matched", "val_reconstruction_matched"
    )

    plot_training_curves(hist, keys=curves)
    plt.tight_layout()
    plt.savefig("training_curves.png", dpi=150)
    print("Training curves saved to training_curves.png")
    
    
if __name__ == "__main__":
    main()