accuracy.py

import torch
import torch.nn as nn
import numpy as np
from torch.nn import Parameter
import torch.optim as optim
from sklearn import metrics
import random
import os
import sys
import utils
from tqdm import tqdm
from collections import defaultdict
import time
import argparse
import dgl
from scipy.sparse import csr_matrix
from scipy.sparse import vstack as s_vstack
from sklearn.preprocessing import StandardScaler
from gensim.models import Word2Vec
import multiprocessing
from concurrent.futures import as_completed
from concurrent.futures import ProcessPoolExecutor
from scipy.sparse import csr_matrix, lil_matrix, csc_matrix

from preprocess.data_load import gen_DGLGraph, gen_weighted_DGLGraph
import preprocess.data_load as dl
from preprocess.batch import DataLoader
from initialize.initial_embedder import MultipleEmbedding
from initialize.random_walk_hyper import random_walk_hyper

from model.Whatsnet import Whatsnet, WhatsnetLayer
from model.layer import FC, Wrap_Embedding

def run_epoch(args, data, dataloader, initembedder, customerembedder, colorembedder, sizeembedder, groupembedder, embedder, scorer, opt="train"):
    total_pred = []
    total_label = []
    num_data = 0
    num_recon_data = 0
    total_loss = 0
    total_ce_loss = 0
    total_recon_loss = 0
    
    # Batch ==============================================================
    ts = time.time()
    batchcount = 0
    pbar = tqdm(enumerate(dataloader), total=len(dataloader))
    for step, (input_nodes, output_nodes, blocks) in pbar:      
        # Wrap up loader
        blocks = [b.to(device) for b in blocks]
        srcs, dsts = blocks[-1].edges(etype='in')
        nodeindices = srcs.to(device)
        hedgeindices = dsts.to(device)
        nodelabels = blocks[-1].edges[('node','in','edge')].data['label'].long().to(device)

        batchcount += 1
        # Get Embedding
        if args.embedder == "whatsnetLSPE":
            v_feat, recon_loss = initembedder(input_nodes['node'].to(device))
            e_feat = data.e_feat[input_nodes['edge']].to(device)
            v_pos = data.v_pos[input_nodes['node']].to(device)
            e_pos = data.e_pos[input_nodes['edge']].to(device)
            v, e = embedder(blocks, v_feat, e_feat, v_pos, e_pos)
        elif args.embedder == "whatsnet":
            if args.att_type_v in ["ITRE", "ShawRE"]:
                vindex = torch.arange(len(input_nodes['node'])).unsqueeze(1).to(device)
                v_feat, recon_loss = initembedder(input_nodes['node'].to(device))
                e_feat = data.e_feat[input_nodes['edge']].to(device)
                v, e = embedder(blocks, v_feat, e_feat, vindex)
            else:
                v_feat, recon_loss = initembedder(input_nodes['node'].to(device))
                e_feat = data.e_feat[input_nodes['edge']].to(device)
                # concat feature embedding
                color_feat, _ = colorembedder(data.node2color[input_nodes['node']].long().to(device))
                size_feat, _ = sizeembedder(data.node2size[input_nodes['node']].long().to(device))
                group_feat, _ = groupembedder(data.node2group[input_nodes['node']].long().to(device))
                v_feat = torch.concat((v_feat, color_feat, size_feat, group_feat), dim=1)
                if not args.custom:
                    v, e = embedder(blocks, v_feat, e_feat, None, None)
                else:
                    v, e = embedder(blocks, v_feat, e_feat, data, customerembedder)
        else:
                v_feat, recon_loss = initembedder(input_nodes['node'].to(device))
                e_feat = data.e_feat[input_nodes['edge']].to(device)
                v, e = embedder(blocks, v_feat, e_feat)
                
        # Predict Class
        if args.scorer == "sm":
          if args.dataset_name == 'task1':
            unique_hedgeindices = torch.unique(hedgeindices)
            label_indices = []
            for i in range(len(hedgeindices)):
                if i == 0:
                  label_indices.append(i)
                  prev_value = hedgeindices[i]
                else:
                  if hedgeindices[i] != prev_value:
                      label_indices.append(i)
                      prev_value = hedgeindices[i]

            label_indices = torch.Tensor(label_indices).long()
            nodelabels = nodelabels[label_indices]   
            hembedding = e[unique_hedgeindices]
            input_embeddings = hembedding
            predictions = scorer(input_embeddings)

          elif args.dataset_name == 'task2':
            if opt == 'train':
              trainable_indices = torch.nonzero((nodelabels == 0) | (nodelabels == 1)).squeeze()
              hedgeindices = hedgeindices[trainable_indices]
              nodeindices = nodeindices[trainable_indices]
              nodelabels = nodelabels[trainable_indices]

              hembedding = e[hedgeindices]
              vembedding = v[nodeindices]
              input_embeddings = torch.cat([hembedding,vembedding], dim=1)
              predictions = scorer(input_embeddings)
            elif opt == 'valid':
              valid_indices = torch.nonzero((nodelabels == 2) | (nodelabels == 3)).squeeze()
              hedgeindices = hedgeindices[valid_indices]
              nodeindices = nodeindices[valid_indices]
              nodelabels = nodelabels[valid_indices]
              nodelabels = nodelabels - 2


              hembedding = e[hedgeindices]
              vembedding = v[nodeindices]
              input_embeddings = torch.cat([hembedding,vembedding], dim=1)
              predictions = scorer(input_embeddings)
              
          else:
            raise ValueError("Not supported data type")
              
        elif args.scorer == "im":
            predictions, nodelabels = scorer(blocks[-1], v, e)
        total_pred.append(predictions.detach())
        total_label.append(nodelabels.detach())

        description = f'Step: {step+1}/{len(dataloader)}'
        pbar.set_description(description)
    
    print("Time : ", time.time() - ts)
    print(num_data)
    
    return total_pred, total_label

def run_test_epoch(args, data, testdataloader, initembedder, customerembedder, colorembedder, sizeembedder, groupembedder, embedder, scorer, loss_fn):
    total_pred = []
    total_label = []
    num_data = 0
    num_recon_data = 0
    total_loss = 0
    total_ce_loss = 0
    total_recon_loss = 0
    
    # Batch ==============================================================
    ts = time.time()
    batchcount = 0
    for input_nodes, output_nodes, blocks in testdataloader:      
        # Wrap up loader
        blocks = [b.to(device) for b in blocks]
        srcs, dsts = blocks[-1].edges(etype='in')
        nodeindices_in_batch = srcs.to(device)
        nodeindices = blocks[-1].srcdata[dgl.NID]['node'][srcs]
        hedgeindices_in_batch = dsts.to(device)
        hedgeindices = blocks[-1].srcdata[dgl.NID]['edge'][dsts]
        hedgeindices = dsts.to(device)
        nodelabels = blocks[-1].edges[('node','in','edge')].data['label'].long().to(device)
        
        batchcount += 1
        # Get Embedding
        if args.embedder == "whatsnetLSPE":
            v_feat, recon_loss = initembedder(input_nodes['node'].to(device))
            e_feat = data.e_feat[input_nodes['edge']].to(device)
            v_pos = data.v_pos[input_nodes['node']].to(device)
            e_pos = data.e_pos[input_nodes['edge']].to(device)
            v, e = embedder(blocks, v_feat, e_feat, v_pos, e_pos)
        elif args.embedder == "whatsnet":
            if args.att_type_v in ["ITRE", "ShawRE", "RafRE"]:
                vindex = torch.arange(len(input_nodes['node'])).unsqueeze(1).to(device)
                v_feat, recon_loss = initembedder(input_nodes['node'].to(device))
                e_feat = data.e_feat[input_nodes['edge']].to(device)
                v, e = embedder(blocks, v_feat, e_feat, vindex)
            else:
                v_feat, recon_loss = initembedder(input_nodes['node'].to(device))
                e_feat = data.e_feat[input_nodes['edge']].to(device)
                v, e = embedder(blocks, v_feat, e_feat)
        else:
                v_feat, recon_loss = initembedder(input_nodes['node'].to(device))
                e_feat = data.e_feat[input_nodes['edge']].to(device)
                v, e = embedder(blocks, v_feat, e_feat)
                
        # Predict Class
        if args.scorer == "sm":
          if args.dataset_name == "task1":
            hembedding = e[hedgeindices_in_batch]
            input_embeddings = hembedding
            predictions = scorer(input_embeddings)
          elif args.dataset_name == "task2":
            hembedding = e[hedgeindices_in_batch]
            vembedding = v[nodeindices_in_batch]
            input_embeddings = torch.cat([hembedding,vembedding], dim=1)
            predictions = scorer(input_embeddings)
          else:
            raise ValueError("Not supported dataset type")
          
        elif args.scorer == "im":
            predictions, nodelabels = scorer(blocks[-1], v, e)
        total_pred.append(predictions.detach())
        pred_cls = torch.argmax(predictions, dim=1)
        total_label.append(nodelabels.detach())
        
        num_data += predictions.shape[0]
        num_recon_data += input_nodes['node'].shape[0]
        ce_loss = loss_fn(predictions, nodelabels)
        loss = ce_loss + args.rw * recon_loss

        total_loss += (loss.item() * predictions.shape[0])
        total_ce_loss += (ce_loss.item() * predictions.shape[0])
        total_recon_loss += (recon_loss.item() * input_nodes['node'].shape[0]) # This is fixed as zero
        
    return total_pred, total_label, total_loss / num_data, total_ce_loss / num_data, total_recon_loss / num_recon_data, initembedder, customerembedder, colorembedder, sizeembedder, groupembedder, embedder, scorer

# Make Output Directory --------------------------------------------------------------------------------------------------------------
initialization = "rw"
args = utils.parse_args()
if args.evaltype == "test":
    assert args.fix_seed
    outputdir = "results_test/" + args.dataset_name + "_" + str(args.k) + "/" + initialization + "/"
    outputParamResFname = outputdir + args.model_name + "/param_result.txt"
    if args.custom:
        outputdir += args.model_name + "/" + args.param_name +"/" + str(args.seed) + "/" + "custom_embedding/"
    else:
        outputdir += args.model_name + "/" + args.param_name +"/" + str(args.seed) + "/"
    if args.recalculate is False and os.path.isfile(outputdir + "log_test_confusion.txt"):
        sys.exit("Already Run")
else:
    outputdir = "results/" + args.dataset_name + "_" + str(args.k) + "/" + initialization + "/"
    outputdir += args.model_name + "/" + args.param_name +"/"
    if args.recalculate is False and os.path.isfile(outputdir + "log_test_confusion.txt"):
        sys.exit("Already Run")
if os.path.isdir(outputdir) is False:
    os.makedirs(outputdir)
print("OutputDir = " + outputdir)

# Initialization --------------------------------------------------------------------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
dataset_name = args.dataset_name #'citeseer' 'cora'

if args.fix_seed:
    random.seed(args.seed)
    np.random.seed(args.seed)
    dgl.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
    dgl.seed(args.seed)

exp_num = args.exp_num
test_epoch = args.test_epoch
plot_epoch = args.epochs

# Data -----------------------------------------------------------------------------
data = dl.Hypergraph(args, dataset_name)
if args.dataset_name == 'task1':
  train_data = data.get_data(0)
  valid_data = data.get_data(1)
elif args.dataset_name == 'task2':
  train_data = data.get_data(0, task2=True)
  valid_data = data.get_data(1, task2=True)
if args.evaltype == "test":
    test_data = data.get_data(2)

ls = [{('node', 'in', 'edge'): -1, ('edge', 'con', 'node'): args.sampling}] * (args.num_layers * 2 + 1)
full_ls = [{('node', 'in', 'edge'): -1, ('edge', 'con', 'node'): -1}] * (args.num_layers * 2 + 1)
if data.weight_flag:
    g = gen_weighted_DGLGraph(args, data.hedge2node, data.hedge2nodePE, data.hedge2nodepos, data.node2hedge, data.node2hedgePE, device)
else:
    g = gen_DGLGraph(args, data.hedge2node, data.hedge2nodepos, data.node2hedge, device)
try:
    sampler = dgl.dataloading.NeighborSampler(ls)
    fullsampler = dgl.dataloading.NeighborSampler(full_ls)
except:
    sampler = dgl.dataloading.MultiLayerNeighborSampler(ls, False)
    fullsampler = dgl.dataloading.MultiLayerNeighborSampler(full_ls)
if args.use_gpu:
    g = g.to(device)
    train_data = train_data.to(device)
    valid_data = valid_data.to(device)
    if args.evaltype == "test":
        test_data = test_data.to(device)
    data.e_feat = data.e_feat.to(device)
dataloader = dgl.dataloading.DataLoader( g, {"edge": train_data}, sampler, batch_size=args.bs, shuffle=True, drop_last=False) # , num_workers=4
validdataloader = dgl.dataloading.DataLoader( g, {"edge": valid_data}, sampler, batch_size=args.bs, shuffle=True, drop_last=False)
if args.evaltype == "test":
    testdataloader = dgl.dataloading.DataLoader(g, {"edge": test_data}, fullsampler, batch_size=args.bs, shuffle=False, drop_last=False)

args.input_vdim = data.v_feat.size(1)
args.input_edim = data.e_feat.size(1)
args.order_dim = data.order_dim

# init embedder
args.input_vdim = 48
if args.orderflag:
    args.input_vdim = 44
savefname = "../%s_%d_wv_%d_%s.npy" % (args.dataset_name, args.k, args.input_vdim, args.walk)
node_list = np.arange(data.numnodes).astype('int')

# A: (data.numnodes, vector_size)의 embedding 생성, 이는 초기 node embedding으로 사용
if os.path.isfile(savefname) is False:
    walk_path = random_walk_hyper(args, node_list, data.hedge2node)
    walks = np.loadtxt(walk_path, delimiter=" ").astype('int')
    print("Start turning path to strs")
    split_num = 20
    pool = ProcessPoolExecutor(max_workers=split_num)
    process_list = []
    walks = np.array_split(walks, split_num)
    result = []
    for walk in walks:
        process_list.append(pool.submit(utils.walkpath2str, walk))
    for p in as_completed(process_list):
        result += p.result()
    pool.shutdown(wait=True)
    walks = result
    # print(walks)
    print("Start Word2vec")
    print("num cpu cores", multiprocessing.cpu_count())
    w2v = Word2Vec(walks, vector_size=args.input_vdim, window=10, min_count=0, sg=1, epochs=1, workers=multiprocessing.cpu_count())
    print(w2v.wv['0'])
    wv = w2v.wv
    A = [wv[str(i)] for i in range(data.numnodes)]
    np.save(savefname, A)
else:
    print("load exist init walks")
    A = np.load(savefname)
A = StandardScaler().fit_transform(A)
A = A.astype('float32')
A = torch.tensor(A).to(device)
# Make embedding vector with input 
print(data.numnodes, data.numcustomers, data.numcolors, data.numsizes, data.numgroups)
initembedder = Wrap_Embedding(data.numnodes, args.input_vdim, scale_grad_by_freq=False, sparse=False).to(device)
initembedder.weight = nn.Parameter(A)

# Make embedding with feature
# 0-342037
customerembedder = Wrap_Embedding(data.numcustomers, 48, scale_grad_by_freq=False, sparse=False).to(device)
# 0-641
colorembedder = Wrap_Embedding(data.numcolors, 8, scale_grad_by_freq=False, sparse=False).to(device)
# 0-28
sizeembedder = Wrap_Embedding(data.numsizes, 8, scale_grad_by_freq=False, sparse=False).to(device)
# 0-31
groupembedder = Wrap_Embedding(data.numgroups, 8, scale_grad_by_freq=False, sparse=False).to(device)
# Add dimension for embeddings
args.input_vdim = args.input_vdim + 24
# args.input_edim = args.input_edim + 24

print("Model:", args.embedder)
# model init
if args.embedder == "whatsnet":  
    input_vdim = args.input_vdim
    pe_ablation_flag = args.pe_ablation
    embedder = Whatsnet(WhatsnetLayer, input_vdim, args.input_edim, args.dim_hidden, args.dim_vertex, args.dim_edge, 
        weight_dim=args.order_dim, num_heads=args.num_heads, num_layers=args.num_layers, num_inds=args.num_inds,
        att_type_v=args.att_type_v, agg_type_v=args.agg_type_v, att_type_e=args.att_type_e, agg_type_e=args.agg_type_e,
        num_att_layer=args.num_att_layer, dropout=args.dropout, weight_flag=data.weight_flag, pe_ablation_flag=pe_ablation_flag, vis_flag=args.analyze_att).to(device)
else:
    print("Only whatsnet embedder supported currently")

    
print("Embedder to Device")
print("Scorer = ", args.scorer)
# pick scorer
if args.scorer == "sm":
    if args.dataset_name == 'task1':
      scorer = FC(args.dim_edge, args.dim_edge, args.output_dim, args.scorer_num_layers, args.dropout).to(device)
    elif args.dataset_name == 'task2':
      scorer = FC(args.dim_vertex + args.dim_edge, args.dim_edge, args.output_dim, args.scorer_num_layers, args.dropout).to(device)
    else:
      raise ValueError("Not supported dataset type")

# Train =================================================================================================================================================================================

initembedder.load_state_dict(torch.load(outputdir + 'initembedder.pt'))
customerembedder.load_state_dict(torch.load(outputdir + 'customerembedder.pt'))
colorembedder.load_state_dict(torch.load(outputdir + 'colorembedder.pt'))
sizeembedder.load_state_dict(torch.load(outputdir + 'sizeembedder.pt'))
groupembedder.load_state_dict(torch.load(outputdir + 'groupembedder.pt'))
embedder.load_state_dict(torch.load(outputdir + 'embedder.pt'))
scorer.load_state_dict(torch.load(outputdir + 'scorer.pt'))

# Test ===========================================================================================================================================================================

initembedder.eval()
customerembedder.eval()
colorembedder.eval()
sizeembedder.eval()
groupembedder.eval()
embedder.eval()
scorer.eval()

with torch.no_grad():
    total_pred, total_label = run_epoch(args, data, validdataloader, initembedder, customerembedder, colorembedder, sizeembedder, groupembedder, embedder, scorer, opt="valid")
# Calculate Accuracy & Epoch Loss
total_label = torch.cat(total_label, dim=0)
total_pred = torch.cat(total_pred)
pred_cls = torch.argmax(total_pred, dim=1)
eval_acc = torch.eq(pred_cls, total_label).sum().item() / len(total_label)
y_test = total_label.cpu().numpy()
pred = pred_cls.cpu().numpy()
confusion, accuracy, precision, recall, f1 = utils.get_clf_eval(y_test, pred, avg='micro', outputdim=args.output_dim)
confusion, accuracy, precision, recall, f1 = utils.get_clf_eval(y_test, pred, avg='macro', outputdim=args.output_dim)
print(accuracy)