<filename>example_reddit/multi_reddit_all.py
from functools import reduce
from GPT_GNN.data import *
from GPT_GNN.model_meta import *
import torch.nn as nn
from warnings import filterwarnings
import numpy as np
from sklearn.metrics import f1_score
import pdb

filterwarnings("ignore")

import argparse

parser = argparse.ArgumentParser(
    description='Multi-task method, the primary task is OAG Paper-Field (L2) classification task')

'''
   Multi-task arguments 
'''
parser.add_argument('--link_ratio', type=float, default=1.0,
                    help='Ratio of loss-loss against primary-loss, range: [0-1]')
parser.add_argument('--attr_ratio', type=float, default=1.0,
                    help='Ratio of attr-loss against primary-loss, range: [0-1]')
parser.add_argument('--attr_type', type=str, default='vec',
                    choices=['text', 'vec'],
                    help='The type of attribute decoder')
parser.add_argument('--neg_samp_num', type=int, default=255,
                    help='Maximum number of negative sample for each target node.')
parser.add_argument('--queue_size', type=int, default=256,
                    help='Max size of adaptive embedding queue.')
parser.add_argument('--w2v_dir', type=str, default='/data/data1/v-bangan/GPT/w2v_all',
                    help='The address of preprocessed graph.')
parser.add_argument('--use_pretrain', help='Whether to use pre-trained model', action='store_true')
parser.add_argument('--pretrain_model_dir', type=str, default='/datadrive/models/gpt_all_cs',
                    help='The address for pretrained model.')

'''
    Dataset arguments
'''
parser.add_argument('--data_dir', type=str, default='/data/data1/v-zhehuang/HGT',
                    help='The address of preprocessed graph.')
parser.add_argument('--model_dir', type=str, default='/data/data1/v-zhehuang/GPT/gpt_all_reddit',
                    help='The address for storing the models and optimization results.')
parser.add_argument('--model_add_name', type=str, default='tmp',
                    help='Additional name.')
parser.add_argument('--task_name', type=str, default='reddit',
                    help='The name of the stored models and optimization results.')
parser.add_argument('--cuda', type=str, default=2,
                    help='Avaiable GPU ID')
parser.add_argument('--sample_depth', type=int, default=6,
                    help='How many layers within a mini-batch subgraph')
parser.add_argument('--sample_width', type=int, default=128,
                    help='How many nodes to be sampled per layer per type')

'''
   Model arguments 
'''
parser.add_argument('--conv_name', type=str, default='hgt',
                    choices=['hgt', 'gcn', 'gat', 'rgcn', 'han', 'hetgnn'],
                    help='The name of GNN filter. By default is Heterogeneous Graph Transformer (hgt)')
parser.add_argument('--n_hid', type=int, default=400,
                    help='Number of hidden dimension')
parser.add_argument('--n_heads', type=int, default=8,
                    help='Number of attention head')
parser.add_argument('--n_layers', type=int, default=3,
                    help='Number of GNN layers')
parser.add_argument('--prev_norm', help='Whether to add layer-norm on the previous layers', action='store_true')
parser.add_argument('--last_norm', help='Whether to add layer-norm on the last layers', action='store_true')
parser.add_argument('--dropout', type=int, default=0.2,
                    help='Dropout ratio')
parser.add_argument('--data_percentage', type=float, default=1.0,
                    help='Percentage of training and validation data to use')
parser.add_argument('--temperature', type=int, default=2, help='weight adaptation')

'''
    Optimization arguments
'''
parser.add_argument('--optimizer', type=str, default='adamw',
                    choices=['adamw', 'adam', 'sgd', 'adagrad'],
                    help='optimizer to use.')
parser.add_argument('--max_lr', type=float, default=1e-3,
                    help='Maximum learning rate.')
parser.add_argument('--scheduler', type=str, default='cosine',
                    help='Name of learning rate scheduler.', choices=['cycle', 'cosine'])
parser.add_argument('--n_epoch', type=int, default=200,
                    help='Number of epoch to run')
parser.add_argument('--n_pool', type=int, default=8,
                    help='Number of process to sample subgraph')
parser.add_argument('--n_batch', type=int, default=16,
                    help='Number of batch (sampled graphs) for each epoch')
parser.add_argument('--batch_size', type=int, default=256,
                    help='Number of output nodes for training')
parser.add_argument('--valid_batch', type=int, default=10,
                    help='Number of valid batch (sampled graphs) for each epoch')
parser.add_argument('--clip', type=float, default=0.5,
                    help='Gradient Norm Clipping')
'''
meta-optimizing arguments
'''
parser.add_argument('--n_fold', type=int, default=1)
parser.add_argument('--weight_emb_dim', type=int, default=1000)
parser.add_argument('--wlr', type=float, default=0.005)
parser.add_argument('--act_type', type=str, default='sigmoid')
parser.add_argument('--seed', type=int, default=43,
                    help='Number of output nodes for training')  


args = parser.parse_args()
args_print(args)
CUDA_STR = "cuda:" + str(args.cuda)
np.random.seed(args.seed)

if args.cuda != -1:
    device = torch.device("cuda:" + str(args.cuda))
else:
    device = torch.device("cpu")
print('Start Loading Graph Data...')
graph = dill.load(open(os.path.join(args.data_dir, 'graph_reddit.pk'), 'rb'))
print('Finish Loading Graph Data!')
if not os.path.isdir(args.model_dir):
    os.mkdir(args.model_dir)

# primary task pre-define
target_type = 'def'
rel_stop_list = ['self']
pre_target_nodes   = graph.pre_target_nodes
train_target_nodes = graph.train_target_nodes
valid_target_nodes = graph.valid_target_nodes
test_target_nodes  = graph.test_target_nodes
repeat_num = int(len(pre_target_nodes) / args.batch_size // args.n_batch)


pre_target_nodes_au = np.concatenate([pre_target_nodes, np.ones(len(pre_target_nodes))]).reshape(2, -1).transpose()
# train_target_nodes_au = np.concatenate([train_target_nodes, np.ones(len(train_target_nodes))]).reshape(2, -1).transpose()

def node_classification_sample(seed, nodes, time_range):
    '''
        sub-graph sampling and label preparation for node classification:
        (1) Sample batch_size number of output nodes (papers) and their time.
    '''
    np.random.seed(seed)
    samp_nodes = np.random.choice(nodes, args.batch_size, replace = False)
    feature, times, edge_list, _, texts = sample_subgraph(graph, time_range, \
                inp = {target_type: np.concatenate([samp_nodes, np.ones(args.batch_size)]).reshape(2, -1).transpose()}, \
                sampled_depth = args.sample_depth, sampled_number = args.sample_width, feature_extractor = feature_reddit)
    
    node_feature, node_type, edge_time, edge_index, edge_type, node_dict, edge_dict = \
            to_torch(feature, times, edge_list, graph)

    x_ids = np.arange(args.batch_size)
    return node_feature, node_type, edge_time, edge_index, edge_type, x_ids, graph.y[samp_nodes]

def GPT_sample(seed, target_nodes, time_range, batch_size, feature_extractor):
    np.random.seed(seed)
    samp_target_nodes = target_nodes[np.random.choice(len(target_nodes), batch_size)]
    threshold   = 0.5
    feature, times, edge_list, _, attr = sample_subgraph(graph, time_range, \
                inp = {target_type: samp_target_nodes}, feature_extractor = feature_extractor, \
                    sampled_depth = args.sample_depth, sampled_number = args.sample_width)
    rem_edge_list = defaultdict(  #source_type
                        lambda: defaultdict(  #relation_type
                            lambda: [] # [target_id, source_id] 
                                ))
    
    ori_list = {}
    for source_type in edge_list[target_type]:
        ori_list[source_type] = {}
        for relation_type in edge_list[target_type][source_type]:
            ori_list[source_type][relation_type] = np.array(edge_list[target_type][source_type][relation_type])
            el = []
            for target_ser, source_ser in edge_list[target_type][source_type][relation_type]:
                if target_ser < source_ser:
                    if relation_type not in rel_stop_list and target_ser < batch_size and \
                           np.random.random() > threshold:
                        rem_edge_list[source_type][relation_type] += [[target_ser, source_ser]]
                        continue
                    el += [[target_ser, source_ser]]
                    el += [[source_ser, target_ser]]
            el = np.array(el)
            edge_list[target_type][source_type][relation_type] = el
            
            if relation_type == 'self':
                continue
                
    '''
        Adding feature nodes:
    '''
    n_target_nodes = len(feature[target_type])
    feature[target_type] = np.concatenate((feature[target_type], np.zeros([batch_size, feature[target_type].shape[1]])))
    times[target_type]   = np.concatenate((times[target_type], times[target_type][:batch_size]))

    for source_type in edge_list[target_type]:
        for relation_type in edge_list[target_type][source_type]:
            el = []
            for target_ser, source_ser in edge_list[target_type][source_type][relation_type]:
                if target_ser < batch_size:
                    if relation_type == 'self':
                        el += [[target_ser + n_target_nodes, target_ser + n_target_nodes]]
                    else:
                        el += [[target_ser + n_target_nodes, source_ser]]
            if len(el) > 0:
                edge_list[target_type][source_type][relation_type] = \
                    np.concatenate((edge_list[target_type][source_type][relation_type], el))


    rem_edge_lists = {}
    for source_type in rem_edge_list:
        rem_edge_lists[source_type] = {}
        for relation_type in rem_edge_list[source_type]:
            rem_edge_lists[source_type][relation_type] = np.array(rem_edge_list[source_type][relation_type])
    del rem_edge_list
          
    return to_torch(feature, times, edge_list, graph), rem_edge_lists, ori_list, \
            attr[:batch_size], (n_target_nodes, n_target_nodes + batch_size)

def prepare_data_pr(pool):
    '''
        Sampled and prepare training and validation data using multi-process parallization.
    '''
    jobs = []
    for _ in np.arange(args.n_batch):
        p = pool.apply_async(node_classification_sample, args=(randint(), train_target_nodes, {1: True}))
        jobs.append(p)
    for _ in range(args.valid_batch):
        p = pool.apply_async(node_classification_sample, args=(randint(), valid_target_nodes, {1: True}))
        jobs.append(p)
    return jobs

def prepare_data_au(pool):
    jobs = []
    for _ in np.arange(args.n_batch):
        jobs.append(pool.apply_async(GPT_sample, args=(randint(), pre_target_nodes_au, {1: True}, args.batch_size, feature_reddit)))
    # jobs.append(pool.apply_async(GPT_sample, args=(randint(), train_target_nodes, {1: True}, args.batch_size, feature_reddit)))
    return jobs


data, rem_edge_list, ori_edge_list, _, _ = GPT_sample(randint(), pre_target_nodes_au, {1: True}, args.batch_size, feature_reddit)
node_feature, node_type, edge_time, edge_index, edge_type, node_dict, edge_dict = data
types = graph.get_types()

# model
'''
    Initialize GNN (model is specified by conv_name)
'''
gnn = GNN(conv_name = args.conv_name, in_dim = len(graph.node_feature[target_type]['emb'].values[0]), n_hid = args.n_hid, \
          n_heads = args.n_heads, n_layers = args.n_layers, dropout = args.dropout, num_types = len(types), \
          num_relations = len(graph.get_meta_graph()) + 1, prev_norm = args.prev_norm, last_norm = args.last_norm, use_RTE = False)
gnn_meta = GNN(conv_name = args.conv_name, in_dim = len(graph.node_feature[target_type]['emb'].values[0]), n_hid = args.n_hid, \
          n_heads = args.n_heads, n_layers = args.n_layers, dropout = args.dropout, num_types = len(types), \
          num_relations = len(graph.get_meta_graph()) + 1, prev_norm = args.prev_norm, last_norm = args.last_norm, use_RTE = False)
if args.use_pretrain:
    gnn.load_state_dict(load_gnn(torch.load(args.pretrain_model_dir, map_location=CUDA_STR)), strict=False)
    print('Load Pre-trained Model from (%s)' % args.pretrain_model_dir)

'''
    Initialize model for the primary task
'''
classifier = Classifier(args.n_hid, graph.y.max().item() + 1)
classifier_meta = Classifier(args.n_hid, graph.y.max().item() + 1)
model_pr = nn.Sequential(gnn, classifier).to(device)
print('Primary Model:\n')
print(model_pr)
print('\n')
params = sum([p.numel() for p in model_pr.parameters()])
print('Parameters: {}'.format(params))

'''
    Initialize model for the auxiliary tasks
'''
if args.attr_type == 'text':
    from gensim.models import Word2Vec

    w2v_model = Word2Vec.load(args.w2v_dir)
    n_tokens = len(w2v_model.wv.vocab)
    attr_decoder = RNNModel(n_word=n_tokens, ninp=gnn.n_hid, \
                            nhid=w2v_model.vector_size, nlayers=2)
    attr_decoder.from_w2v(torch.FloatTensor(w2v_model.wv.vectors))
else:
    attr_decoder = Matcher(gnn.n_hid, gnn.in_dim)

gpt_gnn = GPT_GNN(gnn=gnn, rem_edge_list=rem_edge_list, attr_decoder=attr_decoder, \
                  neg_queue_size=0, types=types, neg_samp_num=args.neg_samp_num, device=device)
gpt_gnn_meta = GPT_GNN(gnn=gnn_meta, rem_edge_list=rem_edge_list, attr_decoder=attr_decoder, \
                  neg_queue_size=0, types=types, neg_samp_num=args.neg_samp_num, device=device)
# gpt_gnn.init_emb.data = node_feature[node_type == node_dict[target_type][1]].mean(dim=0).detach()
init_emb = node_feature[node_type == node_dict[target_type][1]].mean(dim=0).detach()
gpt_gnn = gpt_gnn.to(device)

'''
Initialize weight model for multi-tasks
'''
vnet = Weight(6, args.weight_emb_dim, 1, args.act_type).to(device)
optimizer_v = torch.optim.Adam(vnet.parameters(), lr=args.wlr)

# Optimizer
# params = list(gpt_gnn.parameters()) + list(model_pr.parameters())
params = list(gpt_gnn.parameters()) + list(classifier.parameters())
optimizer = torch.optim.AdamW(params, weight_decay = 1e-2, eps=1e-06, lr = args.max_lr)

params_meta = list(gpt_gnn_meta.parameters()) + list(classifier_meta.parameters())
optimizer_meta = torch.optim.AdamW(params_meta, weight_decay = 1e-2, eps=1e-06, lr = args.max_lr)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, repeat_num * args.n_batch, eta_min=1e-6)
scheduler_meta = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_meta, repeat_num * args.n_batch, eta_min=1e-6)

# prepare training
best_val = 0
train_step = 0
train_step_meta = 0
stats = []
res = []


# prepare for cosine similarty
link_private_param_name = [p[0] for p in gpt_gnn.named_parameters() if 'params' in p[0] and p[1].requires_grad]
attr_private_param_name = [p[0] for p in gpt_gnn.named_parameters() if 'attr_decoder' in p[0] and p[1].requires_grad]
private_param_name = link_private_param_name + attr_private_param_name
share_param_name = []
for m in [gpt_gnn, classifier]:
    for n, p in m.named_parameters():
        if n not in private_param_name and p.requires_grad:
            share_param_name.append(n)

cos_ = nn.CosineSimilarity()
all_param_names = set(link_private_param_name + attr_private_param_name + share_param_name)
dummy_flag = False #激活一下梯度
dummy_flag_meta = False
link_cos_mean = 0
attr_cos_mean = 0

pool = mp.Pool(args.n_pool)
st = time.time()
jobs_pr = prepare_data_pr(pool)
jobs_au = prepare_data_au(pool)
criterion = nn.NLLLoss(reduce=False)

def getCos(args, device, classifier, params, gpt_gnn, optimizer, share_param_name, cos_, loss_pr_mean, loss_link_mean, loss_attr_mean):
    optimizer.zero_grad()
    loss_pr_mean.backward(retain_graph=True)
    torch.nn.utils.clip_grad_norm_(params, args.clip)
    pr_share_grad = record_grad([gpt_gnn, classifier], share_param_name)
    pr_grad_flat = torch.cat([p.clone().flatten().to(device) for n,p in pr_share_grad.items()])

    optimizer.zero_grad()
    loss_link_mean.backward(retain_graph=True)
    torch.nn.utils.clip_grad_norm_(params, args.clip)
    link_share_grad = record_grad([gpt_gnn, classifier], share_param_name)
    link_grad_flat = torch.cat([p.clone().flatten().to(device) for n, p in link_share_grad.items()])

    optimizer.zero_grad()
    loss_attr_mean.backward(retain_graph=True)
    torch.nn.utils.clip_grad_norm_(params, args.clip)
    attr_share_grad = record_grad([gpt_gnn, classifier], share_param_name)
    attr_grad_flat = torch.cat([p.clone().flatten().to(device) for n, p in attr_share_grad.items()])

    pr_link_cos = cos_(pr_grad_flat.view(1, -1), link_grad_flat.view(1, -1))
    pr_attr_cos = cos_(pr_grad_flat.view(1, -1), attr_grad_flat.view(1, -1))

    return pr_link_cos, pr_attr_cos

# Training + validation + testing
print('Start training...')
for epoch in np.arange(args.n_epoch) + 1:
    '''
        Prepare Training and Validation Data for the primary task
    '''
    train_data_pr = [job.get() for job in jobs_pr[:args.n_batch]]
    valid_data_pr = [job.get() for job in jobs_pr[args.n_batch:]]
    pool.close()
    pool.join()
    '''
            Prepare Training and Validation Data for the auxiliary tasks
    '''
    pool = mp.Pool(args.n_pool)
    train_data_au = [job.get() for job in jobs_au]
    pool.close()
    pool.join()

    '''
            After the data is collected, close the pool and then reopen it.
    '''
    pool = mp.Pool(args.n_pool)
    jobs_pr = prepare_data_pr(pool)
    jobs_au = prepare_data_au(pool)
    et = time.time()
    print('Data Preparation: %.1fs' % (et - st))
    '''
            Prepare for Training
    '''
    train_losses = []
    train_pr_losses = []
    train_link_losses = []
    train_attr_losses = []
    gpt_gnn.neg_queue_size = args.queue_size * epoch // args.n_epoch
    gpt_gnn.train()
    model_pr.train()
    gpt_gnn_meta.train()
    gnn_meta.train()
    classifier_meta.train()

    pr_link_cos = 0
    pr_attr_cos = 0
    '''
            Train on primary task and auxiliary tasks (2014 <= time <= 2016)
    '''
    for i in range(args.n_batch):
        node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = train_data_pr[i]
        data, rem_edge_list, ori_edge_list, attr, (start_idx, end_idx) = train_data_au[i]
        node_feature_au, node_type_au, edge_time_au, edge_index_au, edge_type_au, node_dict_au, edge_dict_au = data
        node_feature_au = node_feature_au.detach()
        node_feature_au[start_idx: end_idx] = init_emb
        
        loss_pr_meta = 0 
        for j in range(args.n_fold):
            # meta model
            gnn_meta.to(device)
            gnn_meta.load_state_dict(gnn.state_dict())
            gpt_gnn_meta.to(device)
            gpt_gnn_meta.load_state_dict(gpt_gnn.state_dict())
            classifier_meta.to(device)
            classifier_meta.load_state_dict(classifier.state_dict())


            # primary task
            node_rep = gnn_meta.forward(node_feature.to(device), node_type.to(device), edge_time.to(device), edge_index.to(device), edge_type.to(device))
            res = classifier_meta.forward(node_rep[x_ids])
            loss_pr = criterion(res, ylabel.to(device))   # loss per example

            # auxiliary task

            node_emb_au = gpt_gnn_meta.gnn(node_feature_au.to(device), node_type_au.to(device), edge_time_au.to(device), \
                                edge_index_au.to(device), edge_type_au.to(device))
            loss_link, _ = gpt_gnn_meta.link_loss(node_emb_au, rem_edge_list, ori_edge_list, node_dict_au, target_type,
                                            use_queue=True, update_queue=True)
            if args.attr_type == 'text':
                loss_attr = gpt_gnn_meta.text_loss(node_emb_au[start_idx: end_idx], attr, w2v_model, device)
            else:
                # pdb.set_trace()
                loss_attr = gpt_gnn_meta.feat_loss(node_emb_au[start_idx: end_idx], torch.FloatTensor(attr).to(device))
            # loss_link = 0
            # loss_attr = 0
            # for log
            loss_pr_mean = loss_pr.mean()
            loss_link_mean = loss_link.mean()
            loss_attr_mean = loss_attr.mean()
            if not dummy_flag_meta:
                optimizer_meta.zero_grad()
                loss = loss_pr_mean + loss_link_mean + loss_attr_mean
                loss.backward(retain_graph=True)
                clean_param_name([gpt_gnn_meta, classifier_meta], share_param_name)
                clean_param_name([gpt_gnn_meta, classifier_meta], link_private_param_name)
                clean_param_name([gpt_gnn_meta, classifier_meta], attr_private_param_name)
                dummy_flag_meta = True
            else:
                loss = 0

            pr_link_cos, pr_attr_cos = getCos(args, device, classifier_meta, params_meta, gpt_gnn_meta, optimizer_meta, share_param_name, cos_, loss_pr_mean, loss_link_mean, loss_attr_mean)

            # embeddings for v-net
            loss_pr_emb = torch.stack((loss_pr, \
                    torch.ones([len(loss_pr)]).to(device), \
                    torch.zeros([len(loss_pr)]).to(device), \
                    torch.zeros([len(loss_pr)]).to(device), \
                    torch.full([len(loss_pr)], pr_link_cos.item()).to(device), \
                    torch.full([len(loss_pr)], pr_attr_cos.item()).to(device)
                ))
            loss_link_emb = torch.stack((loss_link, \
                    torch.zeros([len(loss_link)]).to(device), \
                    torch.ones([len(loss_link)]).to(device), \
                    torch.zeros([len(loss_link)]).to(device), \
                    torch.full([len(loss_link)], pr_link_cos.item()).to(device), \
                    torch.full([len(loss_link)], pr_attr_cos.item()).to(device)
                ))
            loss_attr_emb = torch.stack((loss_attr, \
                    torch.zeros([len(loss_attr)]).to(device), \
                    torch.zeros([len(loss_attr)]).to(device), \
                    torch.ones([len(loss_attr)]).to(device), \
                    torch.full([len(loss_attr)], pr_link_cos.item()).to(device), \
                    torch.full([len(loss_attr)], pr_attr_cos.item()).to(device)
                ))

            loss_pr_emb = loss_pr_emb.transpose(1,0)
            loss_link_emb = loss_link_emb.transpose(1,0)
            loss_attr_emb = loss_attr_emb.transpose(1,0)

            # compute weight
            v_pr = vnet(loss_pr_emb)
            v_link = vnet(loss_link_emb)
            v_attr = vnet(loss_attr_emb)

            # compute loss
            loss_pr_avg = (loss_pr * v_pr).mean()
            loss_link_avg = (loss_link * v_link).mean()
            loss_attr_avg = (loss_attr * v_attr).mean()
            loss_meta = loss_pr_avg + loss_link_avg + loss_attr_avg

            # one step update of model parameter (fake) (Eq.6)
            optimizer_meta.zero_grad()
            loss_meta.backward()
            torch.nn.utils.clip_grad_norm_(params_meta, args.clip)
            optimizer_meta.step()
            train_step_meta += 1
            scheduler_meta.step(train_step_meta)

            # primary loss with updated parameter (Eq.7)
            node_rep = gnn_meta.forward(node_feature.to(device), node_type.to(device), edge_time.to(device),
                                        edge_index.to(device), edge_type.to(device))
            res = classifier_meta.forward(node_rep[x_ids])
            loss_pr_meta += criterion(res, ylabel.to(device)).mean() #scalar

        # backward and update v-net params (Eq.9)
        optimizer_v.zero_grad()
        loss_pr_meta.backward()
        optimizer_v.step()

        # with the updated weight, update model parameters (true) (Eq.8)
        # primary task
        node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = train_data_pr[i]
        node_rep = gnn.forward(node_feature.to(device), node_type.to(device), edge_time.to(device),
                               edge_index.to(device), edge_type.to(device))
        res = classifier.forward(node_rep[x_ids])
        loss_pr = criterion(res, ylabel.to(device))   # loss per example
        # auxiliary task
        data, rem_edge_list, ori_edge_list, attr, (start_idx, end_idx) = train_data_au[i]
        node_feature_au, node_type_au, edge_time_au, edge_index_au, edge_type_au, node_dict_au, edge_dict_au = data
        node_feature_au = node_feature_au.detach()
        node_feature_au[start_idx: end_idx] = init_emb
        node_emb_au = gpt_gnn.gnn(node_feature_au.to(device), node_type_au.to(device), edge_time_au.to(device), \
                                  edge_index_au.to(device), edge_type_au.to(device))
        loss_link, _ = gpt_gnn.link_loss(node_emb_au, rem_edge_list, ori_edge_list, node_dict_au, target_type,
                                         use_queue=True, update_queue=True)
        if args.attr_type == 'text':
            loss_attr = gpt_gnn.text_loss(node_emb_au[start_idx: end_idx], attr, w2v_model, device)
        else:
            loss_attr = gpt_gnn.feat_loss(node_emb_au[start_idx: end_idx], torch.FloatTensor(attr).to(device))
        # for log
        loss_pr_mean = loss_pr.mean()
        loss_link_mean = loss_link.mean()
        loss_attr_mean = loss_attr.mean()
        if not dummy_flag:
            optimizer.zero_grad()
            loss = loss_pr_mean + loss_link_mean + loss_attr_mean
            loss.backward(retain_graph=True)
            clean_param_name([gpt_gnn, classifier], share_param_name)
            clean_param_name([gpt_gnn, classifier], link_private_param_name)
            clean_param_name([gpt_gnn, classifier], attr_private_param_name)
            dummy_flag = True

        pr_link_cos, pr_attr_cos = getCos(args, device, classifier, params, gpt_gnn, optimizer, share_param_name, cos_, loss_pr_mean, loss_link_mean, loss_attr_mean)

        print('link cos: %f\tattr cos: %f'%(pr_link_cos.item(), pr_attr_cos.item()))

        # embeddings for v-net
        loss_pr_emb = torch.stack((loss_pr, \
                torch.ones([len(loss_pr)]).to(device), \
                torch.zeros([len(loss_pr)]).to(device), \
                torch.zeros([len(loss_pr)]).to(device), \
                torch.full([len(loss_pr)], pr_link_cos.item()).to(device), \
                torch.full([len(loss_pr)], pr_attr_cos.item()).to(device)
            ))
        loss_link_emb = torch.stack((loss_link, \
                torch.zeros([len(loss_link)]).to(device), \
                torch.ones([len(loss_link)]).to(device), \
                torch.zeros([len(loss_link)]).to(device), \
                torch.full([len(loss_link)], pr_link_cos.item()).to(device), \
                torch.full([len(loss_link)], pr_attr_cos.item()).to(device)
            ))
        loss_attr_emb = torch.stack((loss_attr, \
                torch.zeros([len(loss_attr)]).to(device), \
                torch.zeros([len(loss_attr)]).to(device), \
                torch.ones([len(loss_attr)]).to(device), \
                torch.full([len(loss_attr)], pr_link_cos.item()).to(device), \
                torch.full([len(loss_attr)], pr_attr_cos.item()).to(device)
            ))
        # embeddings for v-net
        loss_pr_emb = loss_pr_emb.transpose(1, 0)
        loss_link_emb = loss_link_emb.transpose(1, 0)
        loss_attr_emb = loss_attr_emb.transpose(1, 0)

        # compute weight
        with torch.no_grad():
            v_pr = vnet(loss_pr_emb)
            v_link = vnet(loss_link_emb)
            v_attr = vnet(loss_attr_emb)

        # compute loss
        loss_pr_avg = (loss_pr * v_pr).mean()
        loss_link_avg = (loss_link * v_link).mean()
        loss_attr_avg = (loss_attr * v_attr).mean()
        loss_pr_avg_weighted = loss_pr.mean()
        loss_link_avg_weighted = loss_link.mean()
        loss_attr_avg_weighted = loss_attr.mean()

        print((
                  "Epoch: %d  Batch: %d  Train Loss Pr: %.2f  Train Loss Link: %.2f  Train Loss Attr: %.2f  Pr_Weight_Mean: %.4f Link_Weight_Mean: %.4f Attr_Weight_Mean: %.4f Pr_Weight_Std: %.4f Link_Weight_Std: %.4f Attr_Weight_Std: %.4f ") %
              (epoch, i, loss_pr_avg_weighted, loss_link_avg_weighted, loss_attr_avg_weighted, v_pr.mean().item(), v_link.mean().item(),
               v_attr.mean().item(), v_pr.std().item(), v_link.std().item(), v_attr.std().item()))

        # total loss
        loss = loss_pr_avg + loss_link_avg + loss_attr_avg
        train_pr_losses += [loss_pr_avg_weighted.cpu().detach().tolist()]
        train_link_losses += [loss_link_avg_weighted.item()]
        train_attr_losses += [loss_attr_avg_weighted.item()]

        # optimize model parameters
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(params, args.clip)
        optimizer.step()
        train_step += 1
        scheduler.step(train_step)
        del res, loss_pr, loss_attr, loss_link, loss_pr_emb, loss_link_emb, loss_attr_emb, v_pr, v_link, v_attr
        del loss_pr_mean, loss_attr_mean, loss_link_mean, loss_pr_avg_weighted, loss_link_avg_weighted, loss_attr_avg_weighted


    '''
        Valid with only primary task (2017 <= time <= 2017)
    '''
    model_pr.eval()
    with torch.no_grad():
        valid_loss_pr_mean = 0
        valid_loss_link_mean = 0
        valid_loss_attr_mean = 0
        valid_res = []
        loss = 0
        for i in range(args.valid_batch):
            node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = valid_data_pr[i]
            node_rep = gnn.forward(node_feature.to(device), node_type.to(device), \
                                    edge_time.to(device), edge_index.to(device), edge_type.to(device))
            res  = classifier.forward(node_rep[x_ids])
            loss += criterion(res, ylabel.to(device)).mean() / args.valid_batch #scalar
            # loss = criterion(res, ylabel.to(device))
        
            '''
                Calculate Valid F1. Update the best model based on highest F1 score.
            '''
            valid_res += [f1_score(res.argmax(dim=1).cpu().tolist(), ylabel.tolist(), average='micro')]
        
        valid_f1 = np.average(valid_res)
        if valid_f1 > best_val:
            best_val = valid_f1
            torch.save(model_pr,
                    os.path.join(args.model_dir, args.task_name + '_' + args.conv_name + '_' + args.model_add_name))
            print('UPDATE!!!')
        
        st = time.time()
        print(("Epoch: %d (%.1fs)  LR: %.5f Train Loss: %.2f  Valid Loss: %.2f  Valid F1: %.4f") % \
              (epoch, (st-et), optimizer.param_groups[0]['lr'], np.average(train_pr_losses), \
                    loss.cpu().detach().tolist(), valid_f1))
        del res, loss
    del train_data_pr, valid_data_pr, train_data_au
    '''
        Test 
    '''
    if epoch > 0 and epoch % 5 == 0:
        best_model = torch.load(os.path.join(args.model_dir, args.task_name + '_' + args.conv_name + '_' + args.model_add_name), map_location=CUDA_STR).to(device)
        best_model.eval()
        gnn_best, classifier_best = best_model
        with torch.no_grad():
            test_res = []
            for _ in range(10):
                node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = \
                            node_classification_sample(randint(), test_target_nodes, {1: True})
                paper_rep = gnn_best.forward(node_feature.to(device), node_type.to(device), \
                            edge_time.to(device), edge_index.to(device), edge_type.to(device))[x_ids]
                res = classifier_best.forward(paper_rep)
                test_f1 = f1_score(res.argmax(dim=1).cpu().tolist(), ylabel.tolist(), average='micro')
                test_res += [test_f1]
        print('Best Test F1: %.4f' % np.average(test_res))

'''
    Evaluate the trained model via test set (time >= 2018)
'''

best_model = torch.load(os.path.join(args.model_dir, args.task_name + '_' + args.conv_name + '_' + args.model_add_name), map_location=CUDA_STR).to(device)
best_model.eval()
gnn, classifier = best_model
with torch.no_grad():
    test_res = []
    for _ in range(10):
        node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = \
                    node_classification_sample(randint(), test_target_nodes, {1: True})
        paper_rep = gnn.forward(node_feature.to(device), node_type.to(device), \
                    edge_time.to(device), edge_index.to(device), edge_type.to(device))[x_ids]
        res = classifier.forward(paper_rep)
        test_f1 = f1_score(res.argmax(dim=1).cpu().tolist(), ylabel.tolist(), average='micro')
        test_res += [test_f1]
    print('Best Test F1: %.4f' % np.average(test_res))

<filename>example_reddit/multi_reddit_all.py
from functools import reduce
from GPT_GNN.data import *
from GPT_GNN.model_meta import *
import torch.nn as nn
from warnings import filterwarnings
import numpy as np
from sklearn.metrics import f1_score
import pdb

filterwarnings("ignore")

import argparse

parser = argparse.ArgumentParser(
    description='Multi-task method, the primary task is OAG Paper-Field (L2) classification task')

'''
   Multi-task arguments 
'''
parser.add_argument('--link_ratio', type=float, default=1.0,
                    help='Ratio of loss-loss against primary-loss, range: [0-1]')
parser.add_argument('--attr_ratio', type=float, default=1.0,
                    help='Ratio of attr-loss against primary-loss, range: [0-1]')
parser.add_argument('--attr_type', type=str, default='vec',
                    choices=['text', 'vec'],
                    help='The type of attribute decoder')
parser.add_argument('--neg_samp_num', type=int, default=255,
                    help='Maximum number of negative sample for each target node.')
parser.add_argument('--queue_size', type=int, default=256,
                    help='Max size of adaptive embedding queue.')
parser.add_argument('--w2v_dir', type=str, default='/data/data1/v-bangan/GPT/w2v_all',
                    help='The address of preprocessed graph.')
parser.add_argument('--use_pretrain', help='Whether to use pre-trained model', action='store_true')
parser.add_argument('--pretrain_model_dir', type=str, default='/datadrive/models/gpt_all_cs',
                    help='The address for pretrained model.')

'''
    Dataset arguments
'''
parser.add_argument('--data_dir', type=str, default='/data/data1/v-zhehuang/HGT',
                    help='The address of preprocessed graph.')
parser.add_argument('--model_dir', type=str, default='/data/data1/v-zhehuang/GPT/gpt_all_reddit',
                    help='The address for storing the models and optimization results.')
parser.add_argument('--model_add_name', type=str, default='tmp',
                    help='Additional name.')
parser.add_argument('--task_name', type=str, default='reddit',
                    help='The name of the stored models and optimization results.')
parser.add_argument('--cuda', type=str, default=2,
                    help='Avaiable GPU ID')
parser.add_argument('--sample_depth', type=int, default=6,
                    help='How many layers within a mini-batch subgraph')
parser.add_argument('--sample_width', type=int, default=128,
                    help='How many nodes to be sampled per layer per type')

'''
   Model arguments 
'''
parser.add_argument('--conv_name', type=str, default='hgt',
                    choices=['hgt', 'gcn', 'gat', 'rgcn', 'han', 'hetgnn'],
                    help='The name of GNN filter. By default is Heterogeneous Graph Transformer (hgt)')
parser.add_argument('--n_hid', type=int, default=400,
                    help='Number of hidden dimension')
parser.add_argument('--n_heads', type=int, default=8,
                    help='Number of attention head')
parser.add_argument('--n_layers', type=int, default=3,
                    help='Number of GNN layers')
parser.add_argument('--prev_norm', help='Whether to add layer-norm on the previous layers', action='store_true')
parser.add_argument('--last_norm', help='Whether to add layer-norm on the last layers', action='store_true')
parser.add_argument('--dropout', type=int, default=0.2,
                    help='Dropout ratio')
parser.add_argument('--data_percentage', type=float, default=1.0,
                    help='Percentage of training and validation data to use')
parser.add_argument('--temperature', type=int, default=2, help='weight adaptation')

'''
    Optimization arguments
'''
parser.add_argument('--optimizer', type=str, default='adamw',
                    choices=['adamw', 'adam', 'sgd', 'adagrad'],
                    help='optimizer to use.')
parser.add_argument('--max_lr', type=float, default=1e-3,
                    help='Maximum learning rate.')
parser.add_argument('--scheduler', type=str, default='cosine',
                    help='Name of learning rate scheduler.', choices=['cycle', 'cosine'])
parser.add_argument('--n_epoch', type=int, default=200,
                    help='Number of epoch to run')
parser.add_argument('--n_pool', type=int, default=8,
                    help='Number of process to sample subgraph')
parser.add_argument('--n_batch', type=int, default=16,
                    help='Number of batch (sampled graphs) for each epoch')
parser.add_argument('--batch_size', type=int, default=256,
                    help='Number of output nodes for training')
parser.add_argument('--valid_batch', type=int, default=10,
                    help='Number of valid batch (sampled graphs) for each epoch')
parser.add_argument('--clip', type=float, default=0.5,
                    help='Gradient Norm Clipping')
'''
meta-optimizing arguments
'''
parser.add_argument('--n_fold', type=int, default=1)
parser.add_argument('--weight_emb_dim', type=int, default=1000)
parser.add_argument('--wlr', type=float, default=0.005)
parser.add_argument('--act_type', type=str, default='sigmoid')
parser.add_argument('--seed', type=int, default=43,
                    help='Number of output nodes for training')  


args = parser.parse_args()
args_print(args)
CUDA_STR = "cuda:" + str(args.cuda)
np.random.seed(args.seed)

if args.cuda != -1:
    device = torch.device("cuda:" + str(args.cuda))
else:
    device = torch.device("cpu")
print('Start Loading Graph Data...')
graph = dill.load(open(os.path.join(args.data_dir, 'graph_reddit.pk'), 'rb'))
print('Finish Loading Graph Data!')
if not os.path.isdir(args.model_dir):
    os.mkdir(args.model_dir)

# primary task pre-define
target_type = 'def'
rel_stop_list = ['self']
pre_target_nodes   = graph.pre_target_nodes
train_target_nodes = graph.train_target_nodes
valid_target_nodes = graph.valid_target_nodes
test_target_nodes  = graph.test_target_nodes
repeat_num = int(len(pre_target_nodes) / args.batch_size // args.n_batch)


pre_target_nodes_au = np.concatenate([pre_target_nodes, np.ones(len(pre_target_nodes))]).reshape(2, -1).transpose()
# train_target_nodes_au = np.concatenate([train_target_nodes, np.ones(len(train_target_nodes))]).reshape(2, -1).transpose()

def node_classification_sample(seed, nodes, time_range):
    '''
        sub-graph sampling and label preparation for node classification:
        (1) Sample batch_size number of output nodes (papers) and their time.
    '''
    np.random.seed(seed)
    samp_nodes = np.random.choice(nodes, args.batch_size, replace = False)
    feature, times, edge_list, _, texts = sample_subgraph(graph, time_range, \
                inp = {target_type: np.concatenate([samp_nodes, np.ones(args.batch_size)]).reshape(2, -1).transpose()}, \
                sampled_depth = args.sample_depth, sampled_number = args.sample_width, feature_extractor = feature_reddit)
    
    node_feature, node_type, edge_time, edge_index, edge_type, node_dict, edge_dict = \
            to_torch(feature, times, edge_list, graph)

    x_ids = np.arange(args.batch_size)
    return node_feature, node_type, edge_time, edge_index, edge_type, x_ids, graph.y[samp_nodes]

def GPT_sample(seed, target_nodes, time_range, batch_size, feature_extractor):
    np.random.seed(seed)
    samp_target_nodes = target_nodes[np.random.choice(len(target_nodes), batch_size)]
    threshold   = 0.5
    feature, times, edge_list, _, attr = sample_subgraph(graph, time_range, \
                inp = {target_type: samp_target_nodes}, feature_extractor = feature_extractor, \
                    sampled_depth = args.sample_depth, sampled_number = args.sample_width)
    rem_edge_list = defaultdict(  #source_type
                        lambda: defaultdict(  #relation_type
                            lambda: [] # [target_id, source_id] 
                                ))
    
    ori_list = {}
    for source_type in edge_list[target_type]:
        ori_list[source_type] = {}
        for relation_type in edge_list[target_type][source_type]:
            ori_list[source_type][relation_type] = np.array(edge_list[target_type][source_type][relation_type])
            el = []
            for target_ser, source_ser in edge_list[target_type][source_type][relation_type]:
                if target_ser < source_ser:
                    if relation_type not in rel_stop_list and target_ser < batch_size and \
                           np.random.random() > threshold:
                        rem_edge_list[source_type][relation_type] += [[target_ser, source_ser]]
                        continue
                    el += [[target_ser, source_ser]]
                    el += [[source_ser, target_ser]]
            el = np.array(el)
            edge_list[target_type][source_type][relation_type] = el
            
            if relation_type == 'self':
                continue
                
    '''
        Adding feature nodes:
    '''
    n_target_nodes = len(feature[target_type])
    feature[target_type] = np.concatenate((feature[target_type], np.zeros([batch_size, feature[target_type].shape[1]])))
    times[target_type]   = np.concatenate((times[target_type], times[target_type][:batch_size]))

    for source_type in edge_list[target_type]:
        for relation_type in edge_list[target_type][source_type]:
            el = []
            for target_ser, source_ser in edge_list[target_type][source_type][relation_type]:
                if target_ser < batch_size:
                    if relation_type == 'self':
                        el += [[target_ser + n_target_nodes, target_ser + n_target_nodes]]
                    else:
                        el += [[target_ser + n_target_nodes, source_ser]]
            if len(el) > 0:
                edge_list[target_type][source_type][relation_type] = \
                    np.concatenate((edge_list[target_type][source_type][relation_type], el))


    rem_edge_lists = {}
    for source_type in rem_edge_list:
        rem_edge_lists[source_type] = {}
        for relation_type in rem_edge_list[source_type]:
            rem_edge_lists[source_type][relation_type] = np.array(rem_edge_list[source_type][relation_type])
    del rem_edge_list
          
    return to_torch(feature, times, edge_list, graph), rem_edge_lists, ori_list, \
            attr[:batch_size], (n_target_nodes, n_target_nodes + batch_size)

def prepare_data_pr(pool):
    '''
        Sampled and prepare training and validation data using multi-process parallization.
    '''
    jobs = []
    for _ in np.arange(args.n_batch):
        p = pool.apply_async(node_classification_sample, args=(randint(), train_target_nodes, {1: True}))
        jobs.append(p)
    for _ in range(args.valid_batch):
        p = pool.apply_async(node_classification_sample, args=(randint(), valid_target_nodes, {1: True}))
        jobs.append(p)
    return jobs

def prepare_data_au(pool):
    jobs = []
    for _ in np.arange(args.n_batch):
        jobs.append(pool.apply_async(GPT_sample, args=(randint(), pre_target_nodes_au, {1: True}, args.batch_size, feature_reddit)))
    # jobs.append(pool.apply_async(GPT_sample, args=(randint(), train_target_nodes, {1: True}, args.batch_size, feature_reddit)))
    return jobs


data, rem_edge_list, ori_edge_list, _, _ = GPT_sample(randint(), pre_target_nodes_au, {1: True}, args.batch_size, feature_reddit)
node_feature, node_type, edge_time, edge_index, edge_type, node_dict, edge_dict = data
types = graph.get_types()

# model
'''
    Initialize GNN (model is specified by conv_name)
'''
gnn = GNN(conv_name = args.conv_name, in_dim = len(graph.node_feature[target_type]['emb'].values[0]), n_hid = args.n_hid, \
          n_heads = args.n_heads, n_layers = args.n_layers, dropout = args.dropout, num_types = len(types), \
          num_relations = len(graph.get_meta_graph()) + 1, prev_norm = args.prev_norm, last_norm = args.last_norm, use_RTE = False)
gnn_meta = GNN(conv_name = args.conv_name, in_dim = len(graph.node_feature[target_type]['emb'].values[0]), n_hid = args.n_hid, \
          n_heads = args.n_heads, n_layers = args.n_layers, dropout = args.dropout, num_types = len(types), \
          num_relations = len(graph.get_meta_graph()) + 1, prev_norm = args.prev_norm, last_norm = args.last_norm, use_RTE = False)
if args.use_pretrain:
    gnn.load_state_dict(load_gnn(torch.load(args.pretrain_model_dir, map_location=CUDA_STR)), strict=False)
    print('Load Pre-trained Model from (%s)' % args.pretrain_model_dir)

'''
    Initialize model for the primary task
'''
classifier = Classifier(args.n_hid, graph.y.max().item() + 1)
classifier_meta = Classifier(args.n_hid, graph.y.max().item() + 1)
model_pr = nn.Sequential(gnn, classifier).to(device)
print('Primary Model:\n')
print(model_pr)
print('\n')
params = sum([p.numel() for p in model_pr.parameters()])
print('Parameters: {}'.format(params))

'''
    Initialize model for the auxiliary tasks
'''
if args.attr_type == 'text':
    from gensim.models import Word2Vec

    w2v_model = Word2Vec.load(args.w2v_dir)
    n_tokens = len(w2v_model.wv.vocab)
    attr_decoder = RNNModel(n_word=n_tokens, ninp=gnn.n_hid, \
                            nhid=w2v_model.vector_size, nlayers=2)
    attr_decoder.from_w2v(torch.FloatTensor(w2v_model.wv.vectors))
else:
    attr_decoder = Matcher(gnn.n_hid, gnn.in_dim)

gpt_gnn = GPT_GNN(gnn=gnn, rem_edge_list=rem_edge_list, attr_decoder=attr_decoder, \
                  neg_queue_size=0, types=types, neg_samp_num=args.neg_samp_num, device=device)
gpt_gnn_meta = GPT_GNN(gnn=gnn_meta, rem_edge_list=rem_edge_list, attr_decoder=attr_decoder, \
                  neg_queue_size=0, types=types, neg_samp_num=args.neg_samp_num, device=device)
# gpt_gnn.init_emb.data = node_feature[node_type == node_dict[target_type][1]].mean(dim=0).detach()
init_emb = node_feature[node_type == node_dict[target_type][1]].mean(dim=0).detach()
gpt_gnn = gpt_gnn.to(device)

'''
Initialize weight model for multi-tasks
'''
vnet = Weight(6, args.weight_emb_dim, 1, args.act_type).to(device)
optimizer_v = torch.optim.Adam(vnet.parameters(), lr=args.wlr)

# Optimizer
# params = list(gpt_gnn.parameters()) + list(model_pr.parameters())
params = list(gpt_gnn.parameters()) + list(classifier.parameters())
optimizer = torch.optim.AdamW(params, weight_decay = 1e-2, eps=1e-06, lr = args.max_lr)

params_meta = list(gpt_gnn_meta.parameters()) + list(classifier_meta.parameters())
optimizer_meta = torch.optim.AdamW(params_meta, weight_decay = 1e-2, eps=1e-06, lr = args.max_lr)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, repeat_num * args.n_batch, eta_min=1e-6)
scheduler_meta = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_meta, repeat_num * args.n_batch, eta_min=1e-6)

# prepare training
best_val = 0
train_step = 0
train_step_meta = 0
stats = []
res = []


# prepare for cosine similarty
link_private_param_name = [p[0] for p in gpt_gnn.named_parameters() if 'params' in p[0] and p[1].requires_grad]
attr_private_param_name = [p[0] for p in gpt_gnn.named_parameters() if 'attr_decoder' in p[0] and p[1].requires_grad]
private_param_name = link_private_param_name + attr_private_param_name
share_param_name = []
for m in [gpt_gnn, classifier]:
    for n, p in m.named_parameters():
        if n not in private_param_name and p.requires_grad:
            share_param_name.append(n)

cos_ = nn.CosineSimilarity()
all_param_names = set(link_private_param_name + attr_private_param_name + share_param_name)
dummy_flag = False #激活一下梯度
dummy_flag_meta = False
link_cos_mean = 0
attr_cos_mean = 0

pool = mp.Pool(args.n_pool)
st = time.time()
jobs_pr = prepare_data_pr(pool)
jobs_au = prepare_data_au(pool)
criterion = nn.NLLLoss(reduce=False)

def getCos(args, device, classifier, params, gpt_gnn, optimizer, share_param_name, cos_, loss_pr_mean, loss_link_mean, loss_attr_mean):
    optimizer.zero_grad()
    loss_pr_mean.backward(retain_graph=True)
    torch.nn.utils.clip_grad_norm_(params, args.clip)
    pr_share_grad = record_grad([gpt_gnn, classifier], share_param_name)
    pr_grad_flat = torch.cat([p.clone().flatten().to(device) for n,p in pr_share_grad.items()])

    optimizer.zero_grad()
    loss_link_mean.backward(retain_graph=True)
    torch.nn.utils.clip_grad_norm_(params, args.clip)
    link_share_grad = record_grad([gpt_gnn, classifier], share_param_name)
    link_grad_flat = torch.cat([p.clone().flatten().to(device) for n, p in link_share_grad.items()])

    optimizer.zero_grad()
    loss_attr_mean.backward(retain_graph=True)
    torch.nn.utils.clip_grad_norm_(params, args.clip)
    attr_share_grad = record_grad([gpt_gnn, classifier], share_param_name)
    attr_grad_flat = torch.cat([p.clone().flatten().to(device) for n, p in attr_share_grad.items()])

    pr_link_cos = cos_(pr_grad_flat.view(1, -1), link_grad_flat.view(1, -1))
    pr_attr_cos = cos_(pr_grad_flat.view(1, -1), attr_grad_flat.view(1, -1))

    return pr_link_cos, pr_attr_cos

# Training + validation + testing
print('Start training...')
for epoch in np.arange(args.n_epoch) + 1:
    '''
        Prepare Training and Validation Data for the primary task
    '''
    train_data_pr = [job.get() for job in jobs_pr[:args.n_batch]]
    valid_data_pr = [job.get() for job in jobs_pr[args.n_batch:]]
    pool.close()
    pool.join()
    '''
            Prepare Training and Validation Data for the auxiliary tasks
    '''
    pool = mp.Pool(args.n_pool)
    train_data_au = [job.get() for job in jobs_au]
    pool.close()
    pool.join()

    '''
            After the data is collected, close the pool and then reopen it.
    '''
    pool = mp.Pool(args.n_pool)
    jobs_pr = prepare_data_pr(pool)
    jobs_au = prepare_data_au(pool)
    et = time.time()
    print('Data Preparation: %.1fs' % (et - st))
    '''
            Prepare for Training
    '''
    train_losses = []
    train_pr_losses = []
    train_link_losses = []
    train_attr_losses = []
    gpt_gnn.neg_queue_size = args.queue_size * epoch // args.n_epoch
    gpt_gnn.train()
    model_pr.train()
    gpt_gnn_meta.train()
    gnn_meta.train()
    classifier_meta.train()

    pr_link_cos = 0
    pr_attr_cos = 0
    '''
            Train on primary task and auxiliary tasks (2014 <= time <= 2016)
    '''
    for i in range(args.n_batch):
        node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = train_data_pr[i]
        data, rem_edge_list, ori_edge_list, attr, (start_idx, end_idx) = train_data_au[i]
        node_feature_au, node_type_au, edge_time_au, edge_index_au, edge_type_au, node_dict_au, edge_dict_au = data
        node_feature_au = node_feature_au.detach()
        node_feature_au[start_idx: end_idx] = init_emb
        
        loss_pr_meta = 0 
        for j in range(args.n_fold):
            # meta model
            gnn_meta.to(device)
            gnn_meta.load_state_dict(gnn.state_dict())
            gpt_gnn_meta.to(device)
            gpt_gnn_meta.load_state_dict(gpt_gnn.state_dict())
            classifier_meta.to(device)
            classifier_meta.load_state_dict(classifier.state_dict())


            # primary task
            node_rep = gnn_meta.forward(node_feature.to(device), node_type.to(device), edge_time.to(device), edge_index.to(device), edge_type.to(device))
            res = classifier_meta.forward(node_rep[x_ids])
            loss_pr = criterion(res, ylabel.to(device))   # loss per example

            # auxiliary task

            node_emb_au = gpt_gnn_meta.gnn(node_feature_au.to(device), node_type_au.to(device), edge_time_au.to(device), \
                                edge_index_au.to(device), edge_type_au.to(device))
            loss_link, _ = gpt_gnn_meta.link_loss(node_emb_au, rem_edge_list, ori_edge_list, node_dict_au, target_type,
                                            use_queue=True, update_queue=True)
            if args.attr_type == 'text':
                loss_attr = gpt_gnn_meta.text_loss(node_emb_au[start_idx: end_idx], attr, w2v_model, device)
            else:
                # pdb.set_trace()
                loss_attr = gpt_gnn_meta.feat_loss(node_emb_au[start_idx: end_idx], torch.FloatTensor(attr).to(device))
            # loss_link = 0
            # loss_attr = 0
            # for log
            loss_pr_mean = loss_pr.mean()
            loss_link_mean = loss_link.mean()
            loss_attr_mean = loss_attr.mean()
            if not dummy_flag_meta:
                optimizer_meta.zero_grad()
                loss = loss_pr_mean + loss_link_mean + loss_attr_mean
                loss.backward(retain_graph=True)
                clean_param_name([gpt_gnn_meta, classifier_meta], share_param_name)
                clean_param_name([gpt_gnn_meta, classifier_meta], link_private_param_name)
                clean_param_name([gpt_gnn_meta, classifier_meta], attr_private_param_name)
                dummy_flag_meta = True
            else:
                loss = 0

            pr_link_cos, pr_attr_cos = getCos(args, device, classifier_meta, params_meta, gpt_gnn_meta, optimizer_meta, share_param_name, cos_, loss_pr_mean, loss_link_mean, loss_attr_mean)

            # embeddings for v-net
            loss_pr_emb = torch.stack((loss_pr, \
                    torch.ones([len(loss_pr)]).to(device), \
                    torch.zeros([len(loss_pr)]).to(device), \
                    torch.zeros([len(loss_pr)]).to(device), \
                    torch.full([len(loss_pr)], pr_link_cos.item()).to(device), \
                    torch.full([len(loss_pr)], pr_attr_cos.item()).to(device)
                ))
            loss_link_emb = torch.stack((loss_link, \
                    torch.zeros([len(loss_link)]).to(device), \
                    torch.ones([len(loss_link)]).to(device), \
                    torch.zeros([len(loss_link)]).to(device), \
                    torch.full([len(loss_link)], pr_link_cos.item()).to(device), \
                    torch.full([len(loss_link)], pr_attr_cos.item()).to(device)
                ))
            loss_attr_emb = torch.stack((loss_attr, \
                    torch.zeros([len(loss_attr)]).to(device), \
                    torch.zeros([len(loss_attr)]).to(device), \
                    torch.ones([len(loss_attr)]).to(device), \
                    torch.full([len(loss_attr)], pr_link_cos.item()).to(device), \
                    torch.full([len(loss_attr)], pr_attr_cos.item()).to(device)
                ))

            loss_pr_emb = loss_pr_emb.transpose(1,0)
            loss_link_emb = loss_link_emb.transpose(1,0)
            loss_attr_emb = loss_attr_emb.transpose(1,0)

            # compute weight
            v_pr = vnet(loss_pr_emb)
            v_link = vnet(loss_link_emb)
            v_attr = vnet(loss_attr_emb)

            # compute loss
            loss_pr_avg = (loss_pr * v_pr).mean()
            loss_link_avg = (loss_link * v_link).mean()
            loss_attr_avg = (loss_attr * v_attr).mean()
            loss_meta = loss_pr_avg + loss_link_avg + loss_attr_avg

            # one step update of model parameter (fake) (Eq.6)
            optimizer_meta.zero_grad()
            loss_meta.backward()
            torch.nn.utils.clip_grad_norm_(params_meta, args.clip)
            optimizer_meta.step()
            train_step_meta += 1
            scheduler_meta.step(train_step_meta)

            # primary loss with updated parameter (Eq.7)
            node_rep = gnn_meta.forward(node_feature.to(device), node_type.to(device), edge_time.to(device),
                                        edge_index.to(device), edge_type.to(device))
            res = classifier_meta.forward(node_rep[x_ids])
            loss_pr_meta += criterion(res, ylabel.to(device)).mean() #scalar

        # backward and update v-net params (Eq.9)
        optimizer_v.zero_grad()
        loss_pr_meta.backward()
        optimizer_v.step()

        # with the updated weight, update model parameters (true) (Eq.8)
        # primary task
        node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = train_data_pr[i]
        node_rep = gnn.forward(node_feature.to(device), node_type.to(device), edge_time.to(device),
                               edge_index.to(device), edge_type.to(device))
        res = classifier.forward(node_rep[x_ids])
        loss_pr = criterion(res, ylabel.to(device))   # loss per example
        # auxiliary task
        data, rem_edge_list, ori_edge_list, attr, (start_idx, end_idx) = train_data_au[i]
        node_feature_au, node_type_au, edge_time_au, edge_index_au, edge_type_au, node_dict_au, edge_dict_au = data
        node_feature_au = node_feature_au.detach()
        node_feature_au[start_idx: end_idx] = init_emb
        node_emb_au = gpt_gnn.gnn(node_feature_au.to(device), node_type_au.to(device), edge_time_au.to(device), \
                                  edge_index_au.to(device), edge_type_au.to(device))
        loss_link, _ = gpt_gnn.link_loss(node_emb_au, rem_edge_list, ori_edge_list, node_dict_au, target_type,
                                         use_queue=True, update_queue=True)
        if args.attr_type == 'text':
            loss_attr = gpt_gnn.text_loss(node_emb_au[start_idx: end_idx], attr, w2v_model, device)
        else:
            loss_attr = gpt_gnn.feat_loss(node_emb_au[start_idx: end_idx], torch.FloatTensor(attr).to(device))
        # for log
        loss_pr_mean = loss_pr.mean()
        loss_link_mean = loss_link.mean()
        loss_attr_mean = loss_attr.mean()
        if not dummy_flag:
            optimizer.zero_grad()
            loss = loss_pr_mean + loss_link_mean + loss_attr_mean
            loss.backward(retain_graph=True)
            clean_param_name([gpt_gnn, classifier], share_param_name)
            clean_param_name([gpt_gnn, classifier], link_private_param_name)
            clean_param_name([gpt_gnn, classifier], attr_private_param_name)
            dummy_flag = True

        pr_link_cos, pr_attr_cos = getCos(args, device, classifier, params, gpt_gnn, optimizer, share_param_name, cos_, loss_pr_mean, loss_link_mean, loss_attr_mean)

        print('link cos: %f\tattr cos: %f'%(pr_link_cos.item(), pr_attr_cos.item()))

        # embeddings for v-net
        loss_pr_emb = torch.stack((loss_pr, \
                torch.ones([len(loss_pr)]).to(device), \
                torch.zeros([len(loss_pr)]).to(device), \
                torch.zeros([len(loss_pr)]).to(device), \
                torch.full([len(loss_pr)], pr_link_cos.item()).to(device), \
                torch.full([len(loss_pr)], pr_attr_cos.item()).to(device)
            ))
        loss_link_emb = torch.stack((loss_link, \
                torch.zeros([len(loss_link)]).to(device), \
                torch.ones([len(loss_link)]).to(device), \
                torch.zeros([len(loss_link)]).to(device), \
                torch.full([len(loss_link)], pr_link_cos.item()).to(device), \
                torch.full([len(loss_link)], pr_attr_cos.item()).to(device)
            ))
        loss_attr_emb = torch.stack((loss_attr, \
                torch.zeros([len(loss_attr)]).to(device), \
                torch.zeros([len(loss_attr)]).to(device), \
                torch.ones([len(loss_attr)]).to(device), \
                torch.full([len(loss_attr)], pr_link_cos.item()).to(device), \
                torch.full([len(loss_attr)], pr_attr_cos.item()).to(device)
            ))
        # embeddings for v-net
        loss_pr_emb = loss_pr_emb.transpose(1, 0)
        loss_link_emb = loss_link_emb.transpose(1, 0)
        loss_attr_emb = loss_attr_emb.transpose(1, 0)

        # compute weight
        with torch.no_grad():
            v_pr = vnet(loss_pr_emb)
            v_link = vnet(loss_link_emb)
            v_attr = vnet(loss_attr_emb)

        # compute loss
        loss_pr_avg = (loss_pr * v_pr).mean()
        loss_link_avg = (loss_link * v_link).mean()
        loss_attr_avg = (loss_attr * v_attr).mean()
        loss_pr_avg_weighted = loss_pr.mean()
        loss_link_avg_weighted = loss_link.mean()
        loss_attr_avg_weighted = loss_attr.mean()

        print((
                  "Epoch: %d  Batch: %d  Train Loss Pr: %.2f  Train Loss Link: %.2f  Train Loss Attr: %.2f  Pr_Weight_Mean: %.4f Link_Weight_Mean: %.4f Attr_Weight_Mean: %.4f Pr_Weight_Std: %.4f Link_Weight_Std: %.4f Attr_Weight_Std: %.4f ") %
              (epoch, i, loss_pr_avg_weighted, loss_link_avg_weighted, loss_attr_avg_weighted, v_pr.mean().item(), v_link.mean().item(),
               v_attr.mean().item(), v_pr.std().item(), v_link.std().item(), v_attr.std().item()))

        # total loss
        loss = loss_pr_avg + loss_link_avg + loss_attr_avg
        train_pr_losses += [loss_pr_avg_weighted.cpu().detach().tolist()]
        train_link_losses += [loss_link_avg_weighted.item()]
        train_attr_losses += [loss_attr_avg_weighted.item()]

        # optimize model parameters
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(params, args.clip)
        optimizer.step()
        train_step += 1
        scheduler.step(train_step)
        del res, loss_pr, loss_attr, loss_link, loss_pr_emb, loss_link_emb, loss_attr_emb, v_pr, v_link, v_attr
        del loss_pr_mean, loss_attr_mean, loss_link_mean, loss_pr_avg_weighted, loss_link_avg_weighted, loss_attr_avg_weighted


    '''
        Valid with only primary task (2017 <= time <= 2017)
    '''
    model_pr.eval()
    with torch.no_grad():
        valid_loss_pr_mean = 0
        valid_loss_link_mean = 0
        valid_loss_attr_mean = 0
        valid_res = []
        loss = 0
        for i in range(args.valid_batch):
            node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = valid_data_pr[i]
            node_rep = gnn.forward(node_feature.to(device), node_type.to(device), \
                                    edge_time.to(device), edge_index.to(device), edge_type.to(device))
            res  = classifier.forward(node_rep[x_ids])
            loss += criterion(res, ylabel.to(device)).mean() / args.valid_batch #scalar
            # loss = criterion(res, ylabel.to(device))
        
            '''
                Calculate Valid F1. Update the best model based on highest F1 score.
            '''
            valid_res += [f1_score(res.argmax(dim=1).cpu().tolist(), ylabel.tolist(), average='micro')]
        
        valid_f1 = np.average(valid_res)
        if valid_f1 > best_val:
            best_val = valid_f1
            torch.save(model_pr,
                    os.path.join(args.model_dir, args.task_name + '_' + args.conv_name + '_' + args.model_add_name))
            print('UPDATE!!!')
        
        st = time.time()
        print(("Epoch: %d (%.1fs)  LR: %.5f Train Loss: %.2f  Valid Loss: %.2f  Valid F1: %.4f") % \
              (epoch, (st-et), optimizer.param_groups[0]['lr'], np.average(train_pr_losses), \
                    loss.cpu().detach().tolist(), valid_f1))
        del res, loss
    del train_data_pr, valid_data_pr, train_data_au
    '''
        Test 
    '''
    if epoch > 0 and epoch % 5 == 0:
        best_model = torch.load(os.path.join(args.model_dir, args.task_name + '_' + args.conv_name + '_' + args.model_add_name), map_location=CUDA_STR).to(device)
        best_model.eval()
        gnn_best, classifier_best = best_model
        with torch.no_grad():
            test_res = []
            for _ in range(10):
                node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = \
                            node_classification_sample(randint(), test_target_nodes, {1: True})
                paper_rep = gnn_best.forward(node_feature.to(device), node_type.to(device), \
                            edge_time.to(device), edge_index.to(device), edge_type.to(device))[x_ids]
                res = classifier_best.forward(paper_rep)
                test_f1 = f1_score(res.argmax(dim=1).cpu().tolist(), ylabel.tolist(), average='micro')
                test_res += [test_f1]
        print('Best Test F1: %.4f' % np.average(test_res))

'''
    Evaluate the trained model via test set (time >= 2018)
'''

best_model = torch.load(os.path.join(args.model_dir, args.task_name + '_' + args.conv_name + '_' + args.model_add_name), map_location=CUDA_STR).to(device)
best_model.eval()
gnn, classifier = best_model
with torch.no_grad():
    test_res = []
    for _ in range(10):
        node_feature, node_type, edge_time, edge_index, edge_type, x_ids, ylabel = \
                    node_classification_sample(randint(), test_target_nodes, {1: True})
        paper_rep = gnn.forward(node_feature.to(device), node_type.to(device), \
                    edge_time.to(device), edge_index.to(device), edge_type.to(device))[x_ids]
        res = classifier.forward(paper_rep)
        test_f1 = f1_score(res.argmax(dim=1).cpu().tolist(), ylabel.tolist(), average='micro')
        test_res += [test_f1]
    print('Best Test F1: %.4f' % np.average(test_res))

Multi-task arguments Dataset arguments Model arguments Optimization arguments meta-optimizing arguments # primary task pre-define # train_target_nodes_au = np.concatenate([train_target_nodes, np.ones(len(train_target_nodes))]).reshape(2, -1).transpose() sub-graph sampling and label preparation for node classification:         (1) Sample batch_size number of output nodes (papers) and their time. #source_type #relation_type # [target_id, source_id] Adding feature nodes: Sampled and prepare training and validation data using multi-process parallization. # jobs.append(pool.apply_async(GPT_sample, args=(randint(), train_target_nodes, {1: True}, args.batch_size, feature_reddit))) # model Initialize GNN (model is specified by conv_name) Initialize model for the primary task Initialize model for the auxiliary tasks # gpt_gnn.init_emb.data = node_feature[node_type == node_dict[target_type][1]].mean(dim=0).detach() Initialize weight model for multi-tasks # Optimizer # params = list(gpt_gnn.parameters()) + list(model_pr.parameters()) # prepare training # prepare for cosine similarty #激活一下梯度 # Training + validation + testing Prepare Training and Validation Data for the primary task Prepare Training and Validation Data for the auxiliary tasks After the data is collected, close the pool and then reopen it. Prepare for Training Train on primary task and auxiliary tasks (2014 <= time <= 2016) # meta model # primary task # loss per example # auxiliary task # pdb.set_trace() # loss_link = 0 # loss_attr = 0 # for log # embeddings for v-net # compute weight # compute loss # one step update of model parameter (fake) (Eq.6) # primary loss with updated parameter (Eq.7) #scalar # backward and update v-net params (Eq.9) # with the updated weight, update model parameters (true) (Eq.8) # primary task # loss per example # auxiliary task # for log # embeddings for v-net # embeddings for v-net # compute weight # compute loss # total loss # optimize model parameters Valid with only primary task (2017 <= time <= 2017) #scalar # loss = criterion(res, ylabel.to(device)) Calculate Valid F1. Update the best model based on highest F1 score. Test Evaluate the trained model via test set (time >= 2018)