wehub-resource-sync
65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
import argparse
|
|
|
|
from ogb.linkproppred import *
|
|
from ogb.nodeproppred import *
|
|
|
|
from dgl.data import CitationGraphDataset
|
|
|
|
|
|
def load_graph(name):
|
|
cite_graphs = ["cora", "citeseer", "pubmed"]
|
|
|
|
if name in cite_graphs:
|
|
dataset = CitationGraphDataset(name)
|
|
graph = dataset[0]
|
|
|
|
nodes = graph.nodes()
|
|
y = graph.ndata["label"]
|
|
train_mask = graph.ndata["train_mask"]
|
|
val_mask = graph.ndata["test_mask"]
|
|
|
|
nodes_train, y_train = nodes[train_mask], y[train_mask]
|
|
nodes_val, y_val = nodes[val_mask], y[val_mask]
|
|
eval_set = [(nodes_train, y_train), (nodes_val, y_val)]
|
|
|
|
elif name.startswith("ogbn"):
|
|
dataset = DglNodePropPredDataset(name)
|
|
graph, y = dataset[0]
|
|
split_nodes = dataset.get_idx_split()
|
|
nodes = graph.nodes()
|
|
|
|
train_idx = split_nodes["train"]
|
|
val_idx = split_nodes["valid"]
|
|
|
|
nodes_train, y_train = nodes[train_idx], y[train_idx]
|
|
nodes_val, y_val = nodes[val_idx], y[val_idx]
|
|
eval_set = [(nodes_train, y_train), (nodes_val, y_val)]
|
|
|
|
else:
|
|
raise ValueError("Dataset name error!")
|
|
|
|
return graph, eval_set
|
|
|
|
|
|
def parse_arguments():
|
|
"""
|
|
Parse arguments
|
|
"""
|
|
parser = argparse.ArgumentParser(description="Node2vec")
|
|
parser.add_argument("--dataset", type=str, default="cora")
|
|
# 'train' for training node2vec model, 'time' for testing speed of random walk
|
|
parser.add_argument("--task", type=str, default="train")
|
|
parser.add_argument("--runs", type=int, default=10)
|
|
parser.add_argument("--device", type=str, default="cpu")
|
|
parser.add_argument("--embedding_dim", type=int, default=128)
|
|
parser.add_argument("--walk_length", type=int, default=50)
|
|
parser.add_argument("--p", type=float, default=0.25)
|
|
parser.add_argument("--q", type=float, default=4.0)
|
|
parser.add_argument("--num_walks", type=int, default=10)
|
|
parser.add_argument("--epochs", type=int, default=100)
|
|
parser.add_argument("--batch_size", type=int, default=128)
|
|
|
|
args = parser.parse_args()
|
|
|
|
return args
|