def __init__(self):
d_name = "ogbl-ppa"
dataset = LinkPropPredDataset(name=d_name)
splitted_edge = dataset.get_edge_split()
graph = dataset[0]
self.num_nodes = graph["num_nodes"]
self.ogb_evaluator = Evaluator(name="ogbl-ppa")
def main(_):
ds = LinkPropPredDataset(FLAGS.dataset)
split_edge = ds.get_edge_split()
train_edges = split_edge['train']['edge']
train_edges = np.concatenate([train_edges, train_edges[:, ::-1]], axis=0)
spa = scipy.sparse.csr_matrix(
(np.ones([len(train_edges)]), (train_edges[:, 0], train_edges[:, 1])))
mult_f = tf_fsvd.WYSDeepWalkPF(spa,
window=FLAGS.wys_window,
mult_degrees=False,
neg_sample_coef=FLAGS.wys_neg_coef)
tt = tqdm.tqdm(range(FLAGS.num_runs))
test_metrics = []
val_metrics = []
for run in tt:
u, s, v = tf_fsvd.fsvd(mult_f,
FLAGS.k,
n_iter=FLAGS.svd_iters,
n_redundancy=FLAGS.k * 3)
dataset = LinkPropPredDataset(FLAGS.dataset)
evaluator = Evaluator(name=FLAGS.dataset)
evaluator.K = FLAGS.hits
split_edge = dataset.get_edge_split()
metrics = []
for split in ('test', 'valid'):
pos_edges = split_edge[split]['edge']
neg_edges = split_edge[split]['edge_neg']
pos_scores = tf.reduce_sum(tf.gather(u * s, pos_edges[:, 0]) *
tf.gather(v, pos_edges[:, 1]),
axis=1).numpy()
neg_scores = tf.reduce_sum(tf.gather(u * s, neg_edges[:, 0]) *
tf.gather(v, neg_edges[:, 1]),
axis=1).numpy()
metric = evaluator.eval({
'y_pred_pos': pos_scores,
'y_pred_neg': neg_scores
})
metrics.append(metric['hits@%i' % FLAGS.hits])
test_metrics.append(metrics[0])
val_metrics.append(metrics[1])
tt.set_description(
'HITS@%i: validate=%g; test=%g' %
(FLAGS.hits, np.mean(val_metrics), np.mean(test_metrics)))
print('\n\n *** Trained for %i times and average metrics are:')
print('HITS@20 test: mean=%g; std=%g' %
(np.mean(test_metrics), np.std(test_metrics)))
print('HITS@20 validate: mean=%g; std=%g' %
(np.mean(val_metrics), np.std(val_metrics)))
def __init__(self, name):
start = time.time()
print("[I] Loading dataset %s..." % (name))
self.name = name
self.dataset = DglLinkPropPredDataset(name='ogbl-collab')
self.graph = self.dataset[0] # single DGL graph
# Create edge feat by concatenating weight and year
self.graph.edata['feat'] = torch.cat(
[self.graph.edata['edge_weight'], self.graph.edata['edge_year']],
dim=1
)
self.split_edge = self.dataset.get_edge_split()
self.train_edges = self.split_edge['train']['edge'] # positive train edges
self.val_edges = self.split_edge['valid']['edge'] # positive val edges
self.val_edges_neg = self.split_edge['valid']['edge_neg'] # negative val edges
self.test_edges = self.split_edge['test']['edge'] # positive test edges
self.test_edges_neg = self.split_edge['test']['edge_neg'] # negative test edges
self.evaluator = Evaluator(name='ogbl-collab')
print("[I] Finished loading.")
print("[I] Data load time: {:.4f}s".format(time.time()-start))
def evaluate_hits(name, pos_pred, neg_pred, K):
"""
Compute hits
Args:
name(str): name of dataset
pos_pred(Tensor): predict value of positive edges
neg_pred(Tensor): predict value of negative edges
K(int): num of hits
Returns:
hits(float): score of hits
"""
evaluator = Evaluator(name)
evaluator.K = K
hits = evaluator.eval({
'y_pred_pos': pos_pred,
'y_pred_neg': neg_pred,
})[f'hits@{K}']
return hits
def __init__(self, name):
start = time.time()
print("[I] Loading dataset %s..." % (name))
self.name = name
self.dataset = DglLinkPropPredDataset(name='ogbl-ddi')
self.graph = self.dataset[0] # single DGL graph
self.split_edge = self.dataset.get_edge_split()
self.train_edges = self.split_edge['train']['edge'] # positive train edges
self.val_edges = self.split_edge['valid']['edge'] # positive val edges
self.val_edges_neg = self.split_edge['valid']['edge_neg'] # negative val edges
self.test_edges = self.split_edge['test']['edge'] # positive test edges
self.test_edges_neg = self.split_edge['test']['edge_neg'] # negative test edges
self.evaluator = Evaluator(name='ogbl-ddi')
print("[I] Finished loading.")
print("[I] Data load time: {:.4f}s".format(time.time()-start))
class OgbEvaluator(object):
def __init__(self):
d_name = "ogbl-ppa"
dataset = LinkPropPredDataset(name=d_name)
splitted_edge = dataset.get_edge_split()
graph = dataset[0]
self.num_nodes = graph["num_nodes"]
self.ogb_evaluator = Evaluator(name="ogbl-ppa")
def eval(self, scores, labels, phase):
labels = np.reshape(labels, [-1])
ret = {}
pos = scores[labels > 0.5].squeeze(-1)
neg = scores[labels < 0.5].squeeze(-1)
for K in [10, 50, 100]:
self.ogb_evaluator.K = K
ret['%s_hits@%s' % (phase, K)] = self.ogb_evaluator.eval({
'y_pred_pos': pos,
'y_pred_neg': neg,
})[f'hits@{K}']
return ret
def main():
parser = argparse.ArgumentParser(description="OGBL-Citation2 (MLP)")
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("--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=100)
parser.add_argument("--eval_steps", type=int, default=10)
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-citation2")
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
if args.use_node_embedding:
embedding = torch.load("embedding.pt", map_location="cpu")
x = torch.cat([x, embedding], dim=-1)
x = x.to(device)
predictor = LinkPredictor(x.size(-1), args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
evaluator = Evaluator(name="ogbl-citation2")
logger = Logger(args.runs, args)
for run in range(args.runs):
predictor.reset_parameters()
optimizer = torch.optim.Adam(predictor.parameters(), lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(predictor, x, 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(predictor, x, 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}")
print("Node2vec" if args.use_node_embedding else "MLP")
logger.print_statistics(run)
print("Node2vec" if args.use_node_embedding else "MLP")
logger.print_statistics()
def main():
parser = argparse.ArgumentParser(description='OGBL-Citation (Cluster-GCN)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--num_partitions', type=int, default=15000)
parser.add_argument('--num_workers', type=int, default=12)
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=256)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--epochs', type=int, default=200)
parser.add_argument('--eval_steps', type=int, default=10)
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]
data.edge_index = to_undirected(data.edge_index, data.num_nodes)
cluster_data = ClusterData(data, num_parts=args.num_partitions,
recursive=False, save_dir=dataset.processed_dir)
loader = ClusterLoader(cluster_data, batch_size=args.batch_size,
shuffle=True, num_workers=args.num_workers)
# 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'],
}
model = GCN(data.x.size(-1), args.hidden_channels, args.hidden_channels,
args.num_layers, args.dropout).to(device)
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, loader, optimizer, device)
print(f'Run: {run + 1:02d}, Epoch: {epoch:02d}, Loss: {loss:.4f}')
if epoch > 49 and epoch % args.eval_steps == 0:
result = test(model, predictor, data, split_edge, evaluator,
batch_size=64 * 1024, device=device)
logger.add_result(run, result)
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='OGBL-COLLAB (GNN)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
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=5e-4)
parser.add_argument('--epochs', type=int, default=200)
parser.add_argument('--eval_steps', type=int, default=1)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--k', type=int, default=100)
parser.add_argument('--gpu_id', type=int, default=0)
args = parser.parse_args()
print(args)
device = gpu_setup(args.gpu_id)
dataset = PygLinkPropPredDataset(name='ogbl-collab')
data = dataset[0]
data.edge_weight = data.edge_weight.view(-1).to(torch.float)
data = T.ToSparseTensor()(data)
data = data.to(device)
split_edge = dataset.get_edge_split()
model = GCNWithAttention(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout, args.k).to(device)
predictor = LinkPredictor(args.hidden_channels, args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
print("model parameters {}".format(
sum(p.numel() for p in model.parameters())))
print("predictor parameters {}".format(
sum(p.numel() for p in predictor.parameters())))
print("total parameters {}".format(
sum(p.numel() for p in model.parameters()) +
sum(p.numel() for p in predictor.parameters())))
evaluator = Evaluator(name='ogbl-collab')
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, data, split_edge, optimizer,
args.batch_size)
if epoch % args.eval_steps == 0:
results = test(model, predictor, data, 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()
args.eval_metric = 'hits'
directed = False
else: # assume other datasets are undirected
args.eval_metric = 'auc'
directed = False
if args.use_valedges_as_input:
val_edge_index = split_edge['valid']['edge'].t()
if not directed:
val_edge_index = to_undirected(val_edge_index)
data.edge_index = torch.cat([data.edge_index, val_edge_index], dim=-1)
val_edge_weight = torch.ones([val_edge_index.size(1), 1], dtype=int)
data.edge_weight = torch.cat([data.edge_weight, val_edge_weight], 0)
if args.dataset.startswith('ogbl'):
evaluator = Evaluator(name=args.dataset)
if args.eval_metric == 'hits':
loggers = {
'Hits@20': Logger(args.runs, args),
'Hits@50': Logger(args.runs, args),
'Hits@100': Logger(args.runs, args),
}
elif args.eval_metric == 'mrr':
loggers = {
'MRR': Logger(args.runs, args),
}
elif args.eval_metric == 'auc':
loggers = {
'AUC': Logger(args.runs, args),
}
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():
parser = argparse.ArgumentParser(description='OGBL-DDI (GNN)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
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=1)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--k', type=int, default=50)
parser.add_argument('--gpu_id', type=int, default=0)
args = parser.parse_args()
print(args)
device = gpu_setup(args.gpu_id)
dataset = PygLinkPropPredDataset(name='ogbl-ddi',
transform=T.ToSparseTensor())
data = dataset[0]
adj_t = data.adj_t.to(device)
split_edge = dataset.get_edge_split()
# 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]}
model = GCNWithAttention(args.hidden_channels, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout, args.k).to(device)
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)
print("model parameters {}".format(sum(p.numel() for p in model.parameters())))
print("predictor parameters {}".format(sum(p.numel() for p in predictor.parameters())))
print("total parameters {}".format(data.num_nodes*args.hidden_channels +
sum(p.numel() for p in model.parameters())+sum(p.numel() for p in predictor.parameters())))
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_t, split_edge,
optimizer, args.batch_size)
if epoch % args.eval_steps == 0:
results = test(model, predictor, emb.weight, adj_t, 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():
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='OGBL-Citation2 (GraphSAINT)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
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=16 * 1024)
parser.add_argument('--walk_length', type=int, default=3)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--epochs', type=int, default=200)
parser.add_argument('--num_steps', type=int, default=100)
parser.add_argument('--eval_steps', type=int, default=10)
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-citation2')
split_edge = dataset.get_edge_split()
data = dataset[0]
data.edge_index = to_undirected(data.edge_index, data.num_nodes)
loader = GraphSAINTRandomWalkSampler(data,
batch_size=args.batch_size,
walk_length=args.walk_length,
num_steps=args.num_steps,
sample_coverage=0,
save_dir=dataset.processed_dir)
# 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'],
}
model = GCN(data.x.size(-1), args.hidden_channels, args.hidden_channels,
args.num_layers, args.dropout).to(device)
predictor = LinkPredictor(args.hidden_channels, args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
evaluator = Evaluator(name='ogbl-citation2')
logger = Logger(args.runs, args)
run_idx = 0
while run_idx < args.runs:
model.reset_parameters()
predictor.reset_parameters()
optimizer = torch.optim.Adam(list(model.parameters()) +
list(predictor.parameters()),
lr=args.lr)
run_success = True
for epoch in range(1, 1 + args.epochs):
loss = train(model, predictor, loader, optimizer, device)
print(
f'Run: {run_idx + 1:02d}, Epoch: {epoch:02d}, Loss: {loss:.4f}'
)
if loss > 2.:
run_success = False
logger.reset(run_idx)
print('Learning failed. Rerun...')
break
if epoch > 49 and epoch % args.eval_steps == 0:
result = test(model,
predictor,
data,
split_edge,
evaluator,
batch_size=64 * 1024,
device=device)
logger.add_result(run_idx, result)
train_mrr, valid_mrr, test_mrr = result
print(f'Run: {run_idx + 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}')
print('GraphSAINT')
if run_success:
logger.print_statistics(run_idx)
run_idx += 1
print('GraphSAINT')
logger.print_statistics()
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]}
# Load model, initial embeddings, the link predictor, and the evaluator
model_hgsage = HG_SAGE(data.num_nodes, dropout=dropout, hidden_dim=hidden_channels).to(device)
model = model_hgsage
emb = torch.nn.Embedding(data.num_nodes, hidden_channels).to(device)
predictor = LinkPredictor(hidden_channels, hidden_channels, 1,
num_layers, dropout).to(device)
evaluator = Evaluator(name='ogbl-ddi')
# Train and evaluate model
best_train, best_val, best_test = [], [], []
for run in range(runs):
best_train_r, best_val_r, best_test_r = None, None, None
torch.nn.init.xavier_uniform_(emb.weight)
predictor.reset_parameters()
optimizer = torch.optim.Adam(
list(model.parameters()) + list(emb.parameters()) +
list(predictor.parameters()), lr=lr)
for epoch in range(1, 1 + epochs):
loss = train(model, predictor, emb.weight, adj_t, theta, split_edge,
optimizer, batch_size)
def main():
parser = argparse.ArgumentParser(description='OGBL-DDI')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--model',
type=str,
default='MAD_GCN',
choices=[
'GCN_Linear', 'SAGE_Linear', 'MAD_GCN', 'MAD_SAGE',
'MAD_Model'
])
parser.add_argument('--train_batch_size', type=int, default=4096)
parser.add_argument('--test_batch_size', type=int, default=1024)
parser.add_argument('--lr', type=float, default=0.005)
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--eval_steps', type=int, default=5)
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 = PygLinkPropPredDataset(name='ogbl-ddi',
transform=T.ToSparseTensor())
data = dataset[0]
adj_t = data.adj_t.to(device)
split_edge = dataset.get_edge_split()
# 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]}
model = models.get_model(args.model)(data.num_nodes, adj_t).to(device)
print(f"Parameters: {count_parameters(model)}")
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):
model.reset_parameters()
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(model, adj_t, split_edge, optimizer,
args.train_batch_size)
if epoch % args.eval_steps == 0:
results = test(model, adj_t, split_edge, evaluator,
args.test_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('---')
print(f'Finished epoch {epoch}')
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='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()
os.path.join(log_dir, "%s.log" % (dataset_name + "_" + setting_name))
])
sys.stderr = Logger([
"%s.log" % (dataset_name + setting_name),
os.path.join(log_dir, "%s.log" % (dataset_name + "_" + setting_name))
])
snapshot_dir = os.path.join(log_dir, "snapshot")
if not os.path.isdir(snapshot_dir):
os.makedirs(snapshot_dir)
print("Process Id:", os.getpid())
print(os.path.join(log_dir, sys.argv[0]))
print(args)
shutil.copyfile(__file__, os.path.join(log_dir, "train.py"))
shutil.copyfile(model_file + ".py", os.path.join(log_dir, model_file + ".py"))
evaluator = Evaluator(name="ogbl-citation")
print(evaluator.expected_input_format)
print(evaluator.expected_output_format)
dataset = DglLinkPropPredDataset(name="ogbl-citation")
split_edge = dataset.get_edge_split()
num_worker = 16
train_edge, valid_edge, test_edge = split_edge["train"], split_edge[
"valid"], split_edge["test"]
graph = dataset[0]
origin_graph = copy.deepcopy(graph)
graph.readonly(False)
graph.add_edges(graph.edges()[1], graph.edges()[0])
graph.add_edges(
torch.arange(0, graph.number_of_nodes()).long(),
torch.arange(0, graph.number_of_nodes()).long())
def main():
parser = argparse.ArgumentParser(description="OGBL-COLLAB (MLP)")
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("--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=200)
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-collab")
split_edge = dataset.get_edge_split()
data = dataset[0]
x = data.x
if args.use_node_embedding:
embedding = torch.load("embedding.pt", map_location="cpu")
x = torch.cat([x, embedding], dim=-1)
x = x.to(device)
predictor = LinkPredictor(x.size(-1), args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
evaluator = Evaluator(name="ogbl-collab")
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):
predictor.reset_parameters()
optimizer = torch.optim.Adam(predictor.parameters(), lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(predictor, x, split_edge, optimizer, args.batch_size)
if epoch % args.eval_steps == 0:
results = test(predictor, x, 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 run(file, data_name, model_name,lr):
parser = argparse.ArgumentParser(description='OGBL-DDI (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=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.001)
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)
parser.add_argument('--use_nd', action='store_true')
parser.add_argument('--use_lgae', action='store_true')
parser.add_argument('--use_vgae', action='store_true')
parser.add_argument('--model', type=str, default='')
parser.add_argument('--dataset', type=str, default='Citeseer')
args = parser.parse_args()
if data_name != None and model_name != None and lr != None:
args.dataset = data_name
args.model = model_name
args.lr = lr
print(args)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
device = torch.device(device)
dataset = CitationFull(os.path.join('citation_data',args.dataset),name=args.dataset,transform=T.ToSparseTensor())
num_training = int(dataset.__len__()*0.8)
num_val = int(dataset.__len__()*0.1)
num_test = dataset.__len__() - (num_training+num_val)
data = dataset[0]
adj_t = data.adj_t.to(device)
edge_index, edge_type = utils.dense_to_sparse(adj_t.to_dense())
data.edge_index = edge_index
data.x = data.x.to(device)
num_nodes = data.x.shape[0]
num_edges = data.edge_index.shape[1]
print(data)
# nx_data = to_networkx(data, to_undirected=True)
# print('graph density='+str(2*num_edges/(num_nodes*(num_nodes-1))))
# print('clustering coefficient='+str(nx.average_clustering(nx_data)))
decoder_enable = args.model[-3:]
if args.model[-3:] == '-nd': model_name = args.model[:-3]
if model_name == 'lgae':
model = LGAE(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'vgae':
model = DeepVGAE(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'gae':
model = GraphAutoEncoder(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'arga':
model = AdversarialGAE(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'arvga':
model = AdversarialVGAE(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'lrga':
model = LRGA(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'sage':
model = SAGEAutoEncoder(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
if decoder_enable == '-nd':
model.decoder = NeuralDecoder( args.hidden_channels,
args.hidden_channels, 1, args.num_layers, args.dropout)
evaluator = Evaluator(name='ogbl-ddi')
model = model.to(device)
loggers = {}
K_list = ['20','50','100']
for k in K_list:
loggers['Hits@'+k] = Logger(args.runs, args)
for run in range(args.runs):
torch.manual_seed(run)
split_edge = utils.train_test_split_edges(data)
# print(split_edge.train_pos_edge_index.shape)
# print(split_edge.val_pos_edge_index.shape)
# exit()
split_edge.edge_index = edge_index
# emb.weight.data = features
model.reset_parameters()
if args.model in ['arga','arga-nd','arvga','arvga-nd']:
args.lr=0.005
optimizer = torch.optim.Adam(
list(model.parameters()), lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(model, data.x, adj_t, split_edge,
optimizer, args.batch_size)
if epoch % args.eval_steps == 0:
results = test(model, data.x, adj_t, 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, test_auc, test_ap, val_auc, val_ap = result
print(key)
print(f'Run: {run + 1:02d}, '
f'Epoch: {epoch:02d}, '
f'Loss: {loss:.4f}, '
f'auc: {100 * test_auc:.2f}%, '
f'ap: {100 * test_ap:.2f}%, '
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)
toWrite = loggers[key].print_statistics()
file.write(str(args.lr)+' ' +key + ' ' +args.model+"'"+str(toWrite)+'\n')
file.flush()
def main():
parser = argparse.ArgumentParser(
description='Link prediction (Cluster-GCN)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--dataset', type=str, default='ogbl-citation')
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--num_partitions', type=int, default=15000)
parser.add_argument('--num_workers', type=int, default=4)
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=256)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--epochs', type=int, default=200)
parser.add_argument('--eval_steps', type=int, default=10)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--negs', type=int, default=1)
parser.add_argument('--gnn_type', type=str, default='gcn')
args = parser.parse_args()
print(args)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
dataset = DglLinkPropPredDataset(name=args.dataset)
# Manually add self-loop link since GCN will wash out the feature of isolated nodes.
n_nodes = dataset[0].number_of_nodes()
g_data = dgl.add_self_loop(dataset[0])
g_data = dgl.to_bidirected(g_data)
for k in dataset[0].node_attr_schemes().keys():
g_data.ndata[k] = dataset[0].ndata[k]
print(g_data.number_of_nodes(), g_data.number_of_edges())
g_data.create_formats_()
cluster_dataset = ClusterIterDataset(args.dataset,
g_data,
args.num_partitions,
use_pp=False)
cluster_iterator = DataLoader(cluster_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=partial(subgraph_collate_fn,
g_data,
negs=args.negs))
model = GCN(g_data.ndata['feat'].size(-1),
args.hidden_channels,
args.hidden_channels,
args.num_layers,
args.dropout,
gnn_type=args.gnn_type).to(device)
predictor = LinkPredictor(args.hidden_channels, args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
evaluator = Evaluator(name=args.dataset)
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)
epoch_time, to_device_time, ff_time, pred_loss_time, bp_time, io_time, memory, part_1 = 0, 0, 0, 0, 0, 0, 0, 0
for epoch in range(1, 1 + args.epochs):
loss, c_epoch_time, c_to_device_time, c_ff_time, c_pred_loss_time, c_bp_time, c_io_time, c_memory, c_part1 = train(
model, predictor, cluster_iterator, optimizer, device)
print(f'Run: {run + 1:02d}, Epoch: {epoch:02d}, Loss: {loss:.4f}')
epoch_time += c_epoch_time
to_device_time += c_to_device_time
ff_time += c_ff_time
pred_loss_time += c_pred_loss_time
bp_time += c_bp_time
io_time += c_io_time
part_1 += c_part1
memory = max(memory, c_memory[0])
if epoch % args.eval_steps == 0:
print('Ave')
print('epoch time: ', epoch_time / args.eval_steps)
print('to_device_time: ', to_device_time / args.eval_steps)
print('ff_time: ', ff_time / args.eval_steps)
print('part1_time: ', part_1 / args.eval_steps)
print('pred_loss_time: ', pred_loss_time / args.eval_steps)
print('bp_time: ', bp_time / args.eval_steps)
print('io_time: ', io_time / args.eval_steps)
print('max memory', memory)
print('\n')
epoch_time, to_device_time, ff_time, pred_loss_time, bp_time, io_time, memory, part_1 = 0, 0, 0, 0, 0, 0, 0, 0
result = test(model, predictor, g_data, split_edge, evaluator,
64 * 4 * args.batch_size, device)
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="OGBL-PPA (MF)")
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--log_steps", type=int, default=1)
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.005)
parser.add_argument("--epochs", type=int, default=1000)
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")
split_edge = dataset.get_edge_split()
data = dataset[0]
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-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):
emb.reset_parameters()
predictor.reset_parameters()
optimizer = torch.optim.Adam(list(emb.parameters()) +
list(predictor.parameters()),
lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(emb.weight, predictor, split_edge, optimizer,
args.batch_size)
if epoch % args.eval_steps == 0:
results = test(emb.weight, predictor, 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}%")
for key in loggers.keys():
print(key)
loggers[key].print_statistics(run)
for key in loggers.keys():
print(key)
loggers[key].print_statistics()
def run(file, data_name, model_name,lr):
parser = argparse.ArgumentParser(description='OGBL-DDI (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=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.001)
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)
parser.add_argument('--use_nd', action='store_true')
parser.add_argument('--use_lgae', action='store_true')
parser.add_argument('--use_vgae', action='store_true')
parser.add_argument('--model', type=str, default='')
parser.add_argument('--dataset', type=str, default='Citeseer')
args = parser.parse_args()
if data_name != None and model_name != None and lr != None:
args.dataset = data_name
args.model = model_name
args.lr = lr
print(args)
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
device = torch.device(device)
dataset = CitationFull(os.path.join('citation_data',args.dataset),name=args.dataset,transform=T.ToSparseTensor())
num_training = int(dataset.__len__()*0.8)
num_val = int(dataset.__len__()*0.1)
num_test = dataset.__len__() - (num_training+num_val)
data = dataset[0]
print('data:',vars(data))
adj_t = data.adj_t.to(device)
edge_index, edge_type = utils.dense_to_sparse(adj_t.to_dense())
data.edge_index = edge_index
data.x = data.x.to(device)
split_edge = utils.train_test_split_edges(data)
split_edge.edge_index = edge_index
print(data)
print(edge_index.shape)
decoder_enable = args.model[-3:]
if args.model[-3:] == '-nd': model_name = args.model[:-3]
if model_name == 'lgae':
model = LGAE(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'vgae':
model = DeepVGAE(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'gae':
model = GraphAutoEncoder(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'arga':
model = AdversarialGAE(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'arvga':
model = AdversarialVGAE(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'lrga':
model = LRGA(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
elif model_name == 'sage':
model = SAGEAutoEncoder(data.num_features, args.hidden_channels,
args.hidden_channels, args.num_layers,
args.dropout)
if decoder_enable == '-nd':
model.decoder = NeuralDecoder( args.hidden_channels,
args.hidden_channels, 1, args.num_layers, args.dropout)
evaluator = Evaluator(name='ogbl-ddi')
model = model.to(device)
loggers = {
'metrics': Logger(args.runs, args)
}
for run in range(args.runs):
torch.manual_seed(run)
model.reset_parameters()
if args.model in ['arga','arga-nd','arvga','arvga-nd']:
args.lr=0.005
optimizer = torch.optim.Adam(
list(model.parameters()), lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(model, data.x, adj_t, split_edge,
optimizer, args.batch_size)
result = test(model, data.x, data, split_edge, evaluator, args.batch_size)
loggers['metrics'].add_result(run, result)
for key in loggers.keys():
print(key)
toWrite = loggers[key].print_statistics()
file.write(args.model+'\t'+'\t'.join(toWrite)+'\n')
file.flush()
os.fsync(file)
def main():
parser = argparse.ArgumentParser()
add_train_args(parser)
add_common_args(parser)
args = parser.parse_args()
add_experiment(args)
device = model_utils.get_device()
# Load dataset from disk
print('Loading train data...')
train_graph, valid_graph, train_edges, eval_edges, valid_edges = model_utils.load_training_data()
if args.train_partial_graph:
train_edges['edge'] = eval_edges['edge']
train_dl = data.DataLoader(
data.TensorDataset(train_edges['edge']),
batch_size=args.train_batch_size,
shuffle=True,
)
dev_dl = data.DataLoader(
data.TensorDataset(
torch.cat([valid_edges['edge'], valid_edges['edge_neg']], dim=0),
torch.cat([torch.ones(valid_edges['edge'].shape[0]),
torch.zeros(valid_edges['edge_neg'].shape[0])], dim=0),
),
batch_size=args.val_batch_size,
shuffle=True,
)
# Initialize node embeddings
print('Computing initial embeddings')
train_graph = model_utils.initialize_embeddings(
train_graph, 'train_embeddings.pt', args.refresh_embeddings)
valid_graph = model_utils.initialize_embeddings(
valid_graph, 'valid_embeddings.pt', args.refresh_embeddings)
if not args.train_partial_graph:
train_graph = valid_graph
# Stats evaluator
evaluator = Evaluator(name='ogbl-ddi')
# Initialize a model
model = models.get_model(args.model)(
# train_graph.x.shape, train_graph.adj_t.to(device)
num_nodes=train_graph.num_nodes, adj_t=train_graph.adj_t.to(device)
)
# load from checkpoint if path specified
if args.load_path is not None:
model = model_utils.load_model(model, args.load_path)
print(f"Parameters: {model_utils.count_parameters(model)}")
# Move model to GPU if necessary
model.to(device)
# Initialize optimizer
optimizer = optim.Adam(
model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay,
)
# Scheduler
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode='min',
factor=0.5,
patience=30,
verbose=True,
)
os.makedirs(f'{args.save_path}/{args.experiment}')
print(f'Created new experiment: {args.experiment}')
save_arguments(args, f'{args.save_path}/{args.experiment}/args.txt')
# Train!
trained_model = train_model(
train_graph,
valid_graph,
train_dl,
dev_dl,
evaluator,
model,
optimizer,
scheduler,
args,
)
# Save trained model
filename = f'{args.save_path}/{args.experiment}/{model.__class__.__name__}_trained.checkpoint'
model_utils.save_model(trained_model, filename)
def main():
parser = argparse.ArgumentParser(description='OGBL-PPA (MLP)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
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')
splitted_edge = dataset.get_edge_split()
data = dataset[0]
x = data.x.to(torch.float)
embedding = torch.load('embedding.pt', map_location='cpu')
x = torch.cat([x, embedding], dim=-1)
x = x.to(device)
predictor = LinkPredictor(x.size(-1), args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
optimizer = torch.optim.Adam(predictor.parameters(), lr=args.lr)
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):
predictor.reset_parameters()
for epoch in range(1, 1 + args.epochs):
loss = train(predictor, x, splitted_edge, optimizer,
args.batch_size)
if epoch % args.eval_steps == 0:
results = test(predictor, x, 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='OGBL-Citation2 (MLP)')
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('--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=100)
parser.add_argument('--eval_steps', type=int, default=10)
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-citation2')
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
if args.use_node_embedding:
embedding = torch.load('embedding.pt', map_location='cpu')
x = torch.cat([x, embedding], dim=-1)
x = x.to(device)
predictor = LinkPredictor(x.size(-1), args.hidden_channels, 1,
args.num_layers, args.dropout).to(device)
evaluator = Evaluator(name='ogbl-citation2')
logger = Logger(args.runs, args)
for run in range(args.runs):
predictor.reset_parameters()
optimizer = torch.optim.Adam(predictor.parameters(), lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(predictor, x, 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(predictor, x, 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}')
print('Node2vec' if args.use_node_embedding else 'MLP')
logger.print_statistics(run)
print('Node2vec' if args.use_node_embedding else 'MLP')
logger.print_statistics()
def main(args):
if (not args.do_train) and (not args.do_valid) and (not args.do_test) and (
not args.evaluate_train):
raise ValueError('one of train/val/test mode must be choosed.')
if args.init_checkpoint:
override_config(args)
args.save_path = 'log/%s/%s/%s-%s/%s' % (
args.dataset, args.model, args.hidden_dim, args.gamma,
time.time()) if args.save_path == None else args.save_path
writer = SummaryWriter(args.save_path)
# Write logs to checkpoint and console
set_logger(args)
dataset = LinkPropPredDataset(name='ogbl-biokg')
split_edge = dataset.get_edge_split()
train_triples, valid_triples, test_triples = split_edge[
"train"], split_edge["valid"], split_edge["test"]
nrelation = int(max(train_triples['relation'])) + 1
entity_dict = dict()
cur_idx = 0
for key in dataset[0]['num_nodes_dict']:
entity_dict[key] = (cur_idx,
cur_idx + dataset[0]['num_nodes_dict'][key])
cur_idx += dataset[0]['num_nodes_dict'][key]
nentity = sum(dataset[0]['num_nodes_dict'].values())
evaluator = Evaluator(name=args.dataset)
args.nentity = nentity
args.nrelation = nrelation
logging.info('Model: %s' % args.model)
logging.info('Dataset: %s' % args.dataset)
logging.info('#entity: %d' % nentity)
logging.info('#relation: %d' % nrelation)
# train_triples = split_dict['train']
logging.info('#train: %d' % len(train_triples['head']))
# valid_triples = split_dict['valid']
logging.info('#valid: %d' % len(valid_triples['head']))
# test_triples = split_dict['test']
logging.info('#test: %d' % len(test_triples['head']))
train_count, train_true_head, train_true_tail = defaultdict(
lambda: 4), defaultdict(list), defaultdict(list)
for i in tqdm(range(len(train_triples['head']))):
head, relation, tail = train_triples['head'][i], train_triples[
'relation'][i], train_triples['tail'][i]
head_type, tail_type = train_triples['head_type'][i], train_triples[
'tail_type'][i]
train_count[(head, relation, head_type)] += 1
train_count[(tail, -relation - 1, tail_type)] += 1
train_true_head[(relation, tail)].append(head)
train_true_tail[(head, relation)].append(tail)
kge_model = KGEModel(
model_name=args.model,
nentity=nentity,
nrelation=nrelation,
hidden_dim=args.hidden_dim,
gamma=args.gamma,
double_entity_embedding=args.double_entity_embedding,
double_relation_embedding=args.double_relation_embedding,
evaluator=evaluator)
logging.info('Model Parameter Configuration:')
for name, param in kge_model.named_parameters():
logging.info('Parameter %s: %s, require_grad = %s' %
(name, str(param.size()), str(param.requires_grad)))
if args.cuda:
kge_model = kge_model.cuda()
if args.init_checkpoint:
# Restore model from checkpoint directory
logging.info('Loading checkpoint %s...' % args.init_checkpoint)
checkpoint = torch.load(
os.path.join(args.init_checkpoint, 'checkpoint'))
entity_dict = checkpoint['entity_dict']
if args.do_train:
# Set training dataloader iterator
train_dataloader_head = DataLoader(
TrainDataset(train_triples, nentity, nrelation,
args.negative_sample_size, 'head-batch', train_count,
train_true_head, train_true_tail, entity_dict),
batch_size=args.batch_size,
shuffle=True,
num_workers=max(1, args.cpu_num // 2),
collate_fn=TrainDataset.collate_fn)
train_dataloader_tail = DataLoader(
TrainDataset(train_triples, nentity, nrelation,
args.negative_sample_size, 'tail-batch', train_count,
train_true_head, train_true_tail, entity_dict),
batch_size=args.batch_size,
shuffle=True,
num_workers=max(1, args.cpu_num // 2),
collate_fn=TrainDataset.collate_fn)
train_iterator = BidirectionalOneShotIterator(train_dataloader_head,
train_dataloader_tail)
# Set training configuration
current_learning_rate = args.learning_rate
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
kge_model.parameters()),
lr=current_learning_rate)
if args.warm_up_steps:
warm_up_steps = args.warm_up_steps
else:
warm_up_steps = args.max_steps // 2
if args.init_checkpoint:
# Restore model from checkpoint directory
# logging.info('Loading checkpoint %s...' % args.init_checkpoint)
# checkpoint = torch.load(os.path.join(args.init_checkpoint, 'checkpoint'))
init_step = checkpoint['step']
kge_model.load_state_dict(checkpoint['model_state_dict'])
# entity_dict = checkpoint['entity_dict']
if args.do_train:
current_learning_rate = checkpoint['current_learning_rate']
warm_up_steps = checkpoint['warm_up_steps']
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
else:
logging.info('Ramdomly Initializing %s Model...' % args.model)
init_step = 0
step = init_step
logging.info('Start Training...')
logging.info('init_step = %d' % init_step)
logging.info('batch_size = %d' % args.batch_size)
logging.info('negative_adversarial_sampling = %d' %
args.negative_adversarial_sampling)
logging.info('hidden_dim = %d' % args.hidden_dim)
logging.info('gamma = %f' % args.gamma)
logging.info('negative_adversarial_sampling = %s' %
str(args.negative_adversarial_sampling))
if args.negative_adversarial_sampling:
logging.info('adversarial_temperature = %f' %
args.adversarial_temperature)
# Set valid dataloader as it would be evaluated during training
if args.do_train:
logging.info('learning_rate = %d' % current_learning_rate)
training_logs = []
#Training Loop
for step in range(init_step, args.max_steps):
log = kge_model.train_step(kge_model, optimizer, train_iterator,
args)
training_logs.append(log)
if step >= warm_up_steps:
current_learning_rate = current_learning_rate / 10
logging.info('Change learning_rate to %f at step %d' %
(current_learning_rate, step))
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
kge_model.parameters()),
lr=current_learning_rate)
warm_up_steps = warm_up_steps * 3
if step % args.save_checkpoint_steps == 0 and step > 0: # ~ 41 seconds/saving
save_variable_list = {
'step': step,
'current_learning_rate': current_learning_rate,
'warm_up_steps': warm_up_steps,
'entity_dict': entity_dict
}
save_model(kge_model, optimizer, save_variable_list, args)
if step % args.log_steps == 0:
metrics = {}
for metric in training_logs[0].keys():
metrics[metric] = sum(
[log[metric]
for log in training_logs]) / len(training_logs)
log_metrics('Train', step, metrics, writer)
training_logs = []
if args.do_valid and step % args.valid_steps == 0 and step > 0:
logging.info('Evaluating on Valid Dataset...')
metrics = kge_model.test_step(kge_model, valid_triples, args,
entity_dict)
log_metrics('Valid', step, metrics, writer)
save_variable_list = {
'step': step,
'current_learning_rate': current_learning_rate,
'warm_up_steps': warm_up_steps
}
save_model(kge_model, optimizer, save_variable_list, args)
if args.do_valid:
logging.info('Evaluating on Valid Dataset...')
metrics = kge_model.test_step(kge_model, valid_triples, args,
entity_dict)
log_metrics('Valid', step, metrics, writer)
if args.do_test:
logging.info('Evaluating on Test Dataset...')
metrics = kge_model.test_step(kge_model, test_triples, args,
entity_dict)
log_metrics('Test', step, metrics, writer)
if args.evaluate_train:
logging.info('Evaluating on Training Dataset...')
small_train_triples = {}
indices = np.random.choice(len(train_triples['head']),
args.ntriples_eval_train,
replace=False)
for i in train_triples:
if 'type' in i:
small_train_triples[i] = [train_triples[i][x] for x in indices]
else:
small_train_triples[i] = train_triples[i][indices]
metrics = kge_model.test_step(kge_model,
small_train_triples,
args,
entity_dict,
random_sampling=True)
log_metrics('Train', step, metrics, writer)
def main():
parser = argparse.ArgumentParser(description="OGBL-COLLAB (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("--use_valedges_as_input", 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.001)
parser.add_argument("--epochs", type=int, default=400)
parser.add_argument("--eval_steps", type=int, default=1)
parser.add_argument("--runs", type=int, default=1)
parser.add_argument("--seed",type=int,default=1)
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-collab")
data = dataset[0]
edge_index = data.edge_index
data.edge_weight = data.edge_weight.view(-1).to(torch.float)
data = T.ToSparseTensor()(data)
split_edge = dataset.get_edge_split()
# Use training + validation edges for inference on test set.
if args.use_valedges_as_input:
val_edge_index = split_edge["valid"]["edge"].t()
full_edge_index = torch.cat([edge_index, val_edge_index], dim=-1)
data.full_adj_t = SparseTensor.from_edge_index(full_edge_index).t()
data.full_adj_t = data.full_adj_t.to_symmetric()
else:
data.full_adj_t = data.adj_t
data = data.to(device)
if args.use_sage:
model = SAGE(
data.num_features,
args.hidden_channels,
args.hidden_channels,
args.num_layers,
args.dropout,
).to(device)
else:
model = GCN(
data.num_features,
args.hidden_channels,
args.hidden_channels,
args.num_layers,
args.dropout,
).to(device)
predictor = LinkPredictor(
args.hidden_channels, args.hidden_channels, 1, args.num_layers, args.dropout
).to(device)
evaluator = Evaluator(name="ogbl-collab")
loggers = {
"Hits@10": Logger(args.runs, args),
"Hits@50": Logger(args.runs, args),
"Hits@100": Logger(args.runs, args),
}
for run in tqdm(range(args.runs)):
torch.manual_seed(args.seed + run)
np.random.seed(args.seed+run)
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, data, split_edge, optimizer, args.batch_size)
if epoch % args.eval_steps == 0:
results = test(
model, predictor, data, 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():
parser = argparse.ArgumentParser(description='OGBL-Reviews (MLP)')
parser.add_argument('--suffix', type=str, default='groc')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--use_node_features', action='store_true')
parser.add_argument('--hidden_channels', type=int, default=64)
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=200)
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=f'ogbl-reviews-{args.suffix}')
splitted_edge = dataset.get_edge_split()
data = dataset[0]
product_idx = (data.x[:, 0] == 0).nonzero().flatten()
user_idx = (data.x[:, 0] == 1).nonzero().flatten()
user_emb = torch.nn.Parameter(
torch.Tensor(user_idx.size(0), args.hidden_channels).to(device))
if args.use_node_features:
product_x = data.x[product_idx, 1:].to(device)
else:
product_x = None
product_emb = torch.nn.Parameter(
torch.Tensor(product_idx.size(0), args.hidden_channels).to(device))
for split in ['train', 'valid', 'test']:
edge = splitted_edge[f'{split}_edge'].t()
edge[1] -= product_emb.size(0)
splitted_edge[f'{split}_edge'] = edge
predictor = LinkPredictor(
args.hidden_channels + (300 if product_x is not None else 0), 1,
args.dropout).to(device)
evaluator = Evaluator(name=f'ogbl-reviews-{args.suffix}')
logger = Logger(args.runs, args)
for run in range(args.runs):
predictor.reset_parameters()
glorot(user_emb)
glorot(product_emb)
optimizer = torch.optim.Adam([user_emb, product_emb] +
list(predictor.parameters()),
lr=args.lr)
for epoch in range(1, 1 + args.epochs):
loss = train(predictor, product_emb, user_emb, product_x,
splitted_edge, optimizer, args.batch_size)
result = test(predictor, product_emb, user_emb, product_x,
splitted_edge, evaluator, args.batch_size)
logger.add_result(run, result)
if epoch % args.log_steps == 0:
train_rmse, valid_rmse, test_rmse = result
print(f'Run: {run + 1:02d}, '
f'Epoch: {epoch:02d}, '
f'Loss: {loss:.4f}, '
f'Train: {train_rmse:.4f}, '
f'Valid: {valid_rmse:.4f}, '
f'Test: {test_rmse:.4f}')
logger.print_statistics(run)
logger.print_statistics()
def train_model(
train_graph: pyg.torch_geometric.data.Data,
valid_graph: pyg.torch_geometric.data.Data,
train_dl: data.DataLoader,
dev_dl: data.DataLoader,
evaluator: Evaluator,
model: nn.Module,
optimizer: optim.Optimizer,
lr_scheduler: optim.lr_scheduler._LRScheduler,
args: argparse.Namespace,
) -> nn.Module:
device = model_utils.get_device()
loss_fn = nn.functional.binary_cross_entropy
val_loss_fn = nn.functional.binary_cross_entropy
best_val_loss = torch.tensor(float('inf'))
best_val_hits = torch.tensor(0.0)
saved_checkpoints = []
writer = SummaryWriter(log_dir=f'{args.log_dir}/{args.experiment}')
for e in range(1, args.train_epochs + 1):
print(f'Training epoch {e}...')
# Training portion
torch.cuda.empty_cache()
torch.set_grad_enabled(True)
with tqdm(total=args.train_batch_size * len(train_dl)) as progress_bar:
model.train()
# Load graph into GPU
adj_t = train_graph.adj_t.to(device)
edge_index = train_graph.edge_index.to(device)
x = train_graph.x.to(device)
pos_pred = []
neg_pred = []
for i, (y_pos_edges,) in enumerate(train_dl):
y_pos_edges = y_pos_edges.to(device).T
y_neg_edges = negative_sampling(
edge_index,
num_nodes=train_graph.num_nodes,
num_neg_samples=y_pos_edges.shape[1]
).to(device)
y_batch = torch.cat([torch.ones(y_pos_edges.shape[1]), torch.zeros(
y_neg_edges.shape[1])], dim=0).to(device) # Ground truth edge labels (1 or 0)
# Forward pass on model
optimizer.zero_grad()
y_pred = model(adj_t, torch.cat(
[y_pos_edges, y_neg_edges], dim=1))
loss = loss_fn(y_pred, y_batch)
# Backward pass and optimization
loss.backward()
optimizer.step()
if args.use_scheduler:
lr_scheduler.step(loss)
batch_acc = torch.mean(
1 - torch.abs(y_batch.detach() - torch.round(y_pred.detach()))).item()
pos_pred += [y_pred[y_batch == 1].detach()]
neg_pred += [y_pred[y_batch == 0].detach()]
progress_bar.update(y_pos_edges.shape[1])
progress_bar.set_postfix(loss=loss.item(), acc=batch_acc)
writer.add_scalar(
"train/Loss", loss, ((e - 1) * len(train_dl) + i) * args.train_batch_size)
writer.add_scalar("train/Accuracy", batch_acc,
((e - 1) * len(train_dl) + i) * args.train_batch_size)
del y_pos_edges
del y_neg_edges
del y_pred
del loss
del adj_t
del edge_index
del x
# Training set evaluation Hits@K Metrics
pos_pred = torch.cat(pos_pred, dim=0)
neg_pred = torch.cat(neg_pred, dim=0)
results = {}
for K in [10, 20, 30]:
evaluator.K = K
hits = evaluator.eval({
'y_pred_pos': pos_pred,
'y_pred_neg': neg_pred,
})[f'hits@{K}']
results[f'Hits@{K}'] = hits
print()
print(f'Train Statistics')
print('*' * 30)
for k, v in results.items():
print(f'{k}: {v}')
writer.add_scalar(
f"train/{k}", v, (pos_pred.shape[0] + neg_pred.shape[0]) * e)
print('*' * 30)
del pos_pred
del neg_pred
# Validation portion
torch.cuda.empty_cache()
torch.set_grad_enabled(False)
with tqdm(total=args.val_batch_size * len(dev_dl)) as progress_bar:
model.eval()
adj_t = valid_graph.adj_t.to(device)
edge_index = valid_graph.edge_index.to(device)
x = valid_graph.x.to(device)
val_loss = 0.0
accuracy = 0
num_samples_processed = 0
pos_pred = []
neg_pred = []
for i, (edges_batch, y_batch) in enumerate(dev_dl):
edges_batch = edges_batch.T.to(device)
y_batch = y_batch.to(device)
# Forward pass on model in validation environment
y_pred = model(adj_t, edges_batch)
loss = val_loss_fn(y_pred, y_batch)
num_samples_processed += edges_batch.shape[1]
batch_acc = torch.mean(
1 - torch.abs(y_batch - torch.round(y_pred))).item()
accuracy += batch_acc * edges_batch.shape[1]
val_loss += loss.item() * edges_batch.shape[1]
pos_pred += [y_pred[y_batch == 1].detach()]
neg_pred += [y_pred[y_batch == 0].detach()]
progress_bar.update(edges_batch.shape[1])
progress_bar.set_postfix(
val_loss=val_loss / num_samples_processed,
acc=accuracy/num_samples_processed)
writer.add_scalar(
"Val/Loss", loss, ((e - 1) * len(dev_dl) + i) * args.val_batch_size)
writer.add_scalar(
"Val/Accuracy", batch_acc, ((e - 1) * len(dev_dl) + i) * args.val_batch_size)
del edges_batch
del y_batch
del y_pred
del loss
del adj_t
del edge_index
del x
# Validation evaluation Hits@K Metrics
pos_pred = torch.cat(pos_pred, dim=0)
neg_pred = torch.cat(neg_pred, dim=0)
results = {}
for K in [10, 20, 30]:
evaluator.K = K
hits = evaluator.eval({
'y_pred_pos': pos_pred,
'y_pred_neg': neg_pred,
})[f'hits@{K}']
results[f'Hits@{K}'] = hits
print()
print(f'Validation Statistics')
print('*' * 30)
for k, v in results.items():
print(f'{k}: {v}')
writer.add_scalar(
f"Val/{k}", v, (pos_pred.shape[0] + neg_pred.shape[0]) * e)
print('*' * 30)
del pos_pred
del neg_pred
# Save model if it's the best one yet.
if results['Hits@20'] > best_val_hits:
best_val_hits = results['Hits@20']
filename = f'{args.save_path}/{args.experiment}/{model.__class__.__name__}_best_val.checkpoint'
model_utils.save_model(model, filename)
print(f'Model saved!')
print(f'Best validation Hits@20 yet: {best_val_hits}')
# Save model on checkpoints.
if e % args.checkpoint_freq == 0:
filename = f'{args.save_path}/{args.experiment}/{model.__class__.__name__}_epoch_{e}.checkpoint'
model_utils.save_model(model, filename)
print(f'Model checkpoint reached!')
saved_checkpoints.append(filename)
# Delete checkpoints if there are too many
while len(saved_checkpoints) > args.num_checkpoints:
os.remove(saved_checkpoints.pop(0))
return model
