os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

if torch.cuda.is_available() and use_gpu:

print('cuda available with GPU:',torch.cuda.get_device_name(0))

device = torch.device("cuda")

else:

print('cuda not available')

device = torch.device("cpu")

model = gnn_model(MODEL_NAME, net_params)

total_param = 0

print("MODEL DETAILS:\n")

#print(model)

for param in model.parameters():

#print(param.data.size())

total_param += np.prod(list(param.data.size()))

print('MODEL/Total parameters:', MODEL_NAME, total_param)

start0 = time.time()

per_epoch_time = []

DATASET_NAME = dataset.name

if MODEL_NAME in ['GCN', 'GAT', 'SGC']:

if net_params['self_loop']:

print("[!] Adding graph self-loops (central node trick).")

dataset._add_self_loops()

root_log_dir, root_ckpt_dir, write_file_name, write_config_file = dirs

device = net_params['device']

print('The seed is ', params['seed'])

random.seed(params['seed'])

torch.manual_seed(params['seed'])

if device == 'cuda':

torch.cuda.manual_seed(params['seed'])

np.random.seed(params['seed'])

num_nodes = dataset.train_mask.size(0)

index = torch.tensor(np.random.permutation(num_nodes))

print('index:', index)

train_index = index[:int(num_nodes*0.6)]

val_index = index[int(num_nodes*0.6):int(num_nodes*0.8)]

test_index = index[int(num_nodes*0.8):]

train_mask = index_to_mask(train_index, size=num_nodes)

val_mask = index_to_mask(val_index, size=num_nodes)

test_mask = index_to_mask(test_index, size=num_nodes)

train_mask = train_mask.to(device)

val_mask = val_mask.to(device)

test_mask = test_mask.to(device)

labels = dataset.labels.to(device)

# Write network and optimization hyper-parameters in folder config/

with open(write_config_file + '.txt', 'w') as f:

f.write("""Dataset: {},\nModel: {}\n\nparams={}\n\nnet_params={}\n\n\nTotal Parameters: {}\n\n""" .format(DATASET_NAME, MODEL_NAME, params, net_params, net_params['total_param']))

log_dir = os.path.join(root_log_dir, "RUN_" + str(0))

writer = SummaryWriter(log_dir=log_dir)

print("Training Nodes: ", train_mask.int().sum().item())

print("Validation Nodes: ", val_mask.int().sum().item())

print("Test Nodes: ", test_mask.int().sum().item())

print("Number of Classes: ", net_params['n_classes'])

model = gnn_model(MODEL_NAME, net_params)

model = model.to(device)

optimizer = optim.Adam(model.parameters(), lr=params['init_lr'], weight_decay=params['weight_decay'])

# scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',

# factor=params['lr_reduce_factor'],

# patience=params['lr_schedule_patience'],

# verbose=True)

epoch_train_losses, epoch_val_losses = [], []

epoch_train_accs, epoch_val_accs = [], []

graph = dataset.graph

nfeat = graph.ndata['feat'].to(device)

efeat = graph.edata['feat'].to(device)

norm_n = dataset.norm_n.to(device)

norm_e = dataset.norm_e.to(device)

# At any point you can hit Ctrl + C to break out of training early.

try:

with tqdm(range(params['epochs'])) as t:

best_val_acc = 0

for epoch in t:

t.set_description('Epoch %d' % epoch)

start = time.time()

epoch_train_loss, epoch_train_acc, optimizer = train_epoch(model, optimizer, device, graph, nfeat, efeat, norm_n, norm_e, train_mask, labels, epoch)

epoch_val_loss, epoch_val_acc = evaluate_network(model, graph, nfeat, efeat, norm_n, norm_e, val_mask, labels, epoch)

epoch_train_losses.append(epoch_train_loss)

epoch_val_losses.append(epoch_val_loss)

epoch_train_accs.append(epoch_train_acc)

epoch_val_accs.append(epoch_val_acc)

writer.add_scalar('train/_loss', epoch_train_loss, epoch)

writer.add_scalar('val/_loss', epoch_val_loss, epoch)

writer.add_scalar('train/_acc', epoch_train_acc, epoch)

writer.add_scalar('val/_acc', epoch_val_acc, epoch)

writer.add_scalar('learning_rate', optimizer.param_groups[0]['lr'], epoch)

_, epoch_test_acc = evaluate_network(model, graph, nfeat, efeat, norm_n, norm_e, test_mask, labels, epoch)

t.set_postfix(time=time.time()-start, lr=optimizer.param_groups[0]['lr'],

train_loss=epoch_train_loss, val_loss=epoch_val_loss,

train_acc=epoch_train_acc, val_acc=epoch_val_acc,

test_acc=epoch_test_acc)

per_epoch_time.append(time.time()-start)

# Saving checkpoint

ckpt_dir = os.path.join(root_ckpt_dir, "RUN_")

if not os.path.exists(ckpt_dir):

os.makedirs(ckpt_dir)

torch.save(model.state_dict(), '{}.pkl'.format(ckpt_dir + "/epoch_" + str(epoch)))

if best_val_acc < epoch_val_acc:

best_val_acc = epoch_val_acc

torch.save(model.state_dict(), '{}.pkl'.format(ckpt_dir + "/best"))

files = glob.glob(ckpt_dir + '/*.pkl')

for file in files:

if file[-8:]=='best.pkl':

continue

else:

epoch_nb = file.split('_')[-1]

epoch_nb = int(epoch_nb.split('.')[0])

if epoch_nb < epoch - 1:

os.remove(file)

#scheduler.step(epoch_val_loss)

if optimizer.param_groups[0]['lr'] < params['min_lr']:

optimizer.param_groups[0]['lr'] = params['min_lr']

#print("\n!! LR SMALLER OR EQUAL TO MIN LR THRESHOLD.")

#break

# Stop training after params['max_time'] hours

if time.time()-start0 > params['max_time']*3600:

print('-' * 89)

print("Max_time for training elapsed {:.2f} hours, so stopping".format(params['max_time']))

break

except KeyboardInterrupt:

print('-' * 89)

print('Exiting from training early because of KeyboardInterrupt')

model.load_state_dict(torch.load('{}.pkl'.format(ckpt_dir + "/best")))

_, test_acc = evaluate_network(model, graph, nfeat, efeat, norm_n, norm_e, test_mask, labels, epoch)

_, val_acc = evaluate_network(model, graph, nfeat, efeat, norm_n, norm_e, val_mask, labels, epoch)

_, train_acc = evaluate_network(model, graph, nfeat, efeat, norm_n, norm_e, train_mask, labels, epoch)

print("Test Accuracy: {:.4f}".format(test_acc))

print("Train Accuracy: {:.4f}".format(train_acc))

print("TOTAL TIME TAKEN: {:.4f}s".format(time.time()-start0))

print("AVG TIME PER EPOCH: {:.4f}s".format(np.mean(per_epoch_time)))

writer.close()

"""

Write the results in out_dir/results folder

"""

with open(write_file_name + '.txt', 'w') as f:

f.write("""Dataset: {},\nModel: {}\n\nparams={}\n\nnet_params={}\n\n{}\n\nTotal Parameters: {}\n\n

FINAL RESULTS\nTEST ACCURACY: {:.4f}\nTRAIN ACCURACY: {:.4f}\n\n

Total Time Taken: {:.4f} hrs\nAverage Time Per Epoch: {:.4f} s\n\n\n"""\

.format(DATASET_NAME, MODEL_NAME, params, net_params, model, net_params['total_param'],

test_acc, train_acc, (time.time()-start0)/3600, np.mean(per_epoch_time)))

# send results to gmail

try:

from gmail import send

subject = 'Result for Dataset: {}, Model: {}'.format(DATASET_NAME, MODEL_NAME)

body = """Dataset: {},\nModel: {}\n\nparams={}\n\nnet_params={}\n\n{}\n\nTotal Parameters: {}\n\n

FINAL RESULTS\nTEST ACCURACY: {:.4f}\nTRAIN ACCURACY: {:.4f}\n\n

Total Time Taken: {:.4f} hrs\nAverage Time Per Epoch: {:.4f} s\n\n\n"""\

.format(DATASET_NAME, MODEL_NAME, params, net_params, model, net_params['total_param'],

test_acc, train_acc, (time.time()-start0)/3600, np.mean(per_epoch_time))

send(subject, body)

except:

pass

"""

USER CONTROLS

"""

parser = argparse.ArgumentParser()

parser.add_argument('--config', help="Please give a config.json file with training/model/data/param details")

parser.add_argument('--gpu_id', help="Please give a value for gpu id")

parser.add_argument('--model', help="Please give a value for model name")

parser.add_argument('--dataset', help="Please give a value for dataset name")

parser.add_argument('--builtin', help="Please give a value for builtin")

parser.add_argument('--out_dir', help="Please give a value for out_dir")

parser.add_argument('--seed', help="Please give a value for seed")

parser.add_argument('--epochs', help="Please give a value for epochs")

parser.add_argument('--runs', help="Please give a value for runs")

parser.add_argument('--batch_size', help="Please give a value for batch_size")

parser.add_argument('--init_lr', help="Please give a value for init_lr")

parser.add_argument('--lr_reduce_factor', help="Please give a value for lr_reduce_factor")

parser.add_argument('--lr_schedule_patience', help="Please give a value for lr_schedule_patience")

parser.add_argument('--min_lr', help="Please give a value for min_lr")

parser.add_argument('--weight_decay', help="Please give a value for weight_decay")

parser.add_argument('--print_epoch_interval', help="Please give a value for print_epoch_interval")

parser.add_argument('--L', help="Please give a value for L")

parser.add_argument('--K', help="Please give a value for K")

parser.add_argument('--num_low', help="Please give a value for number of low filters")

parser.add_argument('--num_high', help="Please give a value for number of high filters")

parser.add_argument('--num_mid', help="Please give a value for number of middle filters")

parser.add_argument('--opt', help="Please give a value for AutoGCN option")

parser.add_argument('--gate', help="Please give a value for AutoGCN gate option")

parser.add_argument('--hidden_dim', help="Please give a value for hidden_dim")

parser.add_argument('--out_dim', help="Please give a value for out_dim")

parser.add_argument('--residual', help="Please give a value for residual")

parser.add_argument('--edge_feat', help="Please give a value for edge_feat")

parser.add_argument('--readout', help="Please give a value for readout")

parser.add_argument('--kernel', help="Please give a value for kernel")

parser.add_argument('--n_heads', help="Please give a value for n_heads")

parser.add_argument('--gated', help="Please give a value for gated")

parser.add_argument('--in_feat_dropout', help="Please give a value for in_feat_dropout")

parser.add_argument('--dropout', help="Please give a value for dropout")

parser.add_argument('--graph_norm', help="Please give a value for graph_norm")

parser.add_argument('--batch_norm', help="Please give a value for batch_norm")

parser.add_argument('--sage_aggregator', help="Please give a value for sage_aggregator")

parser.add_argument('--data_mode', help="Please give a value for data_mode")

parser.add_argument('--num_pool', help="Please give a value for num_pool")

parser.add_argument('--gnn_per_block', help="Please give a value for gnn_per_block")

parser.add_argument('--embedding_dim', help="Please give a value for embedding_dim")

parser.add_argument('--pool_ratio', help="Please give a value for pool_ratio")

parser.add_argument('--linkpred', help="Please give a value for linkpred")

parser.add_argument('--cat', help="Please give a value for cat")

parser.add_argument('--self_loop', help="Please give a value for self_loop")

parser.add_argument('--max_time', help="Please give a value for max_time")

args = parser.parse_args()

with open(args.config) as f:

config = json.load(f)

# device

if args.gpu_id is not None:

config['gpu']['id'] = int(args.gpu_id)

config['gpu']['use'] = True

device = gpu_setup(config['gpu']['use'], config['gpu']['id'])

# model, dataset, out_dir

if args.model is not None:

MODEL_NAME = args.model

else:

MODEL_NAME = config['model']

if args.dataset is not None:

DATASET_NAME = args.dataset

else:

DATASET_NAME = config['dataset']

dataset = LoadData(DATASET_NAME)

if args.out_dir is not None:

out_dir = args.out_dir

else:

out_dir = config['out_dir']

# parameters

params = config['params']

if args.seed is not None:

params['seed'] = int(args.seed)

if args.epochs is not None:

params['epochs'] = int(args.epochs)

if args.runs is not None:

params['runs'] = int(args.runs)

if args.batch_size is not None:

params['batch_size'] = int(args.batch_size)

if args.init_lr is not None:

params['init_lr'] = float(args.init_lr)

if args.lr_reduce_factor is not None:

params['lr_reduce_factor'] = float(args.lr_reduce_factor)

if args.lr_schedule_patience is not None:

params['lr_schedule_patience'] = int(args.lr_schedule_patience)

if args.min_lr is not None:

params['min_lr'] = float(args.min_lr)

if args.weight_decay is not None:

params['weight_decay'] = float(args.weight_decay)

if args.print_epoch_interval is not None:

params['print_epoch_interval'] = int(args.print_epoch_interval)

if args.max_time is not None:

params['max_time'] = float(args.max_time)

# network parameters

net_params = config['net_params']

net_params['device'] = device

net_params['gpu_id'] = config['gpu']['id']

net_params['batch_size'] = params['batch_size']

if args.L is not None:

net_params['L'] = int(args.L)

if args.K is not None:

net_params['K'] = int(args.K)

if args.num_low is not None:

net_params['num_low'] = int(args.num_low)

if args.num_high is not None:

net_params['num_high'] = int(args.num_high)

if args.num_mid is not None:

net_params['num_mid'] = int(args.num_mid)

if args.opt is not None:

net_params['opt'] = args.opt

if args.gate is not None:

net_params['gate'] = True if args.gate=='True' else False

if args.hidden_dim is not None:

net_params['hidden_dim'] = int(args.hidden_dim)

if args.out_dim is not None:

net_params['out_dim'] = int(args.out_dim)

if args.residual is not None:

net_params['residual'] = True if args.residual=='True' else False

if args.edge_feat is not None:

net_params['edge_feat'] = True if args.edge_feat=='True' else False

if args.readout is not None:

net_params['readout'] = args.readout

if args.kernel is not None:

net_params['kernel'] = int(args.kernel)

if args.n_heads is not None:

net_params['n_heads'] = int(args.n_heads)

if args.gated is not None:

net_params['gated'] = True if args.gated=='True' else False

if args.in_feat_dropout is not None:

net_params['in_feat_dropout'] = float(args.in_feat_dropout)

if args.dropout is not None:

net_params['dropout'] = float(args.dropout)

if args.graph_norm is not None:

net_params['graph_norm'] = True if args.graph_norm=='True' else False

if args.batch_norm is not None:

net_params['batch_norm'] = True if args.batch_norm=='True' else False

if args.sage_aggregator is not None:

net_params['sage_aggregator'] = args.sage_aggregator

if args.data_mode is not None:

net_params['data_mode'] = args.data_mode

if args.num_pool is not None:

net_params['num_pool'] = int(args.num_pool)

if args.gnn_per_block is not None:

net_params['gnn_per_block'] = int(args.gnn_per_block)

if args.embedding_dim is not None:

net_params['embedding_dim'] = int(args.embedding_dim)

if args.pool_ratio is not None:

net_params['pool_ratio'] = float(args.pool_ratio)

if args.linkpred is not None:

net_params['linkpred'] = True if args.linkpred=='True' else False

if args.cat is not None:

net_params['cat'] = True if args.cat=='True' else False

if args.self_loop is not None:

net_params['self_loop'] = True if args.self_loop=='True' else False

if args.builtin is not None:

net_params['builtin'] = True if args.builtin == 'True' else False

# CitationGraph

net_params['in_dim'] = dataset.num_dims # node_dim (feat is an integer)

net_params['n_classes'] = dataset.num_classes

if MODEL_NAME in ['MLP', 'MLP_GATED']:

builtin = ''

else:

builtin = 'DGL' if net_params['builtin'] else 'Custom'

root_log_dir = out_dir + 'logs/' + MODEL_NAME + "_" + DATASET_NAME + "_GPU" + str(config['gpu']['id']) + "_" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y') + builtin

root_ckpt_dir = out_dir + 'checkpoints/' + MODEL_NAME + "_" + DATASET_NAME + "_GPU" + str(config['gpu']['id']) + "_" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y') + builtin

write_file_name = out_dir + 'results/result_' + MODEL_NAME + "_" + DATASET_NAME + "_GPU" + str(config['gpu']['id']) + "_" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y') + builtin

write_config_file = out_dir + 'configs/config_' + MODEL_NAME + "_" + DATASET_NAME + "_GPU" + str(config['gpu']['id']) + "_" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y') + builtin

dirs = root_log_dir, root_ckpt_dir, write_file_name, write_config_file

if not os.path.exists(out_dir + 'results'):

os.makedirs(out_dir + 'results')

if not os.path.exists(out_dir + 'configs'):

os.makedirs(out_dir + 'configs')

val_acc = []

test_acc = []

np.random.seed(121)

seeds = [np.random.randint(10000) for i in range(params['runs'])]

print('seeds', seeds)

for i in range(params['runs']):

params['seed'] = seeds[i]

net_params['total_param'] = view_model_param(MODEL_NAME, net_params)

val_ac, test_ac = train_val_pipeline(MODEL_NAME, dataset, params, net_params, dirs)

val_acc.append(val_ac)

test_acc.append(test_ac)

val_acc = np.array(val_acc)

test_acc = np.array(test_acc)

print('\n')

print('-'*89)