def eval_with_partition(args):
model_load_path = args.model_load_path
print("Starting evaluating model stored at", model_load_path)
device = torch.device("cuda")
dataset = PygNodePropPredDataset(name=args.dataset, root=args.data_folder)
graph = dataset[0]
adj = SparseTensor(row=graph.edge_index[0], col=graph.edge_index[1])
if args.self_loop:
adj = adj.set_diag()
graph.edge_index = add_self_loops(edge_index=graph.edge_index,
num_nodes=graph.num_nodes)[0]
split_idx = dataset.get_idx_split()
evaluator = Evaluator(args.dataset)
args.in_channels = graph.x.size(-1)
args.num_tasks = dataset.num_classes
# print('%s' % args)
model = DeeperGCN(args).to(device)
ckpt = torch.load(model_load_path)
model.load_state_dict(ckpt['model_state_dict'])
res = test_with_partition(model,
graph,
adj,
split_idx,
num_clusters=args.eval_cluster_number,
partition_method=args.partition_method,
evaluator=evaluator,
device=device)
print(res)
return res
def get_adj(row, col, N, asymm_norm=False, set_diag=True, remove_diag=False):
adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N))
if set_diag:
print('... setting diagonal entries')
adj = adj.set_diag()
elif remove_diag:
print('... removing diagonal entries')
adj = adj.remove_diag()
else:
print('... keeping diag elements as they are')
if not asymm_norm:
print('... performing symmetric normalization')
deg = adj.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
else:
print('... performing asymmetric normalization')
deg = adj.sum(dim=1).to(torch.float)
deg_inv = deg.pow(-1.0)
deg_inv[deg_inv == float('inf')] = 0
adj = deg_inv.view(-1, 1) * adj
adj = adj.to_scipy(layout='csr')
return adj
def preprocess(data,
preprocess="diffusion",
num_propagations=10,
p=None,
alpha=None,
use_cache=True,
post_fix=""):
if use_cache:
try:
x = torch.load(f'embeddings/{preprocess}{post_fix}.pt')
print('Using cache')
return x
except:
print(
f'embeddings/{preprocess}{post_fix}.pt not found or not enough iterations! Regenerating it now'
)
# Creates a new file
with open(f'embeddings/{preprocess}{post_fix}.pt', 'w') as fp:
pass
if preprocess == "community":
return community(data, post_fix)
if preprocess == "spectral":
return spectral(data, post_fix)
print('Computing adj...')
N = data.num_nodes
data.edge_index = to_undirected(data.edge_index, data.num_nodes)
row, col = data.edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N))
adj = adj.set_diag()
deg = adj.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
adj = adj.to_scipy(layout='csr')
sgc_dict = {}
print(f'Start {preprocess} processing')
if preprocess == "sgc":
result = sgc(data.x.numpy(), adj, num_propagations)
# if preprocess == "lp":
# result = lp(adj, data.y.data, num_propagations, p = p, alpha = alpha, preprocess = preprocess)
if preprocess == "diffusion":
result = diffusion(data.x.numpy(),
adj,
num_propagations,
p=p,
alpha=alpha)
torch.save(result, f'embeddings/{preprocess}{post_fix}.pt')
return result
def process_adj(data):
N = data.num_nodes
data.edge_index = to_undirected(data.edge_index, data.num_nodes)
row, col = data.edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N))
adj = adj.set_diag()
deg = adj.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
return adj
def test(model, predictor, data, split_edge, evaluator, batch_size, device):
predictor.eval()
print('Evaluating full-batch GNN on CPU...')
weights = [(conv.weight.cpu().detach().numpy(),
conv.bias.cpu().detach().numpy()) for conv in model.convs]
model = GCNInference(weights)
x = data.x.numpy()
adj = SparseTensor(row=data.edge_index[0], col=data.edge_index[1])
adj = adj.set_diag()
deg = adj.sum(dim=1)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
adj = adj.to_scipy(layout='csr')
h = torch.from_numpy(model(x, adj)).to(device)
def test_split(split):
source = split_edge[split]['source_node'].to(device)
target = split_edge[split]['target_node'].to(device)
target_neg = split_edge[split]['target_node_neg'].to(device)
pos_preds = []
for perm in DataLoader(range(source.size(0)), batch_size):
src, dst = source[perm], target[perm]
pos_preds += [predictor(h[src], h[dst]).squeeze().cpu()]
pos_pred = torch.cat(pos_preds, dim=0)
neg_preds = []
source = source.view(-1, 1).repeat(1, 1000).view(-1)
target_neg = target_neg.view(-1)
for perm in DataLoader(range(source.size(0)), batch_size):
src, dst_neg = source[perm], target_neg[perm]
neg_preds += [predictor(h[src], h[dst_neg]).squeeze().cpu()]
neg_pred = torch.cat(neg_preds, dim=0).view(-1, 1000)
return evaluator.eval({
'y_pred_pos': pos_pred,
'y_pred_neg': neg_pred,
})['mrr_list'].mean().item()
train_mrr = test_split('eval_train')
valid_mrr = test_split('valid')
test_mrr = test_split('test')
return train_mrr, valid_mrr, test_mrr
def get_adj(row, col, N, asymm_norm=False, set_diag=True, remove_diag=False):
adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N))
if set_diag:
adj = adj.set_diag()
elif remove_diag:
adj = adj.remove_diag()
if not asymm_norm:
deg = adj.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
else:
deg = adj.sum(dim=1).to(torch.float)
deg_inv = deg.pow(-1.0)
deg_inv[deg_inv == float('inf')] = 0
adj = deg_inv.view(-1, 1) * adj
return adj
def main():
parser = argparse.ArgumentParser(description='OGBL-PPA (Full-Batch)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--use_node_embedding', action='store_true')
parser.add_argument('--use_sage', action='store_true')
parser.add_argument('--num_layers', type=int, default=3)
parser.add_argument('--hidden_channels', type=int, default=256)
parser.add_argument('--dropout', type=float, default=0.0)
parser.add_argument('--batch_size', type=int, default=64 * 1024)
parser.add_argument('--lr', type=float, default=0.01)
parser.add_argument('--epochs', type=int, default=20)
parser.add_argument('--eval_steps', type=int, default=1)
parser.add_argument('--runs', type=int, default=10)
args = parser.parse_args()
print(args)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
dataset = PygLinkPropPredDataset(name='ogbl-ppa')
data = dataset[0]
splitted_edge = dataset.get_edge_split()
if args.use_node_embedding:
x = data.x.to(torch.float)
x = torch.cat([x, torch.load('embedding.pt')], dim=-1)
x = x.to(device)
else:
x = data.x.to(torch.float).to(device)
edge_index = data.edge_index.to(device)
adj = SparseTensor(row=edge_index[0], col=edge_index[1])
if args.use_sage:
model = SAGE(x.size(-1), args.hidden_channels, args.hidden_channels,
args.num_layers, args.dropout).to(device)
else:
model = GCN(x.size(-1), args.hidden_channels, args.hidden_channels,
args.num_layers, args.dropout).to(device)
# Pre-compute GCN normalization.
adj = adj.set_diag()
deg = adj.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
predictor = LinkPredictor(args.hidden_channels, args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
evaluator = Evaluator(name='ogbl-ppa')
loggers = {
'Hits@10': Logger(args.runs, args),
'Hits@50': Logger(args.runs, args),
'Hits@100': Logger(args.runs, args),
}
for run in range(args.runs):
model.reset_parameters()
predictor.reset_parameters()
optimizer = torch.optim.Adam(
list(model.parameters()) + list(predictor.parameters()),
lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(model, predictor, x, adj, splitted_edge, optimizer,
args.batch_size)
if epoch % args.eval_steps == 0:
results = test(model, predictor, x, adj, splitted_edge,
evaluator, args.batch_size)
for key, result in results.items():
loggers[key].add_result(run, result)
if epoch % args.log_steps == 0:
for key, result in results.items():
train_hits, valid_hits, test_hits = result
print(key)
print(f'Run: {run + 1:02d}, '
f'Epoch: {epoch:02d}, '
f'Loss: {loss:.4f}, '
f'Train: {100 * train_hits:.2f}%, '
f'Valid: {100 * valid_hits:.2f}%, '
f'Test: {100 * test_hits:.2f}%')
for key in loggers.keys():
print(key)
loggers[key].print_statistics(run)
for key in loggers.keys():
print(key)
loggers[key].print_statistics()
def main():
parser = argparse.ArgumentParser(description='OGBN-papers100M (MLP)')
parser.add_argument('--data_root_dir', type=str, default='../../dataset')
parser.add_argument('--num_propagations', type=int, default=3)
parser.add_argument('--dropedge_rate', type=float, default=0.4)
parser.add_argument('--node_emb_path', type=str, default=None)
parser.add_argument('--output_path', type=str, required=True)
args = parser.parse_args()
# SGC pre-processing ######################################################
dataset = PygNodePropPredDataset(name='ogbn-papers100M',
root=args.data_root_dir)
split_idx = dataset.get_idx_split()
data = dataset[0]
x = None
if args.node_emb_path:
x = np.load(args.node_emb_path)
else:
x = data.x.numpy()
N = data.num_nodes
print('Making the graph undirected.')
### Randomly drop some edges to save computation
data.edge_index, _ = dropout_adj(data.edge_index,
p=args.dropedge_rate,
num_nodes=data.num_nodes)
data.edge_index = to_undirected(data.edge_index, data.num_nodes)
print(data)
row, col = data.edge_index
print('Computing adj...')
adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N))
adj = adj.set_diag()
deg = adj.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
adj = adj.to_scipy(layout='csr')
train_idx, valid_idx, test_idx = split_idx['train'], split_idx[
'valid'], split_idx['test']
all_idx = torch.cat([train_idx, valid_idx, test_idx])
mapped_train_idx = torch.arange(len(train_idx))
mapped_valid_idx = torch.arange(len(train_idx),
len(train_idx) + len(valid_idx))
mapped_test_idx = torch.arange(
len(train_idx) + len(valid_idx),
len(train_idx) + len(valid_idx) + len(test_idx))
sgc_dict = {}
sgc_dict['label'] = data.y.data[all_idx].to(torch.long)
sgc_dict['split_idx'] = {
'train': mapped_train_idx,
'valid': mapped_valid_idx,
'test': mapped_test_idx
}
print('Start SGC processing')
for _ in tqdm(range(args.num_propagations)):
x = adj @ x
sgc_dict['sgc_embedding'] = torch.from_numpy(x[all_idx]).to(torch.float)
torch.save(sgc_dict, args.output_path)
args = parser.parse_args()
print(args)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
dataset = PygNodePropPredDataset(name='ogbn-arxiv')
split_idx = dataset.get_idx_split()
data = dataset[0]
edge_index = data.edge_index.to(device)
edge_index = to_undirected(edge_index, data.num_nodes)
adj_0 = SparseTensor(row=edge_index[0], col=edge_index[1])
# Pre-compute GCN normalization.
adj = adj_0.set_diag()
deg = adj.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
class l_GCN(torch.nn.Module):
def __init__(self, in_channels, out_channels):
super(l_GCN, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.weight = Parameter(torch.Tensor(in_channels, out_channels))
self.bias = Parameter(torch.Tensor(out_channels))
def main():
parser = argparse.ArgumentParser(description='OGBL-Citation (GNN)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--use_sage', action='store_true')
parser.add_argument('--num_layers', type=int, default=3)
parser.add_argument('--hidden_channels', type=int, default=256)
parser.add_argument('--dropout', type=float, default=0)
parser.add_argument('--batch_size', type=int, default=64 * 1024)
parser.add_argument('--lr', type=float, default=0.0005)
parser.add_argument('--epochs', type=int, default=50)
parser.add_argument('--eval_steps', type=int, default=1)
parser.add_argument('--runs', type=int, default=10)
args = parser.parse_args()
print(args)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
dataset = PygLinkPropPredDataset(name='ogbl-citation')
split_edge = dataset.get_edge_split()
data = dataset[0]
# We randomly pick some training samples that we want to evaluate on:
torch.manual_seed(12345)
idx = torch.randperm(split_edge['train']['source_node'].numel())[:86596]
split_edge['eval_train'] = {
'source_node': split_edge['train']['source_node'][idx],
'target_node': split_edge['train']['target_node'][idx],
'target_node_neg': split_edge['valid']['target_node_neg'],
}
x = data.x.to(device)
edge_index = data.edge_index.to(device)
edge_index = to_undirected(edge_index, data.num_nodes)
adj = SparseTensor(row=edge_index[0], col=edge_index[1])
if args.use_sage:
model = SAGE(x.size(-1), args.hidden_channels, args.hidden_channels,
args.num_layers, args.dropout).to(device)
else:
model = GCN(x.size(-1), args.hidden_channels, args.hidden_channels,
args.num_layers, args.dropout).to(device)
# Pre-compute GCN normalization.
adj = adj.set_value(None)
adj = adj.set_diag()
deg = adj.sum(dim=1)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
predictor = LinkPredictor(args.hidden_channels, args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
evaluator = Evaluator(name='ogbl-citation')
logger = Logger(args.runs, args)
for run in range(args.runs):
model.reset_parameters()
predictor.reset_parameters()
optimizer = torch.optim.Adam(list(model.parameters()) +
list(predictor.parameters()),
lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(model, predictor, x, adj, split_edge, optimizer,
args.batch_size)
print(f'Run: {run + 1:02d}, Epoch: {epoch:02d}, Loss: {loss:.4f}')
if epoch % args.eval_steps == 0:
result = test(model, predictor, x, adj, split_edge, evaluator,
args.batch_size)
logger.add_result(run, result)
if epoch % args.log_steps == 0:
train_mrr, valid_mrr, test_mrr = result
print(f'Run: {run + 1:02d}, '
f'Epoch: {epoch:02d}, '
f'Loss: {loss:.4f}, '
f'Train: {train_mrr:.4f}, '
f'Valid: {valid_mrr:.4f}, '
f'Test: {test_mrr:.4f}')
logger.print_statistics(run)
logger.print_statistics()
def main():
parser = argparse.ArgumentParser(description='OGBN-Arxiv (Full-Batch)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--use_sage', action='store_true')
parser.add_argument('--num_layers', type=int, default=3)
parser.add_argument('--hidden_channels', type=int, default=256)
parser.add_argument('--dropout', type=float, default=0.5)
parser.add_argument('--lr', type=float, default=0.01)
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--runs', type=int, default=5)
args = parser.parse_args()
print(args)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
dataset = PygNodePropPredDataset(name='ogbn-arxiv')
split_idx = dataset.get_idx_split()
data = dataset[0]
x = data.x.to(device)
y_true = data.y.to(device)
train_idx = split_idx['train'].to(device)
edge_index = data.edge_index.to(device)
edge_index = to_undirected(edge_index, data.num_nodes)
adj = SparseTensor(row=edge_index[0], col=edge_index[1])
if args.use_sage:
model = SAGE(data.x.size(-1), args.hidden_channels, dataset.num_classes,
args.num_layers, args.dropout).to(device)
else:
model = GCN(data.x.size(-1), args.hidden_channels, dataset.num_classes, args.num_layers,
args.dropout).to(device)
# Pre-compute GCN normalization.
adj = adj.set_diag()
deg = adj.sum(dim=1).to(torch.float)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
evaluator = Evaluator(name='ogbn-arxiv')
logger = Logger(args.runs, args)
for run in range(args.runs):
model.reset_parameters()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(model, x, adj, y_true, train_idx, optimizer)
result = test(model, x, adj, y_true, split_idx, evaluator)
logger.add_result(run, result)
if epoch % args.log_steps == 0:
train_acc, valid_acc, test_acc = result
print(f'Run: {run + 1:02d}, '
f'Epoch: {epoch:02d}, '
f'Loss: {loss:.4f}, '
f'Train: {100 * train_acc:.2f}%, '
f'Valid: {100 * valid_acc:.2f}% '
f'Test: {100 * test_acc:.2f}%')
logger.print_statistics(run)
logger.print_statistics()
def main():
args = ArgsInit().save_exp()
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
dataset = PygNodePropPredDataset(name=args.dataset, root=args.data_folder)
graph = dataset[0]
adj = SparseTensor(row=graph.edge_index[0], col=graph.edge_index[1])
if args.self_loop:
adj = adj.set_diag()
graph.edge_index = add_self_loops(edge_index=graph.edge_index,
num_nodes=graph.num_nodes)[0]
split_idx = dataset.get_idx_split()
train_idx = split_idx["train"].tolist()
evaluator = Evaluator(args.dataset)
sub_dir = 'random-train_{}-full_batch_test'.format(args.cluster_number)
logging.info(sub_dir)
log_dir = os.path.join(args.save, "tensorboard/")
writer = SummaryWriter(log_dir=log_dir)
args.in_channels = graph.x.size(-1)
args.num_tasks = dataset.num_classes
logging.info('%s' % args)
model = DeeperGCN(args).to(device)
logging.info(model)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0
}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
# generate batches
parts = random_partition_graph(graph.num_nodes,
cluster_number=args.cluster_number)
data = generate_sub_graphs(adj,
parts,
cluster_number=args.cluster_number)
# epoch_loss = train(data, model, graph.x, graph.y, train_idx, optimizer, device)
epoch_loss = train_flag(data, model, graph.x, graph.y, train_idx,
optimizer, device, args)
logging.info('Epoch {}, training loss {:.4f}'.format(
epoch, epoch_loss))
writer.add_scalar("Loss/train", epoch_loss, epoch)
model.print_params(epoch=epoch)
# if epoch % args.eval_epochs == 0:
# save_ckpt(model, optimizer,
# round(epoch_loss, 4), epoch,
# args.model_save_path,
# sub_dir, name_post=f'epoch{epoch}')
# logging.info(f"Epoch {epoch}, saved model to checkpoint folder {args.model_save_path}")
if epoch % args.eval_epochs == 0:
# save_ckpt(model, optimizer,
# round(epoch_loss, 4), epoch,
# args.model_save_path,
# sub_dir, name_post='final')
# logging.info(f"Saved model to checkpoint folder {args.model_save_path}")
logging.info(f'---- Evaluating at epoch {epoch} ----')
res = test_with_partition(model,
graph,
adj,
split_idx,
num_clusters=args.eval_cluster_number,
partition_method=args.partition_method,
evaluator=evaluator,
device=device)
# result = test(model, graph.x, graph.edge_index, graph.y, split_idx, evaluator)
logging.info(res)
logging.info(f"---------------------------------")
train_accuracy, valid_accuracy, test_accuracy = res[
"train_acc"], res["valid_acc"], res["test_acc"]
writer.add_scalar("Acc/train", train_accuracy)
writer.add_scalar("Acc/dev", valid_accuracy)
if train_accuracy > results['highest_train']:
results['highest_train'] = train_accuracy
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
sub_dir,
name_post='valid_best')
logging.info(
f"Saved better model to checkpoint folder {args.model_save_path}"
)
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
def main():
parser = argparse.ArgumentParser(description='OGBL-DDI (Full-Batch)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--use_sage', action='store_true')
parser.add_argument('--num_layers', type=int, default=2)
parser.add_argument('--hidden_channels', type=int, default=256)
parser.add_argument('--dropout', type=float, default=0.5)
parser.add_argument('--batch_size', type=int, default=64 * 1024)
parser.add_argument('--lr', type=float, default=0.005)
parser.add_argument('--epochs', type=int, default=200)
parser.add_argument('--eval_steps', type=int, default=5)
parser.add_argument('--runs', type=int, default=10)
args = parser.parse_args()
print(args)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
dataset = PygLinkPropPredDataset(name='ogbl-ddi')
split_edge = dataset.get_edge_split()
data = dataset[0]
# We randomly pick some training samples that we want to evaluate on:
torch.manual_seed(12345)
idx = torch.randperm(split_edge['train']['edge'].size(0))
idx = idx[:split_edge['valid']['edge'].size(0)]
split_edge['eval_train'] = {'edge': split_edge['train']['edge'][idx]}
edge_index = data.edge_index
adj = SparseTensor(row=edge_index[0], col=edge_index[1]).to(device)
if args.use_sage:
model = SAGE(args.hidden_channels, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout).to(device)
else:
model = GCN(args.hidden_channels, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout).to(device)
# Pre-compute GCN normalization.
adj = adj.set_diag()
deg = adj.sum(dim=1)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
adj = deg_inv_sqrt.view(-1, 1) * adj * deg_inv_sqrt.view(1, -1)
emb = torch.nn.Embedding(data.num_nodes, args.hidden_channels).to(device)
predictor = LinkPredictor(args.hidden_channels, args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
evaluator = Evaluator(name='ogbl-ddi')
loggers = {
'Hits@10': Logger(args.runs, args),
'Hits@20': Logger(args.runs, args),
'Hits@30': Logger(args.runs, args),
}
for run in range(args.runs):
torch.nn.init.xavier_uniform_(emb.weight)
model.reset_parameters()
predictor.reset_parameters()
optimizer = torch.optim.Adam(list(model.parameters()) +
list(emb.parameters()) +
list(predictor.parameters()),
lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(model, predictor, emb.weight, adj, data.edge_index,
split_edge, optimizer, args.batch_size)
if epoch % args.eval_steps == 0:
results = test(model, predictor, emb.weight, adj, split_edge,
evaluator, args.batch_size)
for key, result in results.items():
loggers[key].add_result(run, result)
if epoch % args.log_steps == 0:
for key, result in results.items():
train_hits, valid_hits, test_hits = result
print(key)
print(f'Run: {run + 1:02d}, '
f'Epoch: {epoch:02d}, '
f'Loss: {loss:.4f}, '
f'Train: {100 * train_hits:.2f}%, '
f'Valid: {100 * valid_hits:.2f}%, '
f'Test: {100 * test_hits:.2f}%')
print('---')
for key in loggers.keys():
print(key)
loggers[key].print_statistics(run)
for key in loggers.keys():
print(key)
loggers[key].print_statistics()
def main():
args = ArgsInit().save_exp()
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
dataset = PygNodePropPredDataset(name=args.dataset)
graph = dataset[0]
adj = SparseTensor(row=graph.edge_index[0], col=graph.edge_index[1])
if args.self_loop:
adj = adj.set_diag()
graph.edge_index = add_self_loops(edge_index=graph.edge_index,
num_nodes=graph.num_nodes)[0]
split_idx = dataset.get_idx_split()
train_idx = split_idx["train"].tolist()
evaluator = Evaluator(args.dataset)
sub_dir = 'random-train_{}-full_batch_test'.format(args.cluster_number)
logging.info(sub_dir)
args.in_channels = graph.x.size(-1)
args.num_tasks = dataset.num_classes
logging.info('%s' % args)
model = DeeperGCN(args).to(device)
logging.info(model)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0
}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
# generate batches
parts = random_partition_graph(graph.num_nodes,
cluster_number=args.cluster_number)
data = generate_sub_graphs(adj,
parts,
cluster_number=args.cluster_number)
epoch_loss = train(data, model, graph.x, graph.y, train_idx, optimizer,
device)
logging.info('Epoch {}, training loss {:.4f}'.format(
epoch, epoch_loss))
model.print_params(epoch=epoch)
if epoch == args.epochs:
result = test(model, graph.x, graph.edge_index, graph.y, split_idx,
evaluator)
logging.info(result)
train_accuracy, valid_accuracy, test_accuracy = result
if train_accuracy > results['highest_train']:
results['highest_train'] = train_accuracy
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
sub_dir,
name_post='valid_best')
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
