| import argparse |
| import os |
| import random |
| import torch |
| import pandas as pd |
| import numpy as np |
| import time |
| import torch.optim as optim |
| import scipy |
|
|
| from matplotlib import cm |
| import matplotlib.pyplot as plt |
| import json |
| import torch.nn.functional as F |
| from torch.nn.functional import softmax |
|
|
| torch.autograd.set_detect_anomaly(True) |
| import pickle |
| from torch.utils.tensorboard import SummaryWriter |
| import dataset,util |
| from model_new import Smodel |
| import model_new |
|
|
|
|
| import torch.nn as nn |
| import torchvision.transforms as transforms |
| import torchvision.datasets |
| import torchvision.models |
| import math |
| import shutil |
| import time |
| from datetime import date, timedelta,datetime |
| import torch_geometric |
| from torch_geometric.data import Data, DataLoader |
| from torch_geometric.nn import MessagePassing |
| from torch_geometric.utils import add_self_loops |
| from torch_geometric.nn import GIN,GATConv,MLP |
| from torch_geometric.nn.pool import global_mean_pool,global_add_pool |
| import csv |
|
|
| blue = lambda x: '\033[94m' + x + '\033[0m' |
| red = lambda x: '\033[31m' + x + '\033[0m' |
| green = lambda x: '\033[32m' + x + '\033[0m' |
| yellow = lambda x: '\033[33m' + x + '\033[0m' |
| greenline = lambda x: '\033[42m' + x + '\033[0m' |
| yellowline = lambda x: '\033[43m' + x + '\033[0m' |
|
|
| def get_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--model',default="our", type=str) |
| parser.add_argument('--train_batch', default=64, type=int) |
| parser.add_argument('--test_batch', default=128, type=int) |
| parser.add_argument('--share', type=str, default="0") |
| parser.add_argument('--edge_rep', type=str, default="True") |
| parser.add_argument('--batchnorm', type=str, default="True") |
| parser.add_argument('--extent_norm', type=str, default="T") |
| parser.add_argument('--spanning_tree', type=str, default="F") |
| |
| parser.add_argument('--loss_coef', default=0.1, type=float) |
| parser.add_argument('--h_ch', default=512, type=int) |
| parser.add_argument('--localdepth', type=int, default=1) |
| parser.add_argument('--num_interactions', type=int, default=4) |
| parser.add_argument('--finaldepth', type=int, default=4) |
| parser.add_argument('--classifier_depth', type=int, default=4) |
| parser.add_argument('--dropout', type=float, default=0.0) |
|
|
| parser.add_argument('--dataset', type=str, default='mnist') |
| parser.add_argument('--log', type=str, default="True") |
| parser.add_argument('--test_per_round', type=int, default=10) |
| parser.add_argument('--patience', type=int, default=30) |
| parser.add_argument('--nepoch', type=int, default=201) |
| parser.add_argument('--lr', type=float, default=1e-4) |
| parser.add_argument('--manualSeed', type=str, default="False") |
| parser.add_argument('--man_seed', type=int, default=12345) |
| |
| parser.add_argument("--targetfiles", nargs='+', type=str, default=["Dec11-14:44:32.pth","Nov13-14:30:48.pth"]) |
| args = parser.parse_args() |
| args.log=True if args.log=="True" else False |
| args.edge_rep=True if args.edge_rep=="True" else False |
| args.batchnorm=True if args.batchnorm=="True" else False |
| args.save_dir=os.path.join('./save/',args.dataset) |
| args.manualSeed=True if args.manualSeed=="True" else False |
| return args |
|
|
| args = get_args() |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| criterion=nn.CrossEntropyLoss() |
|
|
| def forward_HGT(args,data,model,mlpmodel): |
| data = data.to(device) |
| x,batch=data.pos, data['vertices'].batch |
| data["vertices"]['x']=data.pos |
| label=data.y.long().view(-1) |
| |
| output=model(data.x_dict, data.edge_index_dict) |
| if args.dataset in ["dbp"]: |
| graph_embeddings=global_add_pool(output,batch) |
| else: |
| graph_embeddings=global_add_pool(output,batch) |
| graph_embeddings.clamp_(max=1e6) |
|
|
| output=mlpmodel(graph_embeddings) |
| |
|
|
| loss = criterion(output, label) |
| return loss,output,label, graph_embeddings |
|
|
| def forward(args,data,model,mlpmodel): |
| data = data.to(device) |
| edge_index1=data['vertices', 'inside', 'vertices']['edge_index'] |
| edge_index2=data['vertices', 'apart', 'vertices']['edge_index'] |
| combined_edge_index=torch.cat([data['vertices', 'inside', 'vertices']['edge_index'],data['vertices', 'apart', 'vertices']['edge_index']],1) |
| |
| if args.spanning_tree == 'True': |
| edge_weight=torch.rand(combined_edge_index.shape[1]) + 1 |
| combined_edge_index = util.build_spanning_tree_edge(combined_edge_index, edge_weight,num_nodes=num_nodes,) |
| |
| num_edge_inside=edge_index1.shape[1] |
| x,batch=data.pos, data['vertices'].batch |
| label=data.y.long().view(-1) |
| """ |
| triplets are not the same for graphs when training |
| """ |
| num_nodes=x.shape[0] |
| edge_index_2rd, num_triplets_real, edx_jk, edx_ij = util.triplets(combined_edge_index, num_nodes) |
| |
| input_feature=torch.zeros([x.shape[0],args.h_ch],device=device) |
| output=model(input_feature,x,[edge_index1,edge_index2], edge_index_2rd,edx_jk, edx_ij,batch,num_edge_inside,args.edge_rep) |
| output=torch.cat(output,dim=1) |
| graph_embeddings=global_add_pool(output,batch) |
| graph_embeddings.clamp_(max=1e6) |
| |
| output=mlpmodel(graph_embeddings) |
| |
|
|
| loss = criterion(output, label) |
| return loss,output,label,graph_embeddings |
| def test(args,loader,model,mlpmodel,writer,reverse_mapping ): |
| y_hat, y_true,y_hat_logit = [], [], [], |
| embeddings=[] |
|
|
| loss_total, pred_num = 0, 0 |
| model.eval() |
| mlpmodel.eval() |
| with torch.no_grad(): |
| for data in loader: |
| if args.model=="our": |
| loss,output,label,embedding =forward(args,data,model,mlpmodel) |
| elif args.model in ["HGT","HAN"]: |
| loss,output,label,embedding =forward_HGT(args,data,model,mlpmodel) |
| _, pred = output.topk(1, dim=1, largest=True, sorted=True) |
| pred,label,output=pred.cpu(),label.cpu(),output.cpu() |
| y_hat += list(pred.detach().numpy().reshape(-1)) |
| y_true += list(label.detach().numpy().reshape(-1)) |
| y_hat_logit+=list(output.detach().numpy()) |
| embeddings.append(embedding) |
| |
| pred_num += len(label.reshape(-1, 1)) |
| loss_total += loss.detach() * len(label.reshape(-1, 1)) |
| |
| y_true_str=[reverse_mapping(item) for item in y_true] |
| writer.add_embedding(torch.cat(embeddings,dim=0).detach().cpu(),metadata=y_true_str,tag="numbers") |
| writer.close() |
| return loss_total/pred_num,y_hat, y_true, y_hat_logit |
| |
| def main(args,train_Loader,val_Loader,test_Loader): |
| donefiles=os.listdir(os.path.join(args.save_dir,args.model,'model')) |
| tensorboard_dir=os.path.join(args.save_dir,args.model,'log') |
| if args.dataset in ["mnist","mnist_sparse"]: |
| reverse_mapping=lambda x: x + 10 |
| |
| elif args.dataset in ["building","mbuilding"]: |
| reverse_mapping=lambda x: dataset.reverse_label_mapping[x] |
| elif args.dataset in ["sbuilding"]: |
| reverse_mapping=lambda x: dataset.single_reverse_label_mapping[x] |
| elif args.dataset in ["dbp","smnist"]: |
| reverse_mapping=lambda x: x |
| for file in donefiles: |
| if file not in args.targetfiles: |
| continue |
| else: |
| print(file) |
| saved_dict=torch.load(os.path.join(args.save_dir,args.model,'model',file)) |
| if saved_dict['args'].dataset in ["mnist","mnist_sparse"]: |
| x_out=90 |
| elif saved_dict['args'].dataset in ["building","mbuilding"]: |
| x_out=100 |
| elif saved_dict['args'].dataset in ["sbuilding","smnist"]: |
| x_out=10 |
| elif saved_dict['args'].dataset in ["dbp"]: |
| x_out=2 |
| if saved_dict['args'].model=="our": |
| model=Smodel(h_channel=saved_dict['args'].h_ch,input_featuresize=saved_dict['args'].h_ch,\ |
| localdepth=saved_dict['args'].localdepth,num_interactions=saved_dict['args'].num_interactions,finaldepth=saved_dict['args'].finaldepth,share=saved_dict['args'].share,batchnorm=saved_dict['args'].batchnorm) |
| mlpmodel=MLP(in_channels=saved_dict['args'].h_ch*saved_dict['args'].num_interactions, hidden_channels=saved_dict['args'].h_ch,out_channels=x_out, num_layers=saved_dict['args'].classifier_depth) |
| elif saved_dict['args'].model=="HGT": |
| model=model_new.HGT(hidden_channels=saved_dict['args'].h_ch, out_channels=saved_dict['args'].h_ch, num_heads=2, num_layers=saved_dict['args'].num_interactions) |
| mlpmodel=MLP(in_channels=saved_dict['args'].h_ch, hidden_channels=saved_dict['args'].h_ch,out_channels=x_out, num_layers=saved_dict['args'].classifier_depth,dropout=saved_dict['args'].dropout) |
| elif saved_dict['args'].model=="HAN": |
| model=model_new.HAN(hidden_channels=saved_dict['args'].h_ch, out_channels=saved_dict['args'].h_ch, num_heads=2, num_layers=saved_dict['args'].num_interactions) |
| mlpmodel=MLP(in_channels=saved_dict['args'].h_ch, hidden_channels=saved_dict['args'].h_ch,out_channels=x_out, num_layers=saved_dict['args'].classifier_depth,dropout=saved_dict['args'].dropout) |
| model.to(device), mlpmodel.to(device) |
| try: |
| model.load_state_dict(saved_dict['model'],strict=True) |
| mlpmodel.load_state_dict(saved_dict['mlpmodel'],strict=True) |
| except OSError: |
| print('loadfail: ',file) |
| pass |
| print(saved_dict['args']) |
|
|
| writer = SummaryWriter(os.path.join(tensorboard_dir,file+"_embedding")) |
| test_loss, yhat_test, ytrue_test, yhatlogit_test = test(saved_dict['args'],test_Loader,model,mlpmodel,writer,reverse_mapping) |
| |
| pred_dir=os.path.join(tensorboard_dir,file+"_test_record") |
| to_save_dict={'labels':ytrue_test,'yhat':yhat_test,'yhat_logit':yhatlogit_test} |
| torch.save(to_save_dict, pred_dir) |
| |
| test_acc=util.calculate(yhat_test,ytrue_test,yhatlogit_test) |
| util.print_1(0,'Test', {"loss":test_loss,"acc":test_acc},color=blue) |
|
|
| |
| if __name__ == '__main__': |
| Seed = 0 |
| test_ratio=0.2 |
| print("data splitting Random Seed: ", Seed) |
| if args.dataset in ["mnist"]: |
| args.data_dir='data/multi_mnist_with_index.pkl' |
| train_ds,val_ds,test_ds=dataset.get_mnist_dataset(args.data_dir,Seed,test_ratio=test_ratio) |
| elif args.dataset in ["mnist_sparse"]: |
| args.data_dir='data/multi_mnist_sparse.pkl' |
| train_ds,val_ds,test_ds=dataset.get_mnist_dataset(args.data_dir,Seed,test_ratio=test_ratio) |
| elif args.dataset in ["building"]: |
| args.data_dir='data/building_with_index.pkl' |
| train_ds,val_ds,test_ds=dataset.get_mbuilding_dataset(args.data_dir,Seed,test_ratio=test_ratio) |
| elif args.dataset in ["mbuilding"]: |
| args.data_dir='data/mp_building.pkl' |
| train_ds,val_ds,test_ds=dataset.get_building_dataset(args.data_dir,Seed,test_ratio=test_ratio) |
| elif args.dataset in ["sbuilding"]: |
| args.data_dir='data/single_building.pkl' |
| train_ds,val_ds,test_ds=dataset.get_sbuilding_dataset(args.data_dir,Seed,test_ratio=test_ratio) |
| elif args.dataset in ["smnist"]: |
| args.data_dir='data/single_mnist.pkl' |
| train_ds,val_ds,test_ds=dataset.get_smnist_dataset(args.data_dir,Seed,test_ratio=test_ratio) |
| elif args.dataset in ['dbp']: |
| args.data_dir='data/triple_building_600.pkl' |
| train_ds,val_ds,test_ds=dataset.get_dbp_dataset(args.data_dir,Seed,test_ratio=test_ratio) |
|
|
| if args.extent_norm=="T": |
| train_ds= dataset.affine_transform_to_range(train_ds,target_range=(-1, 1)) |
| val_ds= dataset.affine_transform_to_range(val_ds,target_range=(-1, 1)) |
| test_ds= dataset.affine_transform_to_range(test_ds,target_range=(-1, 1)) |
| train_loader = torch_geometric.loader.DataLoader(train_ds,batch_size=args.train_batch, shuffle=False,pin_memory=True) |
| val_loader = torch_geometric.loader.DataLoader(val_ds, batch_size=args.test_batch, shuffle=False, pin_memory=True) |
| test_loader = torch_geometric.loader.DataLoader(test_ds,batch_size=args.test_batch, shuffle=False,pin_memory=True) |
|
|
| Seed=random.randint(1, 10000) |
| print("Random Seed: ", Seed) |
| random.seed(Seed) |
| torch.manual_seed(Seed) |
| np.random.seed(Seed) |
| main(args,train_loader,val_loader,test_loader) |
