exp_pep_func_5.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 models import Transformer, Mainmodel, Mainmodel_finetuning

# from script_classification import run_node_classification, run_epoch_node_classification, update_evaluation_value, run_node_clustering


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_pretraining(model, pre_train_loader, optimizer, batch_size, device):
	best_model = model
	best_loss = 100000000
	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


from models import Transformer, Mainmodel, Mainmodel_finetuning, Mainmodel_domainadapt
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)
	optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=5e-5)
	best_model = model
	best_loss = 100000000
	best_epoch = 0
	num_adapt = args.adapt_epoches
	for epoch in range(1, num_adapt):
		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(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(1)):
		print(f'\nrun 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)

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


	return 0


from train_pep_func import train_epoch, evaluate_network, train_epoch_graph_classification, train_epoch_domainadaptation


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()
	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:
			if args.task == "graph_regression":
				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_ap: {np.round(epoch_train_auc, 6)}	| Val_ap: {np.round(epoch_val_auc, 6)}	| epoch_test_ap: {np.round(epoch_test_auc, 6)} ')
			elif args.task == "graph_classification":
				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_ap: {np.round(epoch_train_auc, 6)}	| Val_ap: {np.round(epoch_val_auc, 6)}	| epoch_test_ap: {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_ap = 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 AP: {:.4f}".format(train_auc))
	print("Test AP: {:.4f}".format(test_auc))
	print("Convergence Time (Epochs): {:.4f}".format(epoch))
	print("TOTAL TIME TAKEN: {:.4f}s".format(time.time() - t0))
	print("Convergence Time (Epochs): {:.4f}".format(epoch))
	print("AVG TIME PER EPOCH: {:.4f}s".format(np.mean(per_epoch_time)))
	return test_ap.cpu(), best_epoch


##################################################################################################################################
##################################################################################################################################

import collections
from collections import defaultdict
from ogb.graphproppred import PygGraphPropPredDataset
from torch_geometric.datasets import TUDataset, ZINC, LRGBDataset
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 ["Peptides-func"]:
		dataset_train = LRGBDataset(root='original_datasets/' + args.dataset, name=args.dataset, split='train')
 
		args.num_classes = dataset_train.num_classes
		args.num_features = dataset_train.num_features
		num_features = args.num_features
		num_classes = args.num_classes
		print(f"dataset_train: {len(dataset_train)} |   ")
	
		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(  )
			for i in range(1,6):
				save_graphs("pts/" + args.dataset + "_k_transition_" + str(i) + ".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")
	else:
		print(f"Only generating dataset ...")
		raise SystemExit()



	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']


	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)):
		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]

		pair = (current_graph, current_label, current_subgraphs, current_trans_logM)
		samples_all.append(pair)

	#random.shuffle(samples_all)

	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 load_bias(g):
	M, logM = getM_logM(g, kstep=5)
	return M, logM
def generate_graphs():
 
	graph_ds = GraphClassificationDataset()

	graph_labels = []

	set_subgraphs = []
	set_subgraphs2 = []
	set_subgraphs3 = []
	set_subgraphs4 = []
	set_subgraphs5 = []
	trans_logMs = []
	miss = 0
	checking = []
	dataset_list = ['train', 'val', 'test']
	for ds in dataset_list:
		print(f"processing {ds} dataset  ... ")
		dataset  = LRGBDataset(root='original_datasets/' + args.dataset, name=args.dataset, split= ds)
		for i in range(len(dataset)):
			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)
				node_ids = g.nodes()
 
				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 += data.y
				checking.append(data.y[0])

				all_subgraphs = [dgl.khop_in_subgraph(g, individual_node, k=1)[0] for individual_node in node_ids]
				set_subgraphs.append(all_subgraphs)
				all_subgraphs2 = [dgl.khop_in_subgraph(g, individual_node, k=2)[0] for individual_node in node_ids]
				set_subgraphs2.append(all_subgraphs2)

				all_subgraphs3 = [dgl.khop_in_subgraph(g, individual_node, k=3)[0] for individual_node in node_ids]
				set_subgraphs3.append(all_subgraphs3)

				all_subgraphs4 = [dgl.khop_in_subgraph(g, individual_node, k=4)[0] for individual_node in node_ids]
				set_subgraphs4.append(all_subgraphs4)

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

	graph_labels = torch.stack(graph_labels)
	graph_ds.graph_labels = {"glabel": torch.tensor(graph_labels)}

	torch.save({"set_subgraphs": set_subgraphs}, "pts/" + args.dataset +  "_subgraphs_khop_" + str(1) + ".pt")
	torch.save({"set_subgraphs": set_subgraphs2}, "pts/" + args.dataset +  "_subgraphs_khop_" + str(2) + ".pt")
	torch.save({"set_subgraphs": set_subgraphs3}, "pts/" + args.dataset +  "_subgraphs_khop_" + str(3) + ".pt")
	torch.save({"set_subgraphs": set_subgraphs4}, "pts/" + args.dataset +  "_subgraphs_khop_" + str(4) + ".pt")
	torch.save({"set_subgraphs": set_subgraphs5}, "pts/" + args.dataset +  "_subgraphs_khop_" + str(5) + ".pt")

	torch.save({"trans_logMs": trans_logMs}, "pts/" + args.dataset +  "_M_khop_" + str(1) + ".pt")
	torch.save({"trans_logMs": trans_logMs}, "pts/" + args.dataset +  "_M_khop_" + str(2) + ".pt")
	torch.save({"trans_logMs": trans_logMs}, "pts/" + args.dataset + "_M_khop_" + str(3) + ".pt")
	torch.save({"trans_logMs": trans_logMs}, "pts/" + args.dataset + "_M_khop_" + str(4) + ".pt")
	torch.save({"trans_logMs": trans_logMs}, "pts/" + args.dataset + "_M_khop_" + str(5) + ".pt")
	return graph_ds



# 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 >= 200:
# 			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 += data.y		#graph_labels.append(data.y)
# 			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}')
# 	graph_labels = torch.stack(graph_labels)
# 	graph_ds.graph_labels = {"glabel": torch.tensor(graph_labels)}
# 	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


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

	#
	parser.add_argument("--dataset", default="Peptides-func", 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= 128)
	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("--readout_f", default="sum" )	# mean set2set sum
	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("--recons_type", default="adj")
	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")

	# parser.add_argument("--param_1", type=float, default=1)
	# parser.add_argument("--param_2", type=float, default=1)
	# parser.add_argument("--param_3", type=float, default=1)
	# parser.add_argument("--test_node_degree", type=int, default= 0)

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