run_gnn.py

import time
import random
import argparse
import scipy
import numpy as np
import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import f1_score

from model import GCN, GAT

from utils.pytorchtools import EarlyStopping
from utils.data import load_data
from utils.tools import evaluate_results_nc


def score(logits, labels):
    _, indices = torch.max(logits, dim=1)
    prediction = indices.long().cpu().numpy()
    labels = labels.cpu().numpy()
    acc = (prediction == labels).sum() / len(prediction)
    micro_f1 = f1_score(labels, prediction, average='micro')
    macro_f1 = f1_score(labels, prediction, average='macro')
    return acc, micro_f1, macro_f1


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        print('Using CUDA')
        torch.cuda.manual_seed(seed)

set_seed(123)

def run(args):
    adj_lists, features_list, _, _, _, _, _, labels, train_val_test_idx = load_data(args.dataset)
    device = torch.device(f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu')
    features = torch.FloatTensor(features_list[0]).to(device)

    labels = torch.LongTensor(labels).to(device)
    train_idx = train_val_test_idx['train_idx']
    train_idx = np.sort(train_idx)
    val_idx = train_val_test_idx['val_idx']
    val_idx = np.sort(val_idx)
    test_idx = train_val_test_idx['test_idx']
    test_idx = np.sort(test_idx)
    in_dim = features.shape[1]

    print(labels.size())
    print(features.size())

    for i, adj in enumerate(adj_lists):

        g = dgl.DGLGraph(adj)
        g = dgl.remove_self_loop(g)
        g = dgl.add_self_loop(g)
        g = g.to(device)
        print(type(g))
        N = g.number_of_nodes()

        loss_fn = torch.nn.CrossEntropyLoss()

        svm_macro_f1_lists = []
        svm_micro_f1_lists = []
        nmi_mean_list = []
        nmi_std_list = []
        ari_mean_list = []
        ari_std_list = []
        time_used = []
        test_micro_f1s, test_macro_f1s = [], []
        for run in range(args.repeat):
            num_classes = labels.max().item()+1
            if args.model == 'gat':
                net = GAT(g, in_dim, args.hidden_dim, args.num_heads, num_classes, args.num_layers, F.elu, args.dropout)
            elif args.model == 'gcn':
                net = GCN(g, in_dim, args.hidden_dim, num_classes, args.num_layers, F.relu, args.dropout)
            sum_p = sum(p.numel() for p in net.parameters())
            print(sum_p)
            # assert False
            net.to(device)
            optimizer = torch.optim.Adam(net.parameters(), lr=args.lr, weight_decay=args.weight_decay)

            # training loop
            net.train()
            early_stopping = EarlyStopping(patience=args.patience, verbose=False,
                                        save_path='checkpoint/checkpoint_{}.pt'.format(args.save_postfix))
            st = time.perf_counter()
            times_per_epoch = []
            for epoch in range(args.epochs):
                t_start = time.time()
                net.train()
                logits, _ = net(features)
                train_loss = loss_fn(logits[train_idx], labels[train_idx])
                optimizer.zero_grad()
                train_loss.backward()
                optimizer.step()

                net.eval()
                with torch.no_grad():
                    logits, _ = net(features)
                    val_loss = loss_fn(logits[val_idx], labels[val_idx])
                    val_acc, val_mif1, val_maf1 = score(logits[val_idx], labels[val_idx])
                t_end = time.time()
                t_used = t_end - t_start
                times_per_epoch.append(t_used)
                # print validation info
                if args.repeat == 1:
                    print('Epoch {:05d} | Train_Loss {:.4f} | Val_Loss {:.4f}, Val_mif1 {:.4f}, Val_maf1 {:.4f} | Time(s) {:.4f}'.format(
                        epoch, train_loss.item(), val_loss.item(), val_mif1, val_maf1, t_used))
                # early stopping
                early_stopping(val_loss, net)
                if early_stopping.early_stop:
                    # print('Early stopping!')
                    break
            time_used.append(time.perf_counter()-st)
            # testing with evaluate_results_nc
            net.load_state_dict(torch.load('checkpoint/checkpoint_{}.pt'.format(args.save_postfix)))
            net.eval()
            #test_embeddings = []
            with torch.no_grad():
                logits, embeddings = net(features)
                test_embeddings = embeddings[test_idx]
                print('-----------')
                print(f'Run: {run}')
                times_per_epoch = torch.tensor(times_per_epoch)
                print(f'Times per epoch: {times_per_epoch.mean()}s')
                test_acc, test_micro_f1, test_macro_f1 = score(logits[test_idx],labels[test_idx])
                if args.save_emb:
                    np.save(f"plots/plots_data/{args.dataset}_{args.model}_emb.npy", test_embeddings.cpu().numpy())
                    np.save(f"plots/plots_data/{args.dataset}_{args.model}_label.npy", labels[test_idx].cpu().numpy())
                print(test_acc, test_micro_f1, test_macro_f1)
                # svm_macro_f1_list, svm_micro_f1_list, nmi_mean, nmi_std, ari_mean, ari_std = evaluate_results_nc(
                nmi_mean, nmi_std, ari_mean, ari_std = evaluate_results_nc(
                    test_embeddings.cpu().numpy(), labels[test_idx].cpu().numpy(), num_classes=num_classes)
                print('-----------')
            test_micro_f1s.append(test_micro_f1)
            test_macro_f1s.append(test_macro_f1)
            # svm_macro_f1_lists.append(svm_macro_f1_list)
            # svm_micro_f1_lists.append(svm_micro_f1_list)
            nmi_mean_list.append(nmi_mean)
            nmi_std_list.append(nmi_std)
            ari_mean_list.append(ari_mean)
            ari_std_list.append(ari_std)

        # print out a summary of the evaluations
        # svm_macro_f1_lists = np.transpose(np.array(svm_macro_f1_lists), (1, 0, 2))
        # svm_micro_f1_lists = np.transpose(np.array(svm_micro_f1_lists), (1, 0, 2))
        nmi_mean_list = np.array(nmi_mean_list)
        nmi_std_list = np.array(nmi_std_list)
        ari_mean_list = np.array(ari_mean_list)
        ari_std_list = np.array(ari_std_list)
        print('----------------------------------------------------------------')
        test_macro_f1s = np.array(test_macro_f1s)
        test_micro_f1s = np.array(test_micro_f1s)
        print(f"Meta-Path Graph: {i+1}, Test Results")
        print(f"Average Test macro f1: {test_macro_f1s.mean()*100:.2f} ± {test_macro_f1s.std()*100:.2f}")
        print(f"Average Test micro f1: {test_micro_f1s.mean()*100:.2f} ± {test_micro_f1s.std()*100:.2f}")
        print('----------------------------------------------------------------')
        # print('SVM tests summary')
        # print('Macro-F1: ' + ', '.join(['{:.4f}~{:.4f} ({:.2f})'.format(
        #     macro_f1[:, 0].mean(), macro_f1[:, 1].mean(), train_size) for macro_f1, train_size in
        #     zip(svm_macro_f1_lists, [0.8, 0.6, 0.4, 0.2, 0.1, 0.05])]))
        # print('Micro-F1: ' + ', '.join(['{:.4f}~{:.4f} ({:.2f})'.format(
        #     micro_f1[:, 0].mean(), micro_f1[:, 1].mean(), train_size) for micro_f1, train_size in
        #     zip(svm_micro_f1_lists, [0.8, 0.6, 0.4, 0.2, 0.1, 0.05])]))
        print('K-means tests summary')
        print('NMI: {:.2f} ± {:.2f}'.format(nmi_mean_list.mean() * 100, nmi_mean_list.std() * 100))
        print('ARI: {:.2f} ± {:.2f}'.format(ari_mean_list.mean() * 100, ari_mean_list.std() * 100))
        time_used = torch.tensor(time_used)
        print("Used Time:", time_used.mean(), time_used.std())

        filename = f'results/{args.model}-m-' +  f'{args.dataset}.csv'
        print(f"Saving results to {filename}")
        with open(f"{filename}", 'a+') as write_obj:
            write_obj.write(f"{args.model}-m," + 
                            f"meta_path:{i+1}," +
                            f"n_layer:{args.num_layers}," +
                            f"he:{args.num_heads}," +
                            f"hd:{args.hidden_dim}," +
                            f"lr:{args.lr}," + 
                            f"wd:{args.weight_decay}," + 
                            f"dp:{args.dropout}," +
                            f"{test_macro_f1s.mean()*100:.2f} ± {test_macro_f1s.std()*100:.2f}," +
                            f"{test_micro_f1s.mean()*100:.2f} ± {test_micro_f1s.std()*100:.2f}," +
                            f"{nmi_mean_list.mean()*100:.2f} ± {nmi_mean_list.std()*100:.2f}," +
                            f"{ari_mean_list.mean()*100:.2f} ± {ari_mean_list.std()*100:.2f}\n")


if __name__ == '__main__':
    ap = argparse.ArgumentParser(description='MRGNN testing for the ACM dataset')
    ap.add_argument('--dataset', default='ACM', help='ACM, DBLP, IMDB')
    ap.add_argument('--model', default='gcn', help='gcn, gat')
    ap.add_argument('--feats_type', type=int, default=0,
                    help='Type of the node features used. ' +
                         '0 - loaded features; ' +
                         '1 - only target node features (zero vec for others); ' +
                         '2 - only target node features (id vec for others); ' +
                         '3 - all id vec.')
    ap.add_argument('--hidden_dim', type=int, default=64, help='Dimension of the node hidden state. Default is 64.')
    ap.add_argument('--num_heads', type=int, default=8, help='Number of the attention heads. Default is 8.')
    ap.add_argument('--num_layers', type=int, default=4)
    ap.add_argument('--epochs', type=int, default=200, help='Number of epochs. Default is 100.')
    ap.add_argument('--patience', type=int, default=50, help='Patience. Default is 5.')
    ap.add_argument('--repeat', type=int, default=10, help='Repeat the training and testing for N times. Default is 1.')
    ap.add_argument('--save_postfix', default='DBLP', help='Postfix for the saved model and result. Default is DBLP.')
    ap.add_argument('--device', type=int, default=5)
    ap.add_argument('--dropout', type=float, default=0.6)
    ap.add_argument('--lr', type=float, default=0.001)
    ap.add_argument('--weight_decay', type=float, default=0.001)
    ap.add_argument('--R', type=float, default=100)
    ap.add_argument('--no_re', action='store_true')
    ap.add_argument('--save_emb', action='store_true')


    args = ap.parse_args()

    if args.no_re:
        args.R = 1e-10

    print(args)

    run(args)