train.py

import argparse
import os
import time
from collections import defaultdict
from copy import deepcopy

import dgl
import numpy as np
import torch
from tqdm import tqdm

from data.dataset import load_temporal_knowledge_graph, build_dataloaders
from model.layers import MultiAspectEmbedding, TimeIntervalTransform
from model.model import (
    EmbeddingUpdater, Combiner, EdgeModel,
    InterEventTimeModel, DynamicGraphModel,
)
from utils.helpers import (
    set_seed, load_config, get_embedding,
    pack_checkpoint, EarlyStopping, setup_logger,
)
from evaluate import evaluate


def compute_loss(model, optimize, batch_G, static_emb,
                 dyn_entity_emb, dyn_relation_emb, cfg):
    device = cfg.runtime.device
    batch_G = batch_G.to(device)
    nids = batch_G.ndata['_ID'].long().cpu()

    static_e = static_emb.structural[nids].to(device)
    dynamic_e = dyn_entity_emb.structural[nids, -1].to(device)
    combined = model.combiner(static_e, dynamic_e)
    dyn_rel = dyn_relation_emb.structural[:, -1]

    loss_dict = {}
    if optimize in ("edge", "both"):
        edge_loss = model.edge_model(
            batch_G, combined, static_e, dynamic_e, dyn_rel,
        )
        loss_dict["edge"] = -edge_loss

    if optimize in ("time", "both"):
        time_loss = model.inter_event_time_model.log_prob_density(
            batch_G, dyn_entity_emb, static_emb, dyn_relation_emb,
            model.node_latest_event_time,
        )
        loss_dict["time"] = -time_loss

    return loss_dict


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="config/config.yaml")
    args = parser.parse_args()

    cfg = load_config(args.config)
    set_seed(cfg.runtime.seed)
    logger = setup_logger(cfg.runtime.log_dir)
    device = cfg.runtime.device

    logger.info(f"Loading graph: {cfg.data.graph_name}")
    G = load_temporal_knowledge_graph(cfg.data.graph_name, cfg.data.data_root)
    train_dl, val_dl, test_dl = build_dataloaders(G)
    logger.info(f"  Nodes={G.num_nodes()}, Edges={G.num_edges()}, "
                f"Relations={G.num_relations}, NodeTypes={G.num_node_types}")

    dim = cfg.model.embed_dim
    num_nodes = G.num_nodes()
    num_rels = G.num_relations

    node_latest_event_time = torch.zeros(
        num_nodes, num_nodes + 1, 2, dtype=cfg.runtime.inter_event_dtype
    )

    embedding_updater = EmbeddingUpdater(
        num_nodes, dim, dim, dim, node_latest_event_time,
        num_rels, cfg.model.rel_embed_dim,
        model_type=cfg.model.model_type,
        num_gconv_layers=cfg.model.num_gconv_layers,
        num_rnn_layers=cfg.model.num_rnn_layers,
        num_ntypes=G.num_node_types,
        num_heads=cfg.model.num_attn_heads,
        dropout=cfg.model.dropout,
        graph_name=cfg.data.graph_name,
    ).to(device)

    combiner = Combiner(dim, dim, cfg.model.combine_mode, dropout=cfg.model.dropout).to(device)
    edge_model = EdgeModel(num_nodes, num_rels, cfg.model.rel_embed_dim, combiner, dropout=cfg.model.dropout).to(device)

    time_transform = TimeIntervalTransform(log_transform=True)
    inter_event_time_model = InterEventTimeModel(
        dim, dim, num_rels, cfg.model.rel_embed_dim,
        cfg.model.num_mix_components, time_transform,
        cfg.model.inter_event_time_mode, dropout=cfg.model.dropout,
    ).to(device)

    model = DynamicGraphModel(
        embedding_updater, combiner, edge_model,
        inter_event_time_model, node_latest_event_time,
    ).to(device)

    n_rnn = cfg.model.num_rnn_layers
    static_emb = MultiAspectEmbedding(
        structural=get_embedding(num_nodes, dim, zero_init=False),
        temporal=get_embedding(num_nodes, dim, zero_init=False),
    )
    init_dyn_entity = MultiAspectEmbedding(
        structural=get_embedding(num_nodes, [n_rnn, dim], zero_init=True),
        temporal=get_embedding(num_nodes, [n_rnn, dim, 2], zero_init=True),
    )
    init_dyn_rel = MultiAspectEmbedding(
        structural=get_embedding(num_rels, [n_rnn, cfg.model.rel_embed_dim, 2], zero_init=True),
        temporal=get_embedding(num_rels, [n_rnn, cfg.model.rel_embed_dim, 2], zero_init=True),
    )

    params = list(model.parameters()) + [
        static_emb.structural, static_emb.temporal,
        init_dyn_entity.structural, init_dyn_entity.temporal,
        init_dyn_rel.structural, init_dyn_rel.temporal,
    ]
    edge_opt = torch.optim.AdamW(params, lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
    time_opt = torch.optim.AdamW(params, lr=cfg.train.lr, weight_decay=cfg.train.weight_decay)
    stopper = EarlyStopping(patience=cfg.train.patience)

    logger.info(f"Starting training for {cfg.train.epochs} epochs")
    for epoch in range(cfg.train.epochs):
        model.train()
        model.node_latest_event_time.zero_()
        node_latest_event_time.zero_()
        dyn_entity = init_dyn_entity
        dyn_rel = init_dyn_rel

        epoch_losses = defaultdict(list)
        batch_loss_accum = 0
        t0 = time.time()

        pbar = tqdm(train_dl, desc=f"[Epoch {epoch:03d}]")
        for bi, (prior_G, batch_G, cumul_G, batch_t) in enumerate(pbar):
            last = bi == len(train_dl) - 1

            loss_dict = compute_loss(model, cfg.train.optimize, batch_G, static_emb, dyn_entity, dyn_rel, cfg)
            batch_loss_accum += sum(loss_dict.values())
            for k, v in loss_dict.items():
                epoch_losses[k].append(v.item())

            if bi > 0 and (bi % cfg.train.rnn_truncate_every == 0 or last):
                batch_loss_accum.backward()
                batch_loss_accum = 0
                if cfg.train.optimize in ("edge", "both"):
                    edge_opt.step(); edge_opt.zero_grad()
                if cfg.train.optimize in ("time", "both"):
                    time_opt.step(); time_opt.zero_grad()
                torch.cuda.empty_cache()
                for emb in list(dyn_entity) + list(dyn_rel):
                    if emb is not None:
                        emb.detach_()

            dyn_entity, dyn_rel = model.embedding_updater(
                prior_G, batch_G, cumul_G, static_emb, dyn_entity, dyn_rel, device,
            )

        elapsed = time.time() - t0
        total_loss = sum(sum(v) for v in epoch_losses.values())
        logger.info(f"[Epoch {epoch:03d}] loss={total_loss:.4f}  time={elapsed:.1f}s")

        if epoch >= cfg.eval.eval_from and epoch % cfg.eval.eval_every == 0:
            val_metrics = evaluate(model, val_dl, G, static_emb, dyn_entity, dyn_rel, num_rels, cfg, split="Validation")
            score = val_metrics.get("MRR", -val_metrics.get("loss", 0))
            logger.info(f"  [Val] {val_metrics}")

            ckpt = pack_checkpoint(
                model, edge_opt, time_opt, static_emb,
                init_dyn_entity, dyn_entity, init_dyn_rel, dyn_rel,
                deepcopy(model.node_latest_event_time), None, cfg,
            )
            if stopper.step(score, ckpt):
                logger.info(f"  Early stop at epoch {epoch}!")
                break

    logger.info("Training complete.")


if __name__ == "__main__":
    main()