eval.py

import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import psutil
import csv
import os
import time
import gc 

import torch
import torch.multiprocessing as mp

from tgb.linkproppred.evaluate import Evaluator
from tqdm import tqdm

from tgm import DGraph, DGBatch
from tgm.constants import (
    METRIC_TGB_LINKPROPPRED,
    PADDED_NODE_ID,
    RECIPE_TGB_LINK_PRED,
)
from tgm.data import DGData, DGDataLoader
from tgm.hooks import DeduplicationHook, RecencyNeighborHook, RecipeRegistry, HookManager, DGHook
from tgm.nn import LinkPredictor, TGNMemory
from tgm.nn.encoder.tgn import (
    GraphAttentionEmbedding,
    IdentityMessage,
    LastAggregator,
)
from tgm.util.seed import seed_everything

parser = argparse.ArgumentParser(
    description='TGN Multi-Precision Benchmark',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument('--seed', type=int, default=1337, help='random seed to use')
parser.add_argument('--dataset', type=str, default='tgbl-wiki', help='Dataset name')
parser.add_argument('--bsize', type=int, default=200, help='batch size')
parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='torch device')
parser.add_argument('--epochs', type=int, default=30, help='number of epochs')
parser.add_argument('--lr', type=float, default=0.0001, help='learning rate')
parser.add_argument('--time-dim', type=int, default=100, help='time encoding dimension')
parser.add_argument('--embed-dim', type=int, default=100, help='attention dimension')
parser.add_argument('--memory-dim', type=int, default=100, help='memory dimension')
parser.add_argument(
    '--n-nbrs',
    type=int,
    nargs='+',
    default=[10],
    help='num sampled nbrs at each hop',
)
parser.add_argument('--precision', type=str, choices=['fp32', 'fp16', 'int8'], default='int8', help='Resting data precision')
parser.add_argument('--out-csv', type=str, default=None, help='Optional specific path to write CSV logs')
parser.add_argument('--patience', type=int, default=5, help='Early stopping patience (0 = disabled)')
parser.add_argument('--results-dir', type=str, default='results', help='Directory to store all CSV results')

args = parser.parse_args()

# Create results directory
os.makedirs(args.results_dir, exist_ok=True)

# Concurrency safety: Auto-generate unique CSV name if not provided
if args.out_csv is None:
    args.out_csv = os.path.join(args.results_dir, f'results_{args.dataset}_{args.precision}.csv')

print(args.out_csv)

# ==========================================
# 1. The Multi-Precision Hook
# ==========================================
class PrecisionCastingHook(DGHook):
    """Dequantizes/upcasts resting-precision edge features back to FP32.

    Handles both the main batch edge features (edge_x) and the
    neighbor edge features (nbr_edge_x) produced by RecencyNeighborHook.
    """
    requires = {'edge_x'}
    produces = {'edge_x'}

    def __init__(self):
        super().__init__()

    def __call__(self, dg: DGraph, batch: DGBatch) -> DGBatch:
        if batch.edge_x is not None:
            batch.edge_x = self._cast(batch.edge_x)

        if hasattr(batch, 'nbr_edge_x') and batch.nbr_edge_x is not None:
            batch.nbr_edge_x = [
                self._cast(x) if x is not None else x
                for x in batch.nbr_edge_x
            ]
        return batch

    @staticmethod
    def _cast(tensor: torch.Tensor) -> torch.Tensor:
        if tensor.dtype == torch.quint8:
            return tensor.dequantize()
        elif tensor.dtype == torch.float16:
            return tensor.float()
        return tensor

    def reset_state(self) -> None:
        pass

# ==========================================
# 2. Tracking Utility
# ==========================================
def get_memory_stats():
    cpu_ram = psutil.Process().memory_info().rss / (1024**2)
    vram = torch.cuda.max_memory_allocated() / (1024**2) if torch.cuda.is_available() else 0.0
    return cpu_ram, vram

def get_payload_mb(tensor: torch.Tensor) -> float:
    if tensor is None:
        return 0.0
    # Tensor.element_size() returns bytes per element (1 for int8, 2 for fp16, 4 for fp32)
    return (tensor.nelement() * tensor.element_size()) / (1024**2)

def log_to_csv(filepath, epoch, dataset, precision, loss, val_mrr, test_mrr, train_time, cpu_mb, vram_mb, total_payload):
    file_exists = os.path.isfile(filepath)
    with open(filepath, mode='a', newline='') as f:
        writer = csv.writer(f)
        if not file_exists:
            writer.writerow(['Dataset', 'Precision', 'Epoch', 'Loss', 'Val_MRR', 'Test_MRR', 'Train_Time_s', 'Peak_CPU_MB', 'Peak_VRAM_MB', "Total_Payload_MB"])
        writer.writerow([dataset, precision, epoch, f"{loss:.4f}", f"{val_mrr:.4f}", f"{test_mrr:.4f}", f"{train_time:.2f}", f"{cpu_mb:.2f}", f"{vram_mb:.2f}", f"{total_payload:.2f}"])

# ==========================================
# 3. Training & Eval Functions
# ==========================================
def train(
    loader: DGDataLoader,
    memory: nn.Module,
    encoder: nn.Module,
    decoder: nn.Module,
    opt: torch.optim.Optimizer,
) -> float:
    memory.train()
    encoder.train()
    decoder.train()
    total_loss = 0
    num_batches = 0

    for batch in tqdm(loader):
        opt.zero_grad()

        nbr_nodes = batch.nbr_nids[0].flatten()
        nbr_mask = nbr_nodes != PADDED_NODE_ID

        num_nbrs = len(nbr_nodes) // (
            len(batch.edge_src) + len(batch.edge_dst) + len(batch.neg)
        )

        src_nodes = torch.cat(
            [
                batch.edge_src.repeat_interleave(num_nbrs),
                batch.edge_dst.repeat_interleave(num_nbrs),
                batch.neg.repeat_interleave(num_nbrs),
            ]
        )
        nbr_edge_index = torch.stack(
            [
                batch.global_to_local(src_nodes[nbr_mask]),
                batch.global_to_local(nbr_nodes[nbr_mask]),
            ]
        ).to(dtype=torch.int64)

        nbr_edge_time = batch.nbr_edge_time[0].flatten()[nbr_mask]
        nbr_edge_x = batch.nbr_edge_x[0].flatten(0, -2).float()[nbr_mask]

        z, last_update = memory(batch.unique_nids)
        z = encoder(z, last_update, nbr_edge_index, nbr_edge_time, nbr_edge_x)

        inv_src = batch.global_to_local(batch.edge_src)
        inv_dst = batch.global_to_local(batch.edge_dst)
        inv_neg = batch.global_to_local(batch.neg)
        pos_out = decoder(z[inv_src], z[inv_dst])
        neg_out = decoder(z[inv_src], z[inv_neg])

        loss = F.binary_cross_entropy_with_logits(pos_out, torch.ones_like(pos_out))
        loss += F.binary_cross_entropy_with_logits(neg_out, torch.zeros_like(neg_out))

        memory.update_state(
            batch.edge_src, batch.edge_dst, batch.edge_time, batch.edge_x.float()
        )

        loss.backward()
        opt.step()
        total_loss += float(loss)
        num_batches += 1

        memory.detach()

    return total_loss / max(num_batches, 1)

@torch.no_grad()
def evaluate(
    loader: DGDataLoader,
    memory: nn.Module,
    encoder: nn.Module,
    decoder: nn.Module,
    evaluator: Evaluator,
) -> float:
    memory.eval()
    encoder.eval()
    decoder.eval()
    perf_list = []

    for batch in tqdm(loader):
        nbr_nodes = batch.nbr_nids[0].flatten()
        nbr_mask = nbr_nodes != PADDED_NODE_ID

        num_nbrs = len(nbr_nodes) // (
            len(batch.edge_src) + len(batch.edge_dst) + len(batch.neg)
        )
        src_nodes = torch.cat(
            [
                batch.edge_src.repeat_interleave(num_nbrs),
                batch.edge_dst.repeat_interleave(num_nbrs),
                batch.neg.repeat_interleave(num_nbrs),
            ]
        )
        nbr_edge_index = torch.stack(
            [
                batch.global_to_local(src_nodes[nbr_mask]),
                batch.global_to_local(nbr_nodes[nbr_mask]),
            ]
        ).to(dtype=torch.int64)
        
        nbr_edge_time = batch.nbr_edge_time[0].flatten()[nbr_mask]
        nbr_edge_x = batch.nbr_edge_x[0].flatten(0, -2).float()[nbr_mask]

        z, last_update = memory(batch.unique_nids)
        z = encoder(z, last_update, nbr_edge_index, nbr_edge_time, nbr_edge_x)

        for idx, neg_batch in enumerate(batch.neg_batch_list):
            dst_ids = torch.cat([batch.edge_dst[idx].unsqueeze(0), neg_batch])
            src_ids = batch.edge_src[idx].repeat(len(dst_ids))

            inv_src = batch.global_to_local(src_ids)
            inv_dst = batch.global_to_local(dst_ids)
            y_pred = decoder(z[inv_src], z[inv_dst]).sigmoid()

            input_dict = {
                'y_pred_pos': y_pred[0],
                'y_pred_neg': y_pred[1:],
                'eval_metric': [METRIC_TGB_LINKPROPPRED],
            }
            perf_list.append(evaluator.eval(input_dict)[METRIC_TGB_LINKPROPPRED])

        memory.update_state(
            batch.edge_src, batch.edge_dst, batch.edge_time, batch.edge_x.float()
        )

    return float(np.mean(perf_list))

# ==========================================
# 4. Main Execution
# ==========================================
if __name__ == "__main__":
    mp.set_start_method('spawn', force=True)
    seed_everything(args.seed)
    
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()

    evaluator = Evaluator(name=args.dataset)
    full_data = DGData.from_tgb(args.dataset)
    train_data, val_data, test_data = full_data.split()

    print(f"\n--- Initializing Experiment: {args.dataset} | Precision: {args.precision.upper()} ---")
    print(f"    Train edges: {len(train_data.edge_x)} | Val edges: {len(val_data.edge_x)} | Test edges: {len(test_data.edge_x)}")

    if args.precision == 'int8':
        fp32_tensor = full_data.edge_x
        f_min, f_max = fp32_tensor.min().item(), fp32_tensor.max().item()
        scale = max((f_max - f_min) / 255.0, 1e-8)
        zero_point = int(round(0 - (f_min / scale)))
        zero_point = max(0, min(255, zero_point))
        print(f"Global Quantization -> Scale: {scale:.6f}, Zero-Point: {zero_point}")
        
        train_data.edge_x = torch.quantize_per_tensor(train_data.edge_x, scale, zero_point, dtype=torch.quint8)
        val_data.edge_x = torch.quantize_per_tensor(val_data.edge_x, scale, zero_point, dtype=torch.quint8)
        test_data.edge_x = torch.quantize_per_tensor(test_data.edge_x, scale, zero_point, dtype=torch.quint8)

    elif args.precision == 'fp16':
        train_data.edge_x = train_data.edge_x.to(torch.float16)
        val_data.edge_x = val_data.edge_x.to(torch.float16)
        test_data.edge_x = test_data.edge_x.to(torch.float16)

    # garbage collection if compressed
    total_num_nodes = full_data.num_nodes
    
    # 3. Force Python to empty the garbage immediately
    if args.precision != 'fp32':
        import gc
        del full_data.edge_x
        del full_data
        gc.collect()

        if args.precision == 'int8':
            del fp32_tensor

    p_train = get_payload_mb(train_data.edge_x)
    p_val = get_payload_mb(val_data.edge_x)
    p_test = get_payload_mb(test_data.edge_x)
    total_payload = p_train + p_val + p_test
    print(f"True Tensor Payload Size: {total_payload:.2f} MB")
    
    resting_cpu_ram, _ = get_memory_stats()
    print(f"Resting CPU RAM after conversion: {resting_cpu_ram:.2f} MB")

    train_dg = DGraph(train_data, device=args.device)
    val_dg = DGraph(val_data, device=args.device)
    test_dg = DGraph(test_data, device=args.device)

    nbr_hook = RecencyNeighborHook(
        num_nbrs=args.n_nbrs,
        num_nodes=total_num_nodes,
        seed_nodes_keys=['edge_src', 'edge_dst', 'neg'],
        seed_times_keys=['edge_time', 'edge_time', 'neg_time'],
    )

    hm = RecipeRegistry.build(
        RECIPE_TGB_LINK_PRED, dataset_name=args.dataset, train_dg=train_dg
    )
    train_key, val_key, test_key = hm.keys

    hm.register_shared(nbr_hook)
    hm.register_shared(DeduplicationHook())
    hm.register_shared(PrecisionCastingHook())

    train_loader = DGDataLoader(train_dg, args.bsize, hook_manager=hm)
    val_loader = DGDataLoader(val_dg, args.bsize, hook_manager=hm)
    test_loader = DGDataLoader(test_dg, args.bsize, hook_manager=hm)

    memory = TGNMemory(
        total_num_nodes,
        test_dg.edge_x_dim,
        args.memory_dim,
        args.time_dim,
        message_module=IdentityMessage(test_dg.edge_x_dim, args.memory_dim, args.time_dim),
        aggregator_module=LastAggregator(),
    ).to(args.device)
    
    encoder = GraphAttentionEmbedding(
        in_channels=args.memory_dim,
        out_channels=args.embed_dim,
        msg_dim=test_dg.edge_x_dim,
        time_enc=memory.time_enc,
    ).to(args.device)
    
    decoder = LinkPredictor(node_dim=args.embed_dim, hidden_dim=args.embed_dim).to(args.device)
    
    opt = torch.optim.Adam(
        set(memory.parameters()) | set(encoder.parameters()) | set(decoder.parameters()),
        lr=args.lr,
    )

    best_val = 0.0
    patience_counter = 0

    for epoch in range(1, args.epochs + 1):
        start_time = time.time()
        
        memory.reset_state()
        hm.reset_state()

        with hm.activate(train_key):
            loss = train(train_loader, memory, encoder, decoder, opt)

        train_time = time.time() - start_time

        with hm.activate(val_key):
            val_mrr = evaluate(val_loader, memory, encoder, decoder, evaluator)

        test_mrr = 0.0
        if val_mrr > best_val:
            best_val = val_mrr
            patience_counter = 0
            with hm.activate(test_key):
                test_mrr = evaluate(test_loader, memory, encoder, decoder, evaluator)
        else:
            patience_counter += 1

        peak_cpu, peak_vram = get_memory_stats()
        log_to_csv(args.out_csv, epoch, args.dataset, args.precision, loss, val_mrr, test_mrr, train_time, peak_cpu, peak_vram, total_payload)
        
        print(f"Epoch {epoch:02d} | Loss: {loss:.4f} | Val MRR: {val_mrr:.4f} | Test MRR: {test_mrr:.4f} | "
              f"Best Val: {best_val:.4f} | Peak RAM: {peak_cpu:.0f}MB | Peak VRAM: {peak_vram:.0f}MB | Time: {train_time:.1f}s")

        if args.patience > 0 and patience_counter >= args.patience:
            print(f"Early stopping at epoch {epoch} (no improvement for {args.patience} epochs)")
            break

    print(f"\n--- Finished: {args.dataset} | {args.precision.upper()} | Best Val MRR: {best_val:.4f} ---")