main_protein.py

import math
import torch
import paddle
import pgl
import numpy as np
import paddle.fluid as F
import paddle.fluid.layers as L
from pgl.contrib.ogb.nodeproppred.dataset_pgl import PglNodePropPredDataset
import time
import copy
from ogb.nodeproppred import Evaluator

from utils import to_undirected, add_self_loop, linear_warmup_decay
from model import Proteins_baseline_model, Proteins_label_embedding_model
from partition import random_partition_v2 as random_partition

import argparse
from tqdm import tqdm
evaluator = Evaluator(name='ogbn-proteins')

# place=F.CUDAPlace(6)

def get_config():
    parser = argparse.ArgumentParser()
    
    ## model_arg
    model_group=parser.add_argument_group('model_base_arg')
    model_group.add_argument('--num_layers', default=7, type=int)
    model_group.add_argument('--hidden_size', default=64, type=int)
    model_group.add_argument('--num_heads', default=4, type=int)
    model_group.add_argument('--dropout', default=0.1, type=float)
    model_group.add_argument('--attn_dropout', default=0, type=float)
    
    ## label_embed_arg
    embed_group=parser.add_argument_group('embed_arg')
    embed_group.add_argument('--use_label_e', action='store_true')
    embed_group.add_argument('--label_rate', default=0.5, type=float)
    
    ## train_arg
    train_group=parser.add_argument_group('train_arg')
    train_group.add_argument('--runs', default=10, type=int )
    train_group.add_argument('--epochs', default=2000, type=int )
    train_group.add_argument('--lr', default=0.001, type=float)
    train_group.add_argument('--place', default=-1, type=int)
    train_group.add_argument('--log_file', default='result_proteins.txt', type=str)
    return parser.parse_args()

def optimizer_func(lr=0.01):
    return F.optimizer.AdamOptimizer(learning_rate=lr)


def eval_test(parser, program, model, test_exe, graph, y_true, split_idx):

    y_pred = np.zeros_like(y_true)

    graph.node_feat["label"] = y_true 
    graph.node_feat["nid"] = np.arange(0, graph.num_nodes) 
    for subgraph in random_partition(num_clusters=5, graph=graph, shuffle=False): 
        feed_dict = model.gw.to_feed(subgraph)
        if parser.use_label_e:
            feed_dict['label'] = subgraph.node_feat["label"]
            train_idx_temp = set(split_idx['train']) & set(subgraph.node_feat["nid"])
            train_idx_temp = subgraph.reindex_from_parrent_nodes(list(train_idx_temp))
            feed_dict['label_idx'] = train_idx_temp

        batch_y_pred = test_exe.run(
                program=program,
                feed=feed_dict,
                fetch_list=[model.out_feat])[0]

        y_pred[subgraph.node_feat["nid"]] = batch_y_pred
    
    train_acc = evaluator.eval({
        'y_true': y_true[split_idx['train']],
        'y_pred': y_pred[split_idx['train']],
    })['rocauc']
    val_acc = evaluator.eval({
        'y_true': y_true[split_idx['valid']],
        'y_pred': y_pred[split_idx['valid']],
    })['rocauc']
    test_acc = evaluator.eval({
        'y_true': y_true[split_idx['test']],
        'y_pred': y_pred[split_idx['test']],
    })['rocauc']

    return train_acc, val_acc, test_acc

    
    

def train_loop(parser, start_program, main_program, test_program, 
               model, graph, label, split_idx, exe, run_id, wf=None):
    #build up training program
    exe.run(start_program)
    
    max_acc=0  # best test_acc
    max_step=0 # step for best test_acc 
    max_val_acc=0 # best val_acc
    max_cor_acc=0 # test_acc for best val_acc
    max_cor_step=0 # step for best val_acc
    #training loop

    graph.node_feat["label"] = label
    graph.node_feat["nid"] = np.arange(0, graph.num_nodes) 
    
    if parser.use_label_e:
        train_idx=copy.deepcopy(split_idx['train'])
        np.random.shuffle(train_idx[:50125])
        label_idx = train_idx[: int(50125*parser.label_rate)]
        unlabel_idx = train_idx[int(50125*parser.label_rate): ]
        label_idx_total= set(label_idx)
        unlabel_idx_total= set(unlabel_idx)   
    
    for epoch_id in tqdm(range(parser.epochs)):
        for subgraph in random_partition(num_clusters=9, graph=graph, shuffle=True): 
            #start training  
            if parser.use_label_e:
                feed_dict = model.gw.to_feed(subgraph)
                sub_idx = set(subgraph.node_feat["nid"])
                
                train_idx_temp = label_idx_total & sub_idx
                label_idx = subgraph.reindex_from_parrent_nodes(list(train_idx_temp))
                
                train_idx_temp = unlabel_idx_total & sub_idx
                unlabel_idx = subgraph.reindex_from_parrent_nodes(list(train_idx_temp))
                
                feed_dict['label'] = subgraph.node_feat["label"] 
                feed_dict['label_idx'] = label_idx
                feed_dict['train_idx'] = unlabel_idx
            else:
                feed_dict = model.gw.to_feed(subgraph)
                #feed_dict['label'] = label
                train_idx_temp = set(split_idx['train']) & set(subgraph.node_feat["nid"])
                train_idx_temp = subgraph.reindex_from_parrent_nodes(list(train_idx_temp))
                feed_dict['label'] = subgraph.node_feat["label"] 
                feed_dict['train_idx'] = train_idx_temp
            
            loss = exe.run(main_program,
                          feed=feed_dict,
                          fetch_list=[model.avg_cost])
            loss = loss[0]

        #eval result
        if (epoch_id+1) > parser.epochs*0.9:
            result = eval_test(parser, test_program, model, exe, graph, label, split_idx)
            train_acc, valid_acc, test_acc = result

            max_acc = max(test_acc, max_acc)
            if max_acc == test_acc:
                max_step=epoch_id
            max_val_acc=max(valid_acc, max_val_acc)
            if max_val_acc==valid_acc:
                max_cor_acc=test_acc
                max_cor_step=epoch_id
            max_acc=max(result[2], max_acc)
            if max_acc==result[2]:
                max_step=epoch_id
            result_t=(f'Run: {run_id:02d}, '
                      f'Epoch: {epoch_id:02d}, '
                      #f'Loss: {loss[0]:.4f}, '
                      f'Train: {100 * train_acc:.2f}%, '
                      f'Valid: {100 * valid_acc:.2f}%, '
                      f'Test: {100 * test_acc:.2f}% \n'
                      f'max_Test: {100 * max_acc:.2f}%, '
                      f'max_step: {max_step}\n'
                      f'max_val: {100 * max_val_acc:.2f}%, '
                      f'max_val_Test: {100 * max_cor_acc:.2f}%, '
                      f'max_val_step: {max_cor_step}\n'
                     )
            print(result_t)
            wf.write(result_t)
            wf.write('\n')
            wf.flush()
    return max_cor_acc

def np_scatter(idx, vals, target):
    """target[idx] += vals, but allowing for repeats in idx"""
    np.add.at(target, idx, vals)

def aggregate_node_features(graph):
    efeat = graph.edge_feat["feat"]
    graph.edge_feat["feat"] = efeat
    nfeat = np.zeros((graph.num_nodes, efeat.shape[-1]), dtype="float32")
    edges_dst = graph.edges[:, 1]
    np_scatter(edges_dst, efeat, nfeat)
    graph.node_feat["feat"] = nfeat
    


if __name__ == '__main__':
    parser = get_config()
    print('===========args==============')
    print(parser)
    print('=============================')
    
    dataset = PglNodePropPredDataset(name="ogbn-proteins")
    split_idx=dataset.get_idx_split()
    
    graph, label = dataset[0]
    aggregate_node_features(graph)
    
    place=F.CPUPlace() if parser.place <0 else F.CUDAPlace(parser.place)
    
    startup_prog = F.default_startup_program()
    train_prog = F.default_main_program()
    
    with F.program_guard(train_prog, startup_prog):
        with F.unique_name.guard():
            gw = pgl.graph_wrapper.GraphWrapper(
                    name="proteins",
                    node_feat=graph.node_feat_info(),
                    edge_feat=graph.edge_feat_info())
            
            if parser.use_label_e:
                model = Proteins_label_embedding_model(gw, parser.hidden_size, parser.num_heads, 
                                                       parser.dropout, parser.num_layers)
            else:
                model = Proteins_baseline_model(gw, parser.hidden_size, parser.num_heads, 
                                                 parser.dropout, parser.num_layers)
                
            test_prog=train_prog.clone(for_test=True)
            model.train_program()
            
        
            adam_optimizer = optimizer_func(parser.lr)#optimizer
            adam_optimizer.minimize(model.avg_cost)
    
    exe = F.Executor(place)
    
    wf = open(parser.log_file, 'w', encoding='utf-8')
    total_test_acc=0.0
    for run_i in range(parser.runs):
        total_test_acc+=train_loop(parser, startup_prog, train_prog, test_prog, model,
            graph, label, split_idx, exe, run_i, wf)
    wf.write(f'average: {100 * (total_test_acc/parser.runs):.2f}%')
    wf.close()