gumbel_pretrain.py

"""Unified interface to all dynamic graph model experiments"""
import logging
import math
import os
import sys

import numpy as np
import torch
from sklearn.metrics import (accuracy_score, average_precision_score, f1_score,
                             roc_auc_score)
from tqdm import trange

from args_config import gumbel_args_config
from data_util import load_data, load_graph, load_label_data
from graph import NeighborFinder, make_label_data
from gumbel_alpha import GumbelGAN
from util import EarlyStopMonitor, RandEdgeSampler, set_logger

#import numba

# Argument and global variables
try:
    parser = gumbel_args_config()
    args = parser.parse_args()
except:
    parser.print_help()
    sys.exit(0)

# Arguments
if True:
    BATCH_SIZE = args.bs
    FREEZE = args.freeze
    NUM_NEIGHBORS = args.n_degree
    NUM_NEG = 1
    NUM_EPOCH = args.n_epoch
    NUM_HEADS = args.n_head
    DROP_OUT = args.drop_out
    GPU = args.gpu
    UNIFORM = args.uniform
    USE_TIME = args.time
    AGG_METHOD = args.agg_method
    ATTN_MODE = args.attn_mode
    SEQ_LEN = NUM_NEIGHBORS
    DATA = args.data
    TASK = args.task
    HARD = args.hard
    NUM_LAYER = args.n_layer
    LEARNING_RATE = args.lr
    NODE_DIM = args.node_dim
    TIME_DIM = args.time_dim

    # Model initialize
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    device = torch.device('cuda:{}'.format(GPU))

    import socket

    DEVICE_STR = f'{socket.gethostname()}-{device.index}'
    MODEL_SAVE_PATH = f'./sample_cache/{TASK}-{FREEZE}-{args.data}-gumbel-{HARD}.pth'

    def get_checkpoint_path(epoch):
        return f'./ckpt/{TASK}-{args.data}-gumbel-{HARD}-{epoch}.pth'


# set up logger
if True:
    logger = set_logger()
    logger.info(args)


def eval_one_epoch(hint, dps, src, dst, ts, label):
    with torch.no_grad():
        dps = dps.eval()
        TEST_BATCH_SIZE = BATCH_SIZE
        num_test_instance = len(src)
        num_test_batch = math.ceil(len(src) / TEST_BATCH_SIZE)
        scores = []

        for k in range(num_test_batch):
            s_idx = k * TEST_BATCH_SIZE
            e_idx = min(s_idx + TEST_BATCH_SIZE, num_test_instance)
            src_l_cut = src[s_idx:e_idx]
            dst_l_cut = dst[s_idx:e_idx]
            ts_l_cut = ts[s_idx:e_idx]
            prob_score = dps.forward(src_l_cut, dst_l_cut, ts_l_cut).sigmoid()
            scores.extend(list(prob_score.cpu().numpy()))
        pred_label = np.array(scores) > 0.5
        pred_prob = np.array(scores)
    return accuracy_score(label, pred_label), average_precision_score(
        label,
        pred_label), f1_score(label,
                              pred_label), roc_auc_score(label, pred_prob)


# Load data and train val test split
if True:
    logging.info("Begin dataset %s.", DATA)
    if TASK == "edge":
        edges, n_nodes, val_time, test_time = load_graph(DATA)
        g_df = edges[["from_node_id", "to_node_id", "timestamp"]].copy()
        g_df["idx"] = np.arange(1, len(g_df) + 1)
        g_df.columns = ["u", "i", "ts", "idx"]
    elif TASK == "node":
        edges, nodes = load_data(DATA, "format")
        n_nodes = len(nodes)
        # padding node is 0, so add 1 here.
        id2idx = {row.node_id: row.id_map + 1 for row in nodes.itertuples()}
        edges["from_node_id"] = edges["from_node_id"].map(id2idx)
        edges["to_node_id"] = edges["to_node_id"].map(id2idx)
        g_df = edges[["from_node_id", "to_node_id", "timestamp"]].copy()
        g_df["idx"] = np.arange(1, len(edges) + 1)
        g_df.columns = ["u", "i", "ts", "idx"]
        val_time, test_time = list(np.quantile(g_df.ts, [0.70, 0.85]))

    if len(edges.columns) > 4:
        e_feat = edges.iloc[:, 4:].to_numpy()
        padding = np.zeros((1, e_feat.shape[1]))
        e_feat = np.concatenate((padding, e_feat))
    else:
        e_feat = np.zeros((len(g_df) + 1, NODE_DIM))

    if args.freeze:
        n_feat = np.zeros((n_nodes + 1, NODE_DIM))
    else:
        bound = np.sqrt(6 / (2 * NODE_DIM))
        n_feat = np.random.uniform(-bound, bound, (n_nodes + 1, NODE_DIM))

    src_l = g_df.u.values
    dst_l = g_df.i.values
    e_idx_l = g_df.idx.values
    ts_l = g_df.ts.values

    max_src_index = src_l.max()
    max_idx = max(src_l.max(), dst_l.max())

# set train, validation, test datasets
if True:
    valid_train_flag = (ts_l < val_time)

    train_src_l = src_l[valid_train_flag]
    train_dst_l = dst_l[valid_train_flag]
    train_ts_l = ts_l[valid_train_flag]
    train_e_idx_l = e_idx_l[valid_train_flag]

    train_rand_sampler = RandEdgeSampler(train_src_l, train_dst_l)
    val_rand_sampler = RandEdgeSampler(src_l, dst_l)
    test_rand_sampler = RandEdgeSampler(src_l, dst_l)

# set validation, test datasets
if True:
    if TASK == "edge":
        _, val_data, test_data = load_label_data(dataset=DATA)

        val_src_l = val_data.u.values
        val_dst_l = val_data.i.values
        val_ts_l = val_data.ts.values
        val_label_l = val_data.label.values

        test_src_l = test_data.u.values
        test_dst_l = test_data.i.values
        test_ts_l = test_data.ts.values
        test_label_l = test_data.label.values
    elif TASK == "node":
        # select validation and test dataset
        valid_val_flag = (ts_l <= test_time) * (ts_l > val_time)
        valid_test_flag = ts_l > test_time

        val_src_l, val_dst_l, val_ts_l, val_label_l = make_label_data(
            src_l, dst_l, ts_l, valid_val_flag, test_rand_sampler)
        test_src_l, test_dst_l, test_ts_l, test_label_l = make_label_data(
            src_l, dst_l, ts_l, valid_test_flag, test_rand_sampler)
    else:
        raise NotImplementedError(TASK)

# Initialize the data structure for graph and edge sampling
# build the graph for fast query
adj_list = [[] for _ in range(max_idx + 1)]
for src, dst, eidx, ts in zip(train_src_l, train_dst_l, train_e_idx_l,
                              train_ts_l):
    adj_list[src].append((dst, eidx, ts))
    adj_list[dst].append((src, eidx, ts))
train_ngh_finder = NeighborFinder(adj_list, uniform=True)

# full graph with all the data for the test and validation purpose
full_adj_list = [[] for _ in range(max_idx + 1)]
for src, dst, eidx, ts in zip(src_l, dst_l, e_idx_l, ts_l):
    full_adj_list[src].append((dst, eidx, ts))
    full_adj_list[dst].append((src, eidx, ts))
full_ngh_finder = NeighborFinder(full_adj_list, uniform=True)

dps = GumbelGAN(train_ngh_finder,
                n_feat,
                e_feat,
                n_feat_freeze=args.freeze,
                num_layers=NUM_LAYER,
                use_time=USE_TIME,
                agg_method=AGG_METHOD,
                attn_mode=ATTN_MODE,
                seq_len=SEQ_LEN,
                n_head=NUM_HEADS,
                drop_out=DROP_OUT,
                node_dim=NODE_DIM,
                time_dim=TIME_DIM,
                hard=HARD,
                num_neighbors=NUM_NEIGHBORS)
optimizer = torch.optim.Adam(dps.parameters(), lr=LEARNING_RATE)
criterion = torch.nn.BCELoss()
dps = dps.to(device)

num_instance = len(train_src_l)
num_batch = math.ceil(num_instance / BATCH_SIZE)

logger.info('num of training instances: {}'.format(num_instance))
logger.info('num of batches per epoch: {}'.format(num_batch))
idx_list = np.arange(num_instance)
np.random.shuffle(idx_list)

# Pretraining a 1-layer GumbelGAN with uniform sampling.
if True:
    logger.info("Pretraining a 1-layer GumberlGAN with uniform sampling.")
    early_stopper = EarlyStopMonitor()
    epoch_bar = trange(NUM_EPOCH)
    for epoch in epoch_bar:
        # Training
        # training use only training graph
        dps.ngh_finder = train_ngh_finder
        np.random.shuffle(idx_list)
        batch_bar = trange(num_batch)
        for k in batch_bar:
            s_idx = k * BATCH_SIZE
            e_idx = min(num_instance - 1, s_idx + BATCH_SIZE)
            src_l_cut = train_src_l[s_idx:e_idx]
            dst_l_cut = train_dst_l[s_idx:e_idx]
            ts_l_cut = train_ts_l[s_idx:e_idx]
            size = len(src_l_cut)
            src_l_fake, dst_l_fake = train_rand_sampler.sample(size)

            with torch.no_grad():
                pos_label = torch.ones(size, dtype=torch.float, device=device)
                neg_label = torch.zeros(size, dtype=torch.float, device=device)

            optimizer.zero_grad()
            dps = dps.train()
            pos_prob, neg_prob = dps.contrast(src_l_cut,
                                              dst_l_cut,
                                              dst_l_fake,
                                              ts_l_cut,
                                              NUM_NEIGHBORS,
                                              gumbel=False)

            loss = criterion(pos_prob, pos_label)
            loss += criterion(neg_prob, neg_label)

            loss.backward()
            optimizer.step()
            # get training results
            with torch.no_grad():
                dps = dps.eval()
                pred_score = np.concatenate([(pos_prob).cpu().detach().numpy(),
                                             (neg_prob).cpu().detach().numpy()
                                             ])
                pred_label = pred_score > 0.5
                true_label = np.concatenate([np.ones(size), np.zeros(size)])
                acc = accuracy_score(true_label, pred_label)
                ap = average_precision_score(true_label, pred_label)
                f1 = f1_score(true_label, pred_label)
                auc = roc_auc_score(true_label, pred_score)
                batch_bar.set_postfix(loss=loss.item(),
                                      acc=acc,
                                      f1=f1,
                                      auc=auc)

        # validation phase use all information
        dps.ngh_finder = full_ngh_finder
        val_acc, val_ap, val_f1, val_auc = eval_one_epoch(
            'val for old nodes', dps, val_src_l, val_dst_l, val_ts_l,
            val_label_l)
        epoch_bar.update()
        epoch_bar.set_postfix(acc=val_acc, f1=val_f1, auc=val_auc)

        if early_stopper.early_stop_check(val_auc):
            break
        else:
            torch.save(dps.state_dict(), get_checkpoint_path(epoch))

    logger.info('No improvment over {} epochs, stop training'.format(
        early_stopper.max_round))
    logger.info(f'Loading the best model at epoch {early_stopper.best_epoch}')
    best_model_path = get_checkpoint_path(early_stopper.best_epoch)
    dps.load_state_dict(torch.load(best_model_path))
    logger.info(
        f'Loaded the best model at epoch {early_stopper.best_epoch} for inference'
    )

# If only use attention model, we sample the top-k neighbors w.r.t to
# attention scores.
if HARD == "atte":
    logger.info('Saving DPS model')
    torch.save(dps.state_dict(), MODEL_SAVE_PATH)
    logger.info('DPS models saved')
    exit(0)

# Finetune the GumbelGAN model.
if True:
    logger.info("Finetune the GumbelGAN model.")
    TEMP = args.temp
    TEMP_MIN = 0.5
    ANNEAL_RATE = args.anneal
    optimizer = torch.optim.Adam(dps.parameters(),
                                 lr=LEARNING_RATE / BATCH_SIZE)
    early_stopper = EarlyStopMonitor()
    # epoch_bar = trange(NUM_EPOCH)
    epoch_bar = trange(1)
    for epoch in epoch_bar:
        # Training
        # training use only training graph
        dps.ngh_finder = train_ngh_finder
        np.random.shuffle(idx_list)
        batch_bar = trange(num_batch)
        temp = TEMP
        for k in batch_bar:
            s_idx = k * BATCH_SIZE
            e_idx = min(num_instance - 1, s_idx + BATCH_SIZE)
            src_l_cut = train_src_l[s_idx:e_idx]
            dst_l_cut = train_dst_l[s_idx:e_idx]
            ts_l_cut = train_ts_l[s_idx:e_idx]
            size = len(src_l_cut)
            src_l_fake, dst_l_fake = train_rand_sampler.sample(size)

            with torch.no_grad():
                pos_label = torch.ones((1, ), dtype=torch.float, device=device)
                neg_label = torch.zeros((1, ),
                                        dtype=torch.float,
                                        device=device)

            optimizer.zero_grad()
            dps = dps.train()
            src_l_cut = src_l_cut[:, np.newaxis]
            dst_l_cut = dst_l_cut[:, np.newaxis]
            dst_l_fake = dst_l_fake[:, np.newaxis]
            ts_l_cut = ts_l_cut[:, np.newaxis]
            pos_prob = []
            neg_prob = []
            for src, dst, fake, ts in zip(src_l_cut, dst_l_cut, dst_l_fake,
                                          ts_l_cut):
                pos, neg = dps.contrast(src,
                                        dst,
                                        fake,
                                        ts,
                                        NUM_NEIGHBORS,
                                        gumbel=True)

                loss = criterion(pos, pos_label)
                loss += criterion(neg, neg_label)

                pos_prob.append(pos.cpu().detach().item())
                neg_prob.append(neg.cpu().detach().item())
                loss.backward()
            optimizer.step()
            if k % 10 == 1:
                temp = np.maximum(temp * np.exp(-ANNEAL_RATE * k), TEMP_MIN)
            dps.anneal_temp(temp)
            # get training results
            with torch.no_grad():
                dps = dps.eval()
                pred_score = np.concatenate([pos_prob, neg_prob])
                pred_label = pred_score > 0.5
                true_label = np.concatenate([np.ones(size), np.zeros(size)])
                acc = accuracy_score(true_label, pred_label)
                ap = average_precision_score(true_label, pred_label)
                f1 = f1_score(true_label, pred_label)
                auc = roc_auc_score(true_label, pred_score)
                batch_bar.set_postfix(loss=loss.item(),
                                      acc=acc,
                                      f1=f1,
                                      auc=auc)

        # validation phase use all information
        dps.ngh_finder = full_ngh_finder
        val_acc, val_ap, val_f1, val_auc = eval_one_epoch(
            'val for old nodes', dps, val_src_l, val_dst_l, val_ts_l,
            val_label_l)
        epoch_bar.update()
        epoch_bar.set_postfix(acc=val_acc, f1=val_f1, auc=val_auc)

        if early_stopper.early_stop_check(val_auc):
            break
        else:
            torch.save(dps.state_dict(), get_checkpoint_path(epoch))

    logger.info('No improvment over {} epochs, stop training'.format(
        early_stopper.max_round))
    logger.info(f'Loading the best model at epoch {early_stopper.best_epoch}')
    best_model_path = get_checkpoint_path(early_stopper.best_epoch)
    dps.load_state_dict(torch.load(best_model_path))
    logger.info(
        f'Loaded the best model at epoch {early_stopper.best_epoch} for inference'
    )

logger.info('Saving DPS model')
torch.save(dps.state_dict(), MODEL_SAVE_PATH)
logger.info('DPS models saved')