exp_tudataset.py

import argparse
import copy
import logging
import math
import time
from pathlib import Path
from torch_geometric.utils import to_networkx
import networkx as nx
import numpy as np
import scipy.sparse as sp
import torch
from tqdm import tqdm
from sklearn.preprocessing import normalize
import os
import random
import dgl
import warnings

warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
warnings.filterwarnings('ignore', category=RuntimeWarning, message='scipy._lib.messagestream.MessageStream')
from gnnutils import make_masks, train, test, add_original_graph, load_webkb, load_planetoid, load_wiki, load_bgp, \
    load_film, load_airports, load_amazon, load_coauthor, load_WikiCS, load_crocodile, load_Cora_ML

from util import get_B_sim_phi, getM_logM, load_dgl, get_A_D, load_dgl_fromPyG
from train_tudataset import train_epoch, evaluate_network, train_epoch_graph_classification, \
    train_epoch_domainadaptation
from models import Transformer, Mainmodel, Mainmodel_finetuning, Mainmodel_domainadapt

from script_classification import run_node_classification, run_node_clustering, update_evaluation_value

np.seterr(divide='ignore')


def collate(self, samples):
    graphs, labels = map(list, zip(*samples))
    labels = torch.tensor(np.array(labels)).unsqueeze(1)
    batched_graph = dgl.batch(graphs)

    return batched_graph, labels


def filter_rels(data, r):
    data = copy.deepcopy(data)
    mask = data.edge_color <= r
    data.edge_index = data.edge_index[:, mask]
    data.edge_weight = data.edge_weight[mask]
    data.edge_color = data.edge_color[mask]
    return data


from torch.utils.data import DataLoader


def run_domain_adaptation(file_name, pre_train_loader, optimizer, batch_size, device):
    model = Mainmodel_domainadapt(args, args.num_features, hidden_dim=args.dims, num_layers=args.num_layers,
                                  num_heads=args.num_heads, k_transition=args.k_transition,
                                  num_classes=args.num_classes, cp_filename=file_name, encoder=args.encoder).to(device)
    best_model = model
    best_loss = 100000000
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=5e-5)
    best_epoch = 0
    for epoch in range(1, 50):
        epoch_train_loss, reconstruction_loss = train_epoch_domainadaptation(model, args, optimizer, device,
                                                                             pre_train_loader, epoch, 1, batch_size)
        if best_loss >= epoch_train_loss:
            best_model = model
            best_epoch = epoch
            best_loss = epoch_train_loss
        if epoch - best_epoch > 20:
            break
        msg = "Epoch:%d	|Best_epoch:%d	|Train_loss:%0.4f	 |reconstruction_loss:%0.4f	" % (
        epoch, best_epoch, epoch_train_loss, reconstruction_loss)
        print(msg)
    return best_model, best_epoch


def run_pretraining(model, pre_train_loader, optimizer, batch_size, device):
    best_model = model
    best_loss = 100000000
    best_epoch = 0
    for epoch in range(1, args.pt_epoches):
        epoch_train_loss, KL_Loss, contrastive_loss, reconstruction_loss = train_epoch(model, args, optimizer, device,
                                                                                       pre_train_loader, epoch, 1,
                                                                                       batch_size)
        if best_loss >= epoch_train_loss:
            best_model = model
            best_epoch = epoch
            best_loss = epoch_train_loss
        if epoch - best_epoch > 50:
            break
        # msg1 = "out_degree1: SP %0.4f, Std %0.4f , EO %0.4f, Std %0.4f " % (SP.mean(), SP.std(),EO.mean(), EO.std())
        msg = "Epoch:%d	|Best_epoch:%d	|Train_loss:%0.4f	|KL_Loss:%0.4f	|contrastive_loss:%0.4f	|reconstruction_loss:%0.4f	" \
              % (epoch, best_epoch, epoch_train_loss, KL_Loss, contrastive_loss, reconstruction_loss)
        print(msg)
    return best_model, best_epoch


def run(i, dataset_full, num_features, num_classes):
    model = Mainmodel(args, num_features, hidden_dim=args.dims, num_layers=args.num_layers,
                      num_heads=args.num_heads, k_transition=args.k_transition, encoder=args.encoder).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=5e-5)

    best_model = model
    best_loss = 100000000

    trainset, valset, testset = dataset_full.train, dataset_full.val, dataset_full.test
    data_all = dataset_full.data_all
    print("\nTraining Graphs: ", len(trainset))
    print("Validation Graphs: ", len(valset))
    print("Test Graphs: ", len(testset))
    batch_size = args.batch_size
    train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True, collate_fn=dataset_full.collate)
    val_loader = DataLoader(valset, batch_size=batch_size, shuffle=False, collate_fn=dataset_full.collate)
    test_loader = DataLoader(testset, batch_size=batch_size, shuffle=False, collate_fn=dataset_full.collate)
    pre_train_loader = DataLoader(data_all, batch_size=batch_size, shuffle=True,
                                  collate_fn=dataset_full.collate)  # Pre-training model

    # transfer learning
    file_name_cpt = args.output_path + f'{args.dataset}_{args.encoder}_{args.dims}_{args.num_layers}_{args.k_transition}.pt'
    # if pretained is true, only run pretraining
    if args.pretrained_mode == 1:
        if not os.path.exists(file_name_cpt):
            best_model, best_epoch = run_pretraining(model, pre_train_loader, optimizer, batch_size, device)
            torch.save(best_model, file_name_cpt)
            update_evaluation_value(args.file_name, 'Best_epoch', args.index_excel, best_epoch)
            print(f"\nFinished pre-trained model ...")
        else:
            print(f"\nexists pretrained model, quiting ...")
        raise SystemExit()
    # for domain adapt and fine tunning
    else:
        ds = args.pretrained_ds
        file_name_cpt = args.output_path + f'{ds}_{args.encoder}_{args.dims}_{args.num_layers}_{args.k_transition}.pt'
        if args.domain_adapt == 1:
            file_domain_adapt = args.output_path + f'{ds}_{args.encoder}_{args.dims}_{args.num_layers}_{args.k_transition}_{args.dataset}.pt'
            if not os.path.exists(file_domain_adapt):
                print(f"\nDomain adaptation starting ...")
                best_model, best_epoch = run_domain_adaptation(file_name_cpt, pre_train_loader, optimizer, batch_size,
                                                               device)
                file_name_cpt = args.output_path + f'{ds}_{args.encoder}_{args.dims}_{args.num_layers}_{args.k_transition}_{args.dataset}.pt'
                torch.save(best_model, file_name_cpt)
            file_name_cpt = file_domain_adapt
            print(f"\nDomain adaptation finished ...")
    print(f"\nFine tunning the pre-trained model ...")
    time.sleep(0.1)
    # end transfer learning

    runs_acc = []
    for i in tqdm(range(3)):
        print(f'\n run time: {i}')
        acc, best_epoch = run_epoch_graph_classification(best_model, train_loader, val_loader, test_loader,
                                                         num_features, file_name_cpt, batch_size=batch_size)
        runs_acc.append(acc)
        time.sleep(0.1)

    runs_acc = np.array(runs_acc) * 100

    final_msg = "Graph classification: Mean %0.4f, Std %0.4f" % (runs_acc.mean(), runs_acc.std())
    print(final_msg)


    return 0


def run_epoch_graph_classification(model, train_loader, val_loader, test_loader, num_features, file_name, batch_size):
    model = Mainmodel_finetuning(args, num_features, hidden_dim=args.dims, num_layers=args.num_layers,
                                 num_heads=args.num_heads, k_transition=args.k_transition, num_classes=args.num_classes,
                                 cp_filename=file_name, encoder=args.encoder).to(device)

    best_model = model
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
    best_loss = 100000000
    t0 = time.time()
    best_epoch = 0
    per_epoch_time = []
    epoch_train_AUCs, epoch_val_AUCs, epoch_test_AUCs = [], [], []
    epoch_train_losses, epoch_val_losses = [], []
    for epoch in range(1, args.ft_epoches):
        start = time.time()
        epoch_train_loss, epoch_train_auc, optimizer = train_epoch_graph_classification(args, model, optimizer, device,
                                                                                        train_loader, epoch, batch_size)

        epoch_val_loss, epoch_val_auc = evaluate_network(args, model, optimizer, device, val_loader, epoch,
                                                         batch_size)  # (model, device, val_loader, epoch)
        _, epoch_test_auc = evaluate_network(args, model, optimizer, device, test_loader, epoch,
                                             batch_size)  # (model, device, test_loader, epoch)
        epoch_train_losses.append(epoch_train_loss)
        epoch_val_losses.append(epoch_val_loss)
        epoch_train_AUCs.append(epoch_train_auc)
        epoch_val_AUCs.append(epoch_val_auc)
        epoch_test_AUCs.append(epoch_test_auc)
        if best_loss >= epoch_train_loss:
            best_model = model;
            best_epoch = epoch;
            best_loss = epoch_train_loss
        if epoch - best_epoch > 50:
            print(f"Finish epoch - best_epoch > 100")
            break
        if epoch % 1 == 0:
            print(
                f'Epoch: {epoch}	|Best_epoch: {best_epoch}	|Train_loss: {np.round(epoch_train_loss, 6)}	|Val_loss: {np.round(epoch_val_loss, 6)}	| Train_acc: {np.round(epoch_train_auc, 6)}	| Val_acc: {np.round(epoch_val_auc, 6)}	| epoch_test_acc: {np.round(epoch_test_auc, 6)} ')

        per_epoch_time.append(time.time() - start)
        # Stop training after params['max_time'] hours
        if time.time() - t0 > 48 * 3600:
            print('-' * 89)
            print("Max_time for training elapsed")
            break
    _, test_acc = evaluate_network(args, best_model, optimizer, device, test_loader, epoch,
                                   batch_size)  # (best_model, device, test_loader, epoch)
    index = epoch_val_AUCs.index(max(epoch_val_AUCs))
    test_auc = epoch_test_AUCs[index]
    train_auc = epoch_train_AUCs[index]
    print("Train ACC: {:.4f}".format(train_auc))
    print("Test ACC: {:.4f}".format(test_auc))
    print("Convergence Time (Epochs): {:.4f}".format(epoch))
    print("TOTAL TIME TAKEN: {:.4f}s".format(time.time() - t0))
    print("AVG TIME PER EPOCH: {:.4f}s".format(np.mean(per_epoch_time)))

    return test_acc, best_epoch


def run_epoch_graph_classification(model, train_loader, val_loader, test_loader, num_features, file_name, batch_size):
    model = Mainmodel_finetuning(args, num_features, hidden_dim=args.dims, num_layers=args.num_layers,
                                 num_heads=args.num_heads, k_transition=args.k_transition, num_classes=args.num_classes,
                                 cp_filename=file_name, encoder=args.encoder).to(device)

    best_model = model
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
    best_loss = 100000000
    t0 = time.time()
    best_epoch = 0
    per_epoch_time = []
    epoch_train_AUCs, epoch_val_AUCs, epoch_test_AUCs = [], [], []
    epoch_train_losses, epoch_val_losses = [], []
    for epoch in range(1, args.ft_epoches):
        start = time.time()
        epoch_train_loss, epoch_train_auc, optimizer = train_epoch_graph_classification(args, model, optimizer, device,
                                                                                        train_loader, epoch, batch_size)

        epoch_val_loss, epoch_val_auc = evaluate_network(args, model, optimizer, device, val_loader, epoch,
                                                         batch_size)  # (model, device, val_loader, epoch)
        _, epoch_test_auc = evaluate_network(args, model, optimizer, device, test_loader, epoch,
                                             batch_size)  # (model, device, test_loader, epoch)
        epoch_train_losses.append(epoch_train_loss)
        epoch_val_losses.append(epoch_val_loss)
        epoch_train_AUCs.append(epoch_train_auc)
        epoch_val_AUCs.append(epoch_val_auc)
        epoch_test_AUCs.append(epoch_test_auc)
        if best_loss >= epoch_train_loss:
            best_model = model;
            best_epoch = epoch;
            best_loss = epoch_train_loss
        if epoch - best_epoch > 50:
            print(f"Finish epoch - best_epoch > 100")
            break
        if epoch % 1 == 0:
            print(
                f'Epoch: {epoch}	|Best_epoch: {best_epoch}	|Train_loss: {np.round(epoch_train_loss, 6)}	|Val_loss: {np.round(epoch_val_loss, 6)}	| Train_acc: {np.round(epoch_train_auc, 6)}	| Val_acc: {np.round(epoch_val_auc, 6)}	| epoch_test_acc: {np.round(epoch_test_auc, 6)} ')

        per_epoch_time.append(time.time() - start)
        # Stop training after params['max_time'] hours
        if time.time() - t0 > 48 * 3600:
            print('-' * 89)
            print("Max_time for training elapsed")
            break
    _, test_acc = evaluate_network(args, best_model, optimizer, device, test_loader, epoch,
                                   batch_size)  # (best_model, device, test_loader, epoch)
    index = epoch_val_AUCs.index(max(epoch_val_AUCs))
    test_auc = epoch_test_AUCs[index]
    train_auc = epoch_train_AUCs[index]
    print("Train ACC: {:.4f}".format(train_auc))
    print("Test ACC: {:.4f}".format(test_auc))
    print("Convergence Time (Epochs): {:.4f}".format(epoch))
    print("TOTAL TIME TAKEN: {:.4f}s".format(time.time() - t0))
    print("AVG TIME PER EPOCH: {:.4f}s".format(np.mean(per_epoch_time)))

    return test_acc, best_epoch


##################################################################################################################################
##################################################################################################################################
##################################################################################################################################
##################################################################################################################################
import collections
from collections import defaultdict
from ogb.graphproppred import PygGraphPropPredDataset
from torch_geometric.datasets import TUDataset, ZINC
from dgl.data.utils import save_graphs, load_graphs


def main():
    timestr = time.strftime("%Y%m%d-%H%M%S")
    log_file = args.dataset + "-" + timestr + ".log"
    Path("./exp_logs").mkdir(parents=True, exist_ok=True)
    logging.basicConfig(filename="exp_logs/" + log_file, filemode="w", level=logging.INFO)
    logging.info("Starting on device: %s", device)
    logging.info("Config: %s ", args)

    graph_lists = []
    graph_labels = []

    print(f'Checking dataset {args.dataset}')
    if args.dataset in ["ENZYMES", "PROTEINS", "NCI1", "NCI109", "Mutagenicity"]:
        dataset = TUDataset(root='original_datasets/' + args.dataset, use_node_attr=True, name=args.dataset)
        args.num_classes = dataset.num_classes
        args.num_features = dataset.num_features
        num_features = args.num_features
        num_classes = args.num_classes
    else:
        print("args.dataset in ENZYMES PROTEINS NCI1 NCI109")
        raise NotImplementedError
        raise SystemExit()
    path = "pts/" + args.dataset + "_k_transition_" + str(args.k_transition) + ".bin"
    if not os.path.exists(path):
        print(f"# 1. Constructing pts...")
        graph_ds = generate_graphs(dataset, args.k_transition)
        save_graphs("pts/" + args.dataset + "_k_transition_" + str(args.k_transition) + ".bin", graph_ds.graph_lists,
                    graph_ds.graph_labels)

        print(f"# 2. Loading graphs and labels...")
        graph_lists, graph_labels = load_graphs(
            "pts/" + args.dataset + "_k_transition_" + str(args.k_transition) + ".bin")
    else:
        print(f"# 3 Loading graphs and labels...")
        graph_lists, graph_labels = load_graphs(
            "pts/" + args.dataset + "_k_transition_" + str(args.k_transition) + ".bin")

    print(f"DEGUB generated graph process")

    if args.testmode == 1:
        print(f"Only generating dataset ...")
        raise SystemExit()

    print(f"# 4 Loading subgraphs ...")
    subgraph_lists = torch.load("pts/" + args.dataset + "_subgraphs_khop_" + str(args.k_transition) + ".pt")  #
    subgraph_lists = subgraph_lists['set_subgraphs']
    print(f"len(subgraph_lists): {len(subgraph_lists)}")

    print(f"# 5 Loading transition matrices...")
    dic = torch.load("pts/" + args.dataset + "_M_khop_" + str(args.k_transition) + ".pt")
    trans_logMs = dic['trans_logMs']

    # B_list = torch.load("pts/"+ args.dataset + "_B.pt")
    # B_list = B_list['B']

    # adj_list = torch.load("pts/"+ args.dataset + "_adj.pt")
    # adj_list = adj_list['adj']

    # sim_list = torch.load("pts/"+ args.dataset + "_sim.pt")
    # sim_list = sim_list['sim']

    # phi_list = torch.load("pts/"+ args.dataset + "_phi.pt")
    # phi_list = phi_list['phi']

    # dic  = torch.load("pts/"+ args.dataset + "_M.pt")
    # trans_logMs = dic['trans_logMs']

    samples_all = []
    num_features = args.num_features
    print(f'num_features: {num_features}, args.num_classes: {args.num_classes} ')
    samples_all = []
    checking_label = []
    num_node_list = []
    for i in range(len(graph_lists)):
        # if i >= 700:
        # 	print(f" testing small dataset")
        # 	break
        current_graph = graph_lists[i]
        current_label = graph_labels['glabel'][i]
        checking_label.append(current_label)
        num_node_list.append(current_graph.num_nodes())
        current_subgraphs = subgraph_lists[i]
        current_trans_logM = trans_logMs[i]
        # current_B =  B_list[i]
        # current_adj =  adj_list[i]
        # current_sim = sim_list[i]
        # current_phi = phi_list[i]

        # pair = (current_graph, current_label, current_trans_logM, current_B, current_adj,current_sim, current_phi)
        pair = (current_graph, current_label, current_subgraphs, current_trans_logM)
        samples_all.append(pair)
    random.shuffle(samples_all)

    # textfile = open("checking_label.txt", "a")
    # textfile.write(str(checking_label))
    # textfile.close()
    # textfile = open("checking_label.txt", "a")
    # textfile.write(str(checking_label))
    # textfile.close()

    dataset_full = LoadData(samples_all, args.dataset)

    runs_acc = []

    for i in tqdm(range(args.run_times)):
        acc = run(i, dataset_full, num_features, num_classes)
        runs_acc.append(acc)


# runs_acc = np.array(runs_acc) * 100
# final_msg = "Pre-training phase: Mean %0.4f, Std %0.4f" % (runs_acc.mean(), runs_acc.std())
# print(final_msg)
from molecules import MoleculeDataset


def LoadData(samples_all, DATASET_NAME):
    return MoleculeDataset(samples_all, DATASET_NAME)


def generate_graphs(dataset, k_hop):
    graph_ds = GraphClassificationDataset()
    graph_labels = []
    set_subgraphs = []
    trans_logMs = []
    miss = 0
    checking = []

    for i in range(len(dataset)):
        # if i >= 700:
        # 	print(f" testing small dataset")
        # 	break
        if i % 10 == 0:
            print(f'Processing graph_th: {i}')
            time.sleep(0.1)
        data = dataset[i]
        path = "pts/" + args.dataset + "_kstep_" + str(args.k_transition) + ".pt"
        try:
            g = load_dgl_fromPyG(data)
            if not os.path.exists(path):
                M, logM = load_bias(g)
                trans_logM = torch.from_numpy(np.array(logM)).float()
            graph_ds.graph_lists.append(g)
            trans_logMs.append(trans_logM)
            graph_labels.append(data.y)
            checking.append(data.y[0])
            graph_ds.graph_labels = {"glabel": torch.tensor(graph_labels)}
            node_ids = g.nodes()

            all_subgraphs = [dgl.khop_in_subgraph(g, individual_node, k=1)[0] for individual_node in node_ids]
            set_subgraphs.append(all_subgraphs)
        except:
            miss += 1
            print(f'Missing loading dgl graph: {i}')
    print(f"total DGL missing: {miss}")

    torch.save({"set_subgraphs": set_subgraphs}, "pts/" + args.dataset + "_subgraphs_khop_" + str(k_hop) + ".pt")
    torch.save({"trans_logMs": trans_logMs}, "pts/" + args.dataset + "_M_khop_" + str(k_hop) + ".pt")

    return graph_ds


################ checking
class GraphClassificationDataset:
    def __init__(self):
        self.graph_lists = []  # A list of DGLGraph objects
        self.graph_labels = []
        self.subgraphs = []

    # self.trans_logMs = []
    # self.B = []
    # self.adj = []
    # self.sim = []
    # self.phi = []
    def add(self, g):
        self.graph_lists.append(g)

    def __len__(self):
        return len(self.graphs)

    def __getitem__(self, i):
        # Get the i^th sample and label
        # return self.graphs[i], self.labels[i], self.trans_logM[i], self.B[i], self.adj[i], self.sim[i], self.phi[i]
        return self.graphs[i], self.labels[i], self.subgraphs[i]


# def save_info(M, logM, B, sim, phi):
# 	path = "pts/"
# 	torch.save({"M": M, "logM": logM, "B": B, "sim": sim, "phi": phi},
# 			   path + args.dataset + "_kstep_" + str(args.k_transition) + '.pt')
# 	print('save_info done.')


# def load_info(path):
# 	dic = torch.load(path)
# 	M = dic['M']
# 	logM = dic['logM']
# 	B = dic['B']
# 	sim = dic['sim']

# 	phi = dic['phi']
# 	print('load_info done.')
# 	return M, logM, B, sim, phi


def load_bias(g):
    M, logM = getM_logM(g, kstep=args.k_transition)

    return M, logM


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Experiments")

    #
    parser.add_argument("--dataset", default="ENZYMES", help="Dataset")
    parser.add_argument("--model", default="Mainmodel", help="GNN Model")

    parser.add_argument("--run_times", type=int, default=1)

    parser.add_argument("--drop", type=float, default=0.1, help="dropout")
    parser.add_argument("--custom_masks", default=True, action='store_true', help="custom train/val/test masks")

    # adding args
    parser.add_argument("--device", default="cuda:0", help="GPU ids")

    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--testmode", type=int, default=0)

    # transfer learning
    parser.add_argument("--pretrained_mode", type=int, default=0)
    parser.add_argument("--domain_adapt", type=int, default=0)
    parser.add_argument("--pretrained_ds", default="pre_training_v1", help="Loading pretrained model ")
    parser.add_argument("--d_transfer", type=int, default=32)
    parser.add_argument("--layer_relax", type=int, default=0)
    parser.add_argument("--adapt_epoches", type=int, default=50)
    # transfer learning

    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
    parser.add_argument("--pt_epoches", type=int, default=50)
    parser.add_argument("--ft_epoches", type=int, default=50)
    parser.add_argument("--useAtt", type=int, default=1)
    parser.add_argument("--dims", type=int, default=64, help="hidden dims")
    parser.add_argument("--task", default="graph_classification")
    parser.add_argument("--encoder", default="GIN")  #
    parser.add_argument("--readout_f", default="sum")  # mean set2set sum
    parser.add_argument("--recons_type", default="logM")
    parser.add_argument("--k_transition", type=int, default=1)
    parser.add_argument("--num_layers", type=int, default=5)
    parser.add_argument("--num_heads", type=int, default=4)
    parser.add_argument("--output_path", default="outputs/", help="outputs model")

    parser.add_argument("--pre_training", default="1", help="pre_training or not")
    parser.add_argument("--index_excel", type=int, default="-1", help="index_excel")
    parser.add_argument("--file_name", default="outputs_excels.xlsx", help="file_name dataset")

    ################################################################################################	args.k_transition		 python script_run.py
    #		all_subgraphs = [dgl.khop_in_subgraph(g, individual_node, k = args.k_transition)[0] for individual_node in node_ids]

    args = parser.parse_args()
    print(args)
    device = torch.device(args.device)
    main()