chore: import upstream snapshot with attribution

examples/pytorch/node2vec/main.py

@@ -0,0 +1,62 @@
+import time
+
+from dgl.sampling import node2vec_random_walk
+
+from model import Node2vecModel
+from utils import load_graph, parse_arguments
+
+
+def time_randomwalk(graph, args):
+    """
+    Test cost time of random walk
+    """
+
+    start_time = time.time()
+
+    # default setting for testing
+    params = {"p": 0.25, "q": 4, "walk_length": 50}
+
+    for i in range(args.runs):
+        node2vec_random_walk(graph, graph.nodes(), **params)
+    end_time = time.time()
+    cost_time_avg = (end_time - start_time) / args.runs
+    print(
+        "Run dataset {} {} trials, mean run time: {:.3f}s".format(
+            args.dataset, args.runs, cost_time_avg
+        )
+    )
+
+
+def train_node2vec(graph, eval_set, args):
+    """
+    Train node2vec model
+    """
+    trainer = Node2vecModel(
+        graph,
+        embedding_dim=args.embedding_dim,
+        walk_length=args.walk_length,
+        p=args.p,
+        q=args.q,
+        num_walks=args.num_walks,
+        eval_set=eval_set,
+        eval_steps=1,
+        device=args.device,
+    )
+
+    trainer.train(
+        epochs=args.epochs, batch_size=args.batch_size, learning_rate=0.01
+    )
+
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    graph, eval_set = load_graph(args.dataset)
+
+    if args.task == "train":
+        print("Perform training node2vec model")
+        train_node2vec(graph, eval_set, args)
+    elif args.task == "time":
+        print("Timing random walks")
+        time_randomwalk(graph, args)
+    else:
+        raise ValueError("Task type error!")