[Example]Add sparse embedding in GATNE-T example to support big graph (#2234)

* add sparse and mutil gpu support

* update multi gpu

* update readme

* remove multi gpu

* add multiple gpus support
Co-authored-by: Quan (Andy) Gan <coin2028@hotmail.com>
Co-authored-by: Zihao Ye <expye@outlook.com>

examples/pytorch/GATNE-T/README.md

@@ -37,15 +37,28 @@ python src/main.py --input data/example
 
 To run on "twitter" dataset, use
 ```bash
-python src/main.py --input data/twitter --eval-type 1
+python src/main.py --input data/twitter --eval-type 1 --gpu 0
 ```
 
+For a big dataset, use sparse to avoid cuda out of memory in backward
+```bash
+python src/main_sparse.py --input data/example --gpu 0
+```
+
+If you have multiple GPUs, you can also accelerate training with [`DistributedDataParallel`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html)
+```bash
+python src/main_sparse_multi_gpus.py --input data/example --gpu 0,1
+```
+
+**It is worth noting that DistributedDataParallel will cause more cuda memory consumption and a certain loss of preformance.**
+
+
 Results
 -------
 All the results match the [official code](https://github.com/THUDM/GATNE/blob/master/src/main_pytorch.py) with the same hyper parameter values, including twiiter dataset (auc, pr, f1 is 76.29, 76.17, 69.34, respectively).
 
 |         | auc   | pr    | f1    |
-|  ------  |  ----  |  ---  | ----- |
+| ------- | ----- | ----- | ----- |
 | amazon  | 96.88 | 96.31 | 92.12 |
 | youtube | 82.29 | 80.35 | 74.63 |
 | twitter | 72.40 | 74.40 | 65.89 |

examples/pytorch/GATNE-T/scripts/run_example.sh

-python src/main.py --input data/example
\ No newline at end of file
+python src/main.py --input data/example --gpu 0

examples/pytorch/GATNE-T/scripts/run_example_sparse.sh

+python src/main_sparse.py --input data/example --gpu 0

examples/pytorch/GATNE-T/scripts/run_example_sparse_multi_gpus.sh

+python src/main_sparse_multi_gpus.py --input data/example

examples/pytorch/GATNE-T/src/main.py

@@ -8,7 +8,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import tqdm
+from tqdm.auto import tqdm
 from numpy import random
 from torch.nn.parameter import Parameter
 import dgl
@@ -18,7 +18,7 @@ from utils import *
 
 
 def get_graph(network_data, vocab):
-    ''' Build graph, treat all nodes as the same type
+    """ Build graph, treat all nodes as the same type
 
     Parameters
     ----------
@@ -30,10 +30,10 @@ def get_graph(network_data, vocab):
     ------
     DGLHeteroGraph
         a heterogenous graph, with one node type and different edge types
-    '''
+    """
     graphs = []
 
-    node_type = '_N' # '_N' can be replaced by an arbitrary name
+    node_type = "_N"  # '_N' can be replaced by an arbitrary name
     data_dict = dict()
     num_nodes_dict = {node_type: len(vocab)}
 
@@ -64,11 +64,24 @@ class NeighborSampler(object):
             sampled_block = dgl.to_block(sampled_graph, seeds)
             seeds = sampled_block.srcdata[dgl.NID]
             blocks.insert(0, sampled_block)
-        return blocks, torch.LongTensor(head_invmap), torch.LongTensor(tails), torch.LongTensor(types)
+        return (
+            blocks,
+            torch.LongTensor(head_invmap),
+            torch.LongTensor(tails),
+            torch.LongTensor(types),
+        )
 
 
 class DGLGATNE(nn.Module):
-    def __init__(self, num_nodes, embedding_size, embedding_u_size, edge_types, edge_type_count, dim_a):
+    def __init__(
+        self,
+        num_nodes,
+        embedding_size,
+        embedding_u_size,
+        edge_types,
+        edge_type_count,
+        dim_a,
+    ):
         super(DGLGATNE, self).__init__()
         self.num_nodes = num_nodes
         self.embedding_size = embedding_size
@@ -78,9 +91,15 @@ class DGLGATNE(nn.Module):
         self.dim_a = dim_a
 
         self.node_embeddings = Parameter(torch.FloatTensor(num_nodes, embedding_size))
-        self.node_type_embeddings = Parameter(torch.FloatTensor(num_nodes, edge_type_count, embedding_u_size))
-        self.trans_weights = Parameter(torch.FloatTensor(edge_type_count, embedding_u_size, embedding_size))
-        self.trans_weights_s1 = Parameter(torch.FloatTensor(edge_type_count, embedding_u_size, dim_a))
+        self.node_type_embeddings = Parameter(
+            torch.FloatTensor(num_nodes, edge_type_count, embedding_u_size)
+        )
+        self.trans_weights = Parameter(
+            torch.FloatTensor(edge_type_count, embedding_u_size, embedding_size)
+        )
+        self.trans_weights_s1 = Parameter(
+            torch.FloatTensor(edge_type_count, embedding_u_size, dim_a)
+        )
         self.trans_weights_s2 = Parameter(torch.FloatTensor(edge_type_count, dim_a, 1))
 
         self.reset_parameters()
@@ -105,29 +124,53 @@ class DGLGATNE(nn.Module):
                 edge_type = self.edge_types[i]
                 block.srcdata[edge_type] = self.node_type_embeddings[input_nodes, i]
                 block.dstdata[edge_type] = self.node_type_embeddings[output_nodes, i]
-                block.update_all(fn.copy_u(edge_type, 'm'), fn.sum('m', edge_type), etype=edge_type)
+                block.update_all(
+                    fn.copy_u(edge_type, "m"), fn.sum("m", edge_type), etype=edge_type
+                )
                 node_type_embed.append(block.dstdata[edge_type])
 
             node_type_embed = torch.stack(node_type_embed, 1)
-            tmp_node_type_embed = node_type_embed.unsqueeze(2).view(-1, 1, self.embedding_u_size)  
-            trans_w = self.trans_weights.unsqueeze(0).repeat(batch_size, 1, 1, 1).view(
-                -1, self.embedding_u_size, self.embedding_size
+            tmp_node_type_embed = node_type_embed.unsqueeze(2).view(
+                -1, 1, self.embedding_u_size
             )
-            trans_w_s1 = self.trans_weights_s1.unsqueeze(0).repeat(batch_size, 1, 1, 1).view(
-                -1, self.embedding_u_size, self.dim_a
+            trans_w = (
+                self.trans_weights.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.embedding_u_size, self.embedding_size)
+            )
+            trans_w_s1 = (
+                self.trans_weights_s1.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.embedding_u_size, self.dim_a)
+            )
+            trans_w_s2 = (
+                self.trans_weights_s2.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.dim_a, 1)
             )
-            trans_w_s2 = self.trans_weights_s2.unsqueeze(0).repeat(batch_size, 1, 1, 1).view(-1, self.dim_a, 1)
 
-            attention = F.softmax(
+            attention = (
+                F.softmax(
                     torch.matmul(
-                    torch.tanh(torch.matmul(tmp_node_type_embed, trans_w_s1)), trans_w_s2
-                ).squeeze(2).view(-1, self.edge_type_count),
+                        torch.tanh(torch.matmul(tmp_node_type_embed, trans_w_s1)),
+                        trans_w_s2,
+                    )
+                    .squeeze(2)
+                    .view(-1, self.edge_type_count),
                     dim=1,
-            ).unsqueeze(1).repeat(1, self.edge_type_count, 1)  
+                )
+                .unsqueeze(1)
+                .repeat(1, self.edge_type_count, 1)
+            )
 
-            node_type_embed = torch.matmul(attention, node_type_embed).view(-1, 1, self.embedding_u_size) 
-            node_embed = node_embed[output_nodes].unsqueeze(1).repeat(1, self.edge_type_count, 1) + \
-                torch.matmul(node_type_embed, trans_w).view(-1, self.edge_type_count, self.embedding_size)
+            node_type_embed = torch.matmul(attention, node_type_embed).view(
+                -1, 1, self.embedding_u_size
+            )
+            node_embed = node_embed[output_nodes].unsqueeze(1).repeat(
+                1, self.edge_type_count, 1
+            ) + torch.matmul(node_type_embed, trans_w).view(
+                -1, self.edge_type_count, self.embedding_size
+            )
             last_node_embed = F.normalize(node_embed, dim=2)
 
             return last_node_embed  # [batch_size, edge_type_count, embedding_size]
@@ -188,27 +231,41 @@ def train_model(network_data):
     dim_a = args.att_dim
     att_head = 1
     neighbor_samples = args.neighbor_samples
+    num_workers = args.workers
 
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
+    device = torch.device(
+        "cuda" if args.gpu is not None and torch.cuda.is_available() else "cpu"
+    )
 
     g = get_graph(network_data, vocab)
     all_walks = []
     for i in range(edge_type_count):
         nodes = torch.LongTensor(type_nodes[i] * args.num_walks)
-        traces, types = dgl.sampling.random_walk(g, nodes, metapath=[edge_types[i]]*(neighbor_samples-1))
+        traces, types = dgl.sampling.random_walk(
+            g, nodes, metapath=[edge_types[i]] * (neighbor_samples - 1)
+        )
         all_walks.append(traces)
 
-    train_pairs = generate_pairs(all_walks, args.window_size) 
+    train_pairs = generate_pairs(all_walks, args.window_size, num_workers)
     neighbor_sampler = NeighborSampler(g, [neighbor_samples])
     train_dataloader = torch.utils.data.DataLoader(
-        train_pairs, batch_size=batch_size, collate_fn=neighbor_sampler.sample, shuffle=True
+        train_pairs,
+        batch_size=batch_size,
+        collate_fn=neighbor_sampler.sample,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+    model = DGLGATNE(
+        num_nodes, embedding_size, embedding_u_size, edge_types, edge_type_count, dim_a
     )
-    model = DGLGATNE(num_nodes, embedding_size, embedding_u_size, edge_types, edge_type_count, dim_a)
     nsloss = NSLoss(num_nodes, num_sampled, embedding_size)
     model.to(device)
     nsloss.to(device)
 
-    optimizer = torch.optim.Adam([{"params": model.parameters()}, {"params": nsloss.parameters()}], lr=1e-4)
+    optimizer = torch.optim.Adam(
+        [{"params": model.parameters()}, {"params": nsloss.parameters()}], lr=1e-3
+    )
 
     best_score = 0
     patience = 0
@@ -216,11 +273,10 @@ def train_model(network_data):
         model.train()
         random.shuffle(train_pairs)
 
-        data_iter = tqdm.tqdm(
+        data_iter = tqdm(
             train_dataloader,
             desc="epoch %d" % (epoch),
             total=(len(train_pairs) + (batch_size - 1)) // batch_size,
-            bar_format="{l_bar}{r_bar}",
         )
         avg_loss = 0.0
 
@@ -229,33 +285,54 @@ def train_model(network_data):
             # embs: [batch_size, edge_type_count, embedding_size]
             block_types = block_types.to(device)
             embs = model(block[0].to(device))[head_invmap]
-            embs = embs.gather(1, block_types.view(-1, 1, 1).expand(embs.shape[0], 1, embs.shape[2]))[:, 0]
-            loss = nsloss(block[0].dstdata[dgl.NID][head_invmap].to(device), embs, tails.to(device))
+            embs = embs.gather(
+                1, block_types.view(-1, 1, 1).expand(embs.shape[0], 1, embs.shape[2])
+            )[:, 0]
+            loss = nsloss(
+                block[0].dstdata[dgl.NID][head_invmap].to(device),
+                embs,
+                tails.to(device),
+            )
             loss.backward()
             optimizer.step()
             avg_loss += loss.item()
 
-            if i % 5000 == 0:
             post_fix = {
                 "epoch": epoch,
                 "iter": i,
                 "avg_loss": avg_loss / (i + 1),
                 "loss": loss.item(),
             }
-                data_iter.write(str(post_fix))
+            data_iter.set_postfix(post_fix)
 
         model.eval()
         # {'1': {}, '2': {}}
         final_model = dict(zip(edge_types, [dict() for _ in range(edge_type_count)]))
         for i in range(num_nodes):
-            train_inputs = torch.tensor([i for _ in range(edge_type_count)]).unsqueeze(1).to(device)  # [i, i]
-            train_types = torch.tensor(list(range(edge_type_count))).unsqueeze(1).to(device)  # [0, 1]
-            pairs = torch.cat((train_inputs, train_inputs, train_types), dim=1)  # (2, 3)
-            train_blocks, train_invmap, fake_tails, train_types = neighbor_sampler.sample(pairs)
+            train_inputs = (
+                torch.tensor([i for _ in range(edge_type_count)])
+                .unsqueeze(1)
+                .to(device)
+            )  # [i, i]
+            train_types = (
+                torch.tensor(list(range(edge_type_count))).unsqueeze(1).to(device)
+            )  # [0, 1]
+            pairs = torch.cat(
+                (train_inputs, train_inputs, train_types), dim=1
+            )  # (2, 3)
+            (
+                train_blocks,
+                train_invmap,
+                fake_tails,
+                train_types,
+            ) = neighbor_sampler.sample(pairs)
 
             node_emb = model(train_blocks[0].to(device))[train_invmap]
-            node_emb = node_emb.gather(1, train_types.to(device).view(-1, 1, 1
-                ).expand(node_emb.shape[0], 1, node_emb.shape[2])
+            node_emb = node_emb.gather(
+                1,
+                train_types.to(device)
+                .view(-1, 1, 1)
+                .expand(node_emb.shape[0], 1, node_emb.shape[2]),
             )[:, 0]
 
             for j in range(edge_type_count):
@@ -271,6 +348,7 @@ def train_model(network_data):
                     final_model[edge_types[i]],
                     valid_true_data_by_edge[edge_types[i]],
                     valid_false_data_by_edge[edge_types[i]],
+                    num_workers,
                 )
                 valid_aucs.append(tmp_auc)
                 valid_f1s.append(tmp_f1)
@@ -280,6 +358,7 @@ def train_model(network_data):
                     final_model[edge_types[i]],
                     testing_true_data_by_edge[edge_types[i]],
                     testing_false_data_by_edge[edge_types[i]],
+                    num_workers,
                 )
                 test_aucs.append(tmp_auc)
                 test_f1s.append(tmp_f1)
@@ -323,4 +402,4 @@ if __name__ == "__main__":
     print("Overall ROC-AUC:", average_auc)
     print("Overall PR-AUC", average_pr)
     print("Overall F1:", average_f1)
-    print("Training Time", end-start)
+    print("Training Time", end - start)

examples/pytorch/GATNE-T/src/main_sparse.py

+from collections import defaultdict
+import math
+import os
+import sys
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import tqdm
+from numpy import random
+from torch.nn.parameter import Parameter
+import dgl
+import dgl.function as fn
+
+from utils import *
+
+
+def get_graph(network_data, vocab):
+    """ Build graph, treat all nodes as the same type
+
+    Parameters
+    ----------
+    network_data: a dict
+        keys describing the edge types, values representing edges
+    vocab: a dict
+        mapping node IDs to node indices
+    Output
+    ------
+    DGLHeteroGraph
+        a heterogenous graph, with one node type and different edge types
+    """
+    graphs = []
+
+    node_type = "_N"  # '_N' can be replaced by an arbitrary name
+    data_dict = dict()
+    num_nodes_dict = {node_type: len(vocab)}
+
+    for edge_type in network_data:
+        tmp_data = network_data[edge_type]
+        src = []
+        dst = []
+        for edge in tmp_data:
+            src.extend([vocab[edge[0]], vocab[edge[1]]])
+            dst.extend([vocab[edge[1]], vocab[edge[0]]])
+        data_dict[(node_type, edge_type, node_type)] = (src, dst)
+    graph = dgl.heterograph(data_dict, num_nodes_dict)
+
+    return graph
+
+
+class NeighborSampler(object):
+    def __init__(self, g, num_fanouts):
+        self.g = g
+        self.num_fanouts = num_fanouts
+
+    def sample(self, pairs):
+        pairs = np.stack(pairs)
+        heads, tails, types = pairs[:, 0], pairs[:, 1], pairs[:, 2]
+        seeds, head_invmap = torch.unique(torch.LongTensor(heads), return_inverse=True)
+        blocks = []
+        for fanout in reversed(self.num_fanouts):
+            sampled_graph = dgl.sampling.sample_neighbors(self.g, seeds, fanout)
+            sampled_block = dgl.to_block(sampled_graph, seeds)
+            seeds = sampled_block.srcdata[dgl.NID]
+            blocks.insert(0, sampled_block)
+        return (
+            blocks,
+            torch.LongTensor(head_invmap),
+            torch.LongTensor(tails),
+            torch.LongTensor(types),
+        )
+
+
+class DGLGATNE(nn.Module):
+    def __init__(
+        self,
+        num_nodes,
+        embedding_size,
+        embedding_u_size,
+        edge_types,
+        edge_type_count,
+        dim_a,
+    ):
+        super(DGLGATNE, self).__init__()
+        self.num_nodes = num_nodes
+        self.embedding_size = embedding_size
+        self.embedding_u_size = embedding_u_size
+        self.edge_types = edge_types
+        self.edge_type_count = edge_type_count
+        self.dim_a = dim_a
+
+        self.node_embeddings = nn.Embedding(num_nodes, embedding_size, sparse=True)
+        self.node_type_embeddings = nn.Embedding(
+            num_nodes * edge_type_count, embedding_u_size, sparse=True
+        )
+        self.trans_weights = Parameter(
+            torch.FloatTensor(edge_type_count, embedding_u_size, embedding_size)
+        )
+        self.trans_weights_s1 = Parameter(
+            torch.FloatTensor(edge_type_count, embedding_u_size, dim_a)
+        )
+        self.trans_weights_s2 = Parameter(torch.FloatTensor(edge_type_count, dim_a, 1))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.node_embeddings.weight.data.uniform_(-1.0, 1.0)
+        self.node_type_embeddings.weight.data.uniform_(-1.0, 1.0)
+        self.trans_weights.data.normal_(std=1.0 / math.sqrt(self.embedding_size))
+        self.trans_weights_s1.data.normal_(std=1.0 / math.sqrt(self.embedding_size))
+        self.trans_weights_s2.data.normal_(std=1.0 / math.sqrt(self.embedding_size))
+
+    # embs: [batch_size, embedding_size]
+    def forward(self, block):
+        input_nodes = block.srcdata[dgl.NID]
+        output_nodes = block.dstdata[dgl.NID]
+        batch_size = block.number_of_dst_nodes()
+        node_type_embed = []
+
+        with block.local_scope():
+            for i in range(self.edge_type_count):
+                edge_type = self.edge_types[i]
+                block.srcdata[edge_type] = self.node_type_embeddings(
+                    input_nodes * self.edge_type_count + i
+                )
+                block.dstdata[edge_type] = self.node_type_embeddings(
+                    output_nodes * self.edge_type_count + i
+                )
+                block.update_all(
+                    fn.copy_u(edge_type, "m"), fn.sum("m", edge_type), etype=edge_type
+                )
+                node_type_embed.append(block.dstdata[edge_type])
+
+            node_type_embed = torch.stack(node_type_embed, 1)
+            tmp_node_type_embed = node_type_embed.unsqueeze(2).view(
+                -1, 1, self.embedding_u_size
+            )
+            trans_w = (
+                self.trans_weights.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.embedding_u_size, self.embedding_size)
+            )
+            trans_w_s1 = (
+                self.trans_weights_s1.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.embedding_u_size, self.dim_a)
+            )
+            trans_w_s2 = (
+                self.trans_weights_s2.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.dim_a, 1)
+            )
+
+            attention = (
+                F.softmax(
+                    torch.matmul(
+                        torch.tanh(torch.matmul(tmp_node_type_embed, trans_w_s1)),
+                        trans_w_s2,
+                    )
+                    .squeeze(2)
+                    .view(-1, self.edge_type_count),
+                    dim=1,
+                )
+                .unsqueeze(1)
+                .repeat(1, self.edge_type_count, 1)
+            )
+
+            node_type_embed = torch.matmul(attention, node_type_embed).view(
+                -1, 1, self.embedding_u_size
+            )
+            node_embed = self.node_embeddings(output_nodes).unsqueeze(1).repeat(
+                1, self.edge_type_count, 1
+            ) + torch.matmul(node_type_embed, trans_w).view(
+                -1, self.edge_type_count, self.embedding_size
+            )
+            last_node_embed = F.normalize(node_embed, dim=2)
+
+            return last_node_embed  # [batch_size, edge_type_count, embedding_size]
+
+
+class NSLoss(nn.Module):
+    def __init__(self, num_nodes, num_sampled, embedding_size):
+        super(NSLoss, self).__init__()
+        self.num_nodes = num_nodes
+        self.num_sampled = num_sampled
+        self.embedding_size = embedding_size
+
+        # [ (log(i+2) - log(i+1)) / log(num_nodes + 1)]
+        self.sample_weights = F.normalize(
+            torch.Tensor(
+                [
+                    (math.log(k + 2) - math.log(k + 1)) / math.log(num_nodes + 1)
+                    for k in range(num_nodes)
+                ]
+            ),
+            dim=0,
+        )
+        self.weights = nn.Embedding(num_nodes, embedding_size, sparse=True)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.weights.weight.data.normal_(std=1.0 / math.sqrt(self.embedding_size))
+
+    def forward(self, input, embs, label):
+        n = input.shape[0]
+        log_target = torch.log(
+            torch.sigmoid(torch.sum(torch.mul(embs, self.weights(label)), 1))
+        )
+        negs = (
+            torch.multinomial(
+                self.sample_weights, self.num_sampled * n, replacement=True
+            )
+            .view(n, self.num_sampled)
+            .to(input.device)
+        )
+        noise = torch.neg(self.weights(negs))
+        sum_log_sampled = torch.sum(
+            torch.log(torch.sigmoid(torch.bmm(noise, embs.unsqueeze(2)))), 1
+        ).squeeze()
+
+        loss = log_target + sum_log_sampled
+        return -loss.sum() / n
+
+
+def train_model(network_data):
+    index2word, vocab, type_nodes = generate_vocab(network_data)
+
+    edge_types = list(network_data.keys())
+    num_nodes = len(index2word)
+    edge_type_count = len(edge_types)
+    epochs = args.epoch
+    batch_size = args.batch_size
+    embedding_size = args.dimensions
+    embedding_u_size = args.edge_dim
+    u_num = edge_type_count
+    num_sampled = args.negative_samples
+    dim_a = args.att_dim
+    att_head = 1
+    neighbor_samples = args.neighbor_samples
+    num_workers = args.workers
+
+    device = torch.device(
+        "cuda" if args.gpu is not None and torch.cuda.is_available() else "cpu"
+    )
+
+    g = get_graph(network_data, vocab)
+    all_walks = []
+    for i in range(edge_type_count):
+        nodes = torch.LongTensor(type_nodes[i] * args.num_walks)
+        traces, types = dgl.sampling.random_walk(
+            g, nodes, metapath=[edge_types[i]] * (neighbor_samples - 1)
+        )
+        all_walks.append(traces)
+
+    train_pairs = generate_pairs(all_walks, args.window_size, num_workers)
+    neighbor_sampler = NeighborSampler(g, [neighbor_samples])
+    train_dataloader = torch.utils.data.DataLoader(
+        train_pairs,
+        batch_size=batch_size,
+        collate_fn=neighbor_sampler.sample,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+
+    model = DGLGATNE(
+        num_nodes, embedding_size, embedding_u_size, edge_types, edge_type_count, dim_a,
+    )
+
+    nsloss = NSLoss(num_nodes, num_sampled, embedding_size)
+
+    model.to(device)
+    nsloss.to(device)
+
+    embeddings_params = list(map(id, model.node_embeddings.parameters())) + list(
+        map(id, model.node_type_embeddings.parameters())
+    )
+    weights_params = list(map(id, nsloss.weights.parameters()))
+
+    optimizer = torch.optim.Adam(
+        [
+            {
+                "params": filter(
+                    lambda p: id(p) not in embeddings_params, model.parameters(),
+                )
+            },
+            {
+                "params": filter(
+                    lambda p: id(p) not in weights_params, nsloss.parameters(),
+                )
+            },
+        ],
+        lr=1e-3,
+    )
+
+    sparse_optimizer = torch.optim.SparseAdam(
+        [
+            {"params": model.node_embeddings.parameters()},
+            {"params": model.node_type_embeddings.parameters()},
+            {"params": nsloss.weights.parameters()},
+        ],
+        lr=1e-3,
+    )
+
+    best_score = 0
+    patience = 0
+    for epoch in range(epochs):
+        model.train()
+
+        random.shuffle(train_pairs)
+
+        data_iter = tqdm.tqdm(
+            train_dataloader,
+            desc="epoch %d" % (epoch),
+            total=(len(train_pairs) + (batch_size - 1)) // batch_size,
+        )
+        avg_loss = 0.0
+
+        for i, (block, head_invmap, tails, block_types) in enumerate(data_iter):
+            optimizer.zero_grad()
+            sparse_optimizer.zero_grad()
+            # embs: [batch_size, edge_type_count, embedding_size]
+            block_types = block_types.to(device)
+            embs = model(block[0].to(device))[head_invmap]
+            embs = embs.gather(
+                1, block_types.view(-1, 1, 1).expand(embs.shape[0], 1, embs.shape[2]),
+            )[:, 0]
+            loss = nsloss(
+                block[0].dstdata[dgl.NID][head_invmap].to(device),
+                embs,
+                tails.to(device),
+            )
+            loss.backward()
+            optimizer.step()
+            sparse_optimizer.step()
+            avg_loss += loss.item()
+
+            post_fix = {
+                "epoch": epoch,
+                "iter": i,
+                "avg_loss": avg_loss / (i + 1),
+                "loss": loss.item(),
+            }
+            data_iter.set_postfix(post_fix)
+
+        model.eval()
+        # {'1': {}, '2': {}}
+        final_model = dict(zip(edge_types, [dict() for _ in range(edge_type_count)]))
+        for i in range(num_nodes):
+            train_inputs = (
+                torch.tensor([i for _ in range(edge_type_count)])
+                .unsqueeze(1)
+                .to(device)
+            )  # [i, i]
+            train_types = (
+                torch.tensor(list(range(edge_type_count))).unsqueeze(1).to(device)
+            )  # [0, 1]
+            pairs = torch.cat(
+                (train_inputs, train_inputs, train_types), dim=1
+            )  # (2, 3)
+            (
+                train_blocks,
+                train_invmap,
+                fake_tails,
+                train_types,
+            ) = neighbor_sampler.sample(pairs.cpu())
+
+            node_emb = model(train_blocks[0].to(device))[train_invmap]
+            node_emb = node_emb.gather(
+                1,
+                train_types.to(device)
+                .view(-1, 1, 1)
+                .expand(node_emb.shape[0], 1, node_emb.shape[2]),
+            )[:, 0]
+
+            for j in range(edge_type_count):
+                final_model[edge_types[j]][index2word[i]] = (
+                    node_emb[j].cpu().detach().numpy()
+                )
+
+        valid_aucs, valid_f1s, valid_prs = [], [], []
+        test_aucs, test_f1s, test_prs = [], [], []
+        for i in range(edge_type_count):
+            if args.eval_type == "all" or edge_types[i] in args.eval_type.split(","):
+                tmp_auc, tmp_f1, tmp_pr = evaluate(
+                    final_model[edge_types[i]],
+                    valid_true_data_by_edge[edge_types[i]],
+                    valid_false_data_by_edge[edge_types[i]],
+                    num_workers,
+                )
+                valid_aucs.append(tmp_auc)
+                valid_f1s.append(tmp_f1)
+                valid_prs.append(tmp_pr)
+
+                tmp_auc, tmp_f1, tmp_pr = evaluate(
+                    final_model[edge_types[i]],
+                    testing_true_data_by_edge[edge_types[i]],
+                    testing_false_data_by_edge[edge_types[i]],
+                    num_workers,
+                )
+                test_aucs.append(tmp_auc)
+                test_f1s.append(tmp_f1)
+                test_prs.append(tmp_pr)
+        print("valid auc:", np.mean(valid_aucs))
+        print("valid pr:", np.mean(valid_prs))
+        print("valid f1:", np.mean(valid_f1s))
+
+        average_auc = np.mean(test_aucs)
+        average_f1 = np.mean(test_f1s)
+        average_pr = np.mean(test_prs)
+
+        cur_score = np.mean(valid_aucs)
+        if cur_score > best_score:
+            best_score = cur_score
+            patience = 0
+        else:
+            patience += 1
+            if patience > args.patience:
+                print("Early Stopping")
+                break
+    return average_auc, average_f1, average_pr
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    file_name = args.input
+    print(args)
+
+    training_data_by_type = load_training_data(file_name + "/train.txt")
+    valid_true_data_by_edge, valid_false_data_by_edge = load_testing_data(
+        file_name + "/valid.txt"
+    )
+    testing_true_data_by_edge, testing_false_data_by_edge = load_testing_data(
+        file_name + "/test.txt"
+    )
+
+    start = time.time()
+    average_auc, average_f1, average_pr = train_model(training_data_by_type)
+    end = time.time()
+
+    print("Overall ROC-AUC:", average_auc)
+    print("Overall PR-AUC", average_pr)
+    print("Overall F1:", average_f1)
+    print("Training Time", end - start)

examples/pytorch/GATNE-T/src/main_sparse_multi_gpus.py

+from collections import defaultdict
+import math
+import os
+import sys
+import time
+import datetime
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.multiprocessing as mp
+from torch.nn.parallel import DistributedDataParallel
+from tqdm.auto import tqdm
+from numpy import random
+from torch.nn.parameter import Parameter
+import dgl
+import dgl.function as fn
+
+from utils import *
+
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def get_graph(network_data, vocab):
+    """ Build graph, treat all nodes as the same type
+
+    Parameters
+    ----------
+    network_data: a dict
+        keys describing the edge types, values representing edges
+    vocab: a dict
+        mapping node IDs to node indices
+    Output
+    ------
+    DGLHeteroGraph
+        a heterogenous graph, with one node type and different edge types
+    """
+    graphs = []
+
+    node_type = "_N"  # '_N' can be replaced by an arbitrary name
+    data_dict = dict()
+    num_nodes_dict = {node_type: len(vocab)}
+
+    for edge_type in network_data:
+        tmp_data = network_data[edge_type]
+        src = []
+        dst = []
+        for edge in tmp_data:
+            src.extend([vocab[edge[0]], vocab[edge[1]]])
+            dst.extend([vocab[edge[1]], vocab[edge[0]]])
+        data_dict[(node_type, edge_type, node_type)] = (src, dst)
+    graph = dgl.heterograph(data_dict, num_nodes_dict)
+
+    return graph
+
+
+class NeighborSampler(object):
+    def __init__(self, g, num_fanouts):
+        self.g = g
+        self.num_fanouts = num_fanouts
+
+    def sample(self, pairs):
+        pairs = np.stack(pairs)
+        heads, tails, types = pairs[:, 0], pairs[:, 1], pairs[:, 2]
+        seeds, head_invmap = torch.unique(torch.LongTensor(heads), return_inverse=True)
+        blocks = []
+        for fanout in reversed(self.num_fanouts):
+            sampled_graph = dgl.sampling.sample_neighbors(self.g, seeds, fanout)
+            sampled_block = dgl.to_block(sampled_graph, seeds)
+            seeds = sampled_block.srcdata[dgl.NID]
+            blocks.insert(0, sampled_block)
+        return (
+            blocks,
+            torch.LongTensor(head_invmap),
+            torch.LongTensor(tails),
+            torch.LongTensor(types),
+        )
+
+
+class DGLGATNE(nn.Module):
+    def __init__(
+        self,
+        num_nodes,
+        embedding_size,
+        embedding_u_size,
+        edge_types,
+        edge_type_count,
+        dim_a,
+    ):
+        super(DGLGATNE, self).__init__()
+        self.num_nodes = num_nodes
+        self.embedding_size = embedding_size
+        self.embedding_u_size = embedding_u_size
+        self.edge_types = edge_types
+        self.edge_type_count = edge_type_count
+        self.dim_a = dim_a
+
+        self.node_embeddings = nn.Embedding(num_nodes, embedding_size, sparse=True)
+        self.node_type_embeddings = nn.Embedding(
+            num_nodes * edge_type_count, embedding_u_size, sparse=True
+        )
+        self.trans_weights = Parameter(
+            torch.FloatTensor(edge_type_count, embedding_u_size, embedding_size)
+        )
+        self.trans_weights_s1 = Parameter(
+            torch.FloatTensor(edge_type_count, embedding_u_size, dim_a)
+        )
+        self.trans_weights_s2 = Parameter(torch.FloatTensor(edge_type_count, dim_a, 1))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.node_embeddings.weight.data.uniform_(-1.0, 1.0)
+        self.node_type_embeddings.weight.data.uniform_(-1.0, 1.0)
+        self.trans_weights.data.normal_(std=1.0 / math.sqrt(self.embedding_size))
+        self.trans_weights_s1.data.normal_(std=1.0 / math.sqrt(self.embedding_size))
+        self.trans_weights_s2.data.normal_(std=1.0 / math.sqrt(self.embedding_size))
+
+    # embs: [batch_size, embedding_size]
+    def forward(self, block):
+        input_nodes = block.srcdata[dgl.NID]
+        output_nodes = block.dstdata[dgl.NID]
+        batch_size = block.number_of_dst_nodes()
+        node_type_embed = []
+
+        with block.local_scope():
+            for i in range(self.edge_type_count):
+                edge_type = self.edge_types[i]
+                block.srcdata[edge_type] = self.node_type_embeddings(
+                    input_nodes * self.edge_type_count + i
+                )
+                block.dstdata[edge_type] = self.node_type_embeddings(
+                    output_nodes * self.edge_type_count + i
+                )
+                block.update_all(
+                    fn.copy_u(edge_type, "m"), fn.sum("m", edge_type), etype=edge_type
+                )
+                node_type_embed.append(block.dstdata[edge_type])
+
+            node_type_embed = torch.stack(node_type_embed, 1)
+            tmp_node_type_embed = node_type_embed.unsqueeze(2).view(
+                -1, 1, self.embedding_u_size
+            )
+            trans_w = (
+                self.trans_weights.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.embedding_u_size, self.embedding_size)
+            )
+            trans_w_s1 = (
+                self.trans_weights_s1.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.embedding_u_size, self.dim_a)
+            )
+            trans_w_s2 = (
+                self.trans_weights_s2.unsqueeze(0)
+                .repeat(batch_size, 1, 1, 1)
+                .view(-1, self.dim_a, 1)
+            )
+
+            attention = (
+                F.softmax(
+                    torch.matmul(
+                        torch.tanh(torch.matmul(tmp_node_type_embed, trans_w_s1)),
+                        trans_w_s2,
+                    )
+                    .squeeze(2)
+                    .view(-1, self.edge_type_count),
+                    dim=1,
+                )
+                .unsqueeze(1)
+                .repeat(1, self.edge_type_count, 1)
+            )
+
+            node_type_embed = torch.matmul(attention, node_type_embed).view(
+                -1, 1, self.embedding_u_size
+            )
+            node_embed = self.node_embeddings(output_nodes).unsqueeze(1).repeat(
+                1, self.edge_type_count, 1
+            ) + torch.matmul(node_type_embed, trans_w).view(
+                -1, self.edge_type_count, self.embedding_size
+            )
+            last_node_embed = F.normalize(node_embed, dim=2)
+
+            return last_node_embed  # [batch_size, edge_type_count, embedding_size]
+
+
+class NSLoss(nn.Module):
+    def __init__(self, num_nodes, num_sampled, embedding_size):
+        super(NSLoss, self).__init__()
+        self.num_nodes = num_nodes
+        self.num_sampled = num_sampled
+        self.embedding_size = embedding_size
+
+        # [ (log(i+2) - log(i+1)) / log(num_nodes + 1)]
+        self.sample_weights = F.normalize(
+            torch.Tensor(
+                [
+                    (math.log(k + 2) - math.log(k + 1)) / math.log(num_nodes + 1)
+                    for k in range(num_nodes)
+                ]
+            ),
+            dim=0,
+        )
+        self.weights = nn.Embedding(num_nodes, embedding_size, sparse=True)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.weights.weight.data.normal_(std=1.0 / math.sqrt(self.embedding_size))
+
+    def forward(self, input, embs, label):
+        n = input.shape[0]
+        log_target = torch.log(
+            torch.sigmoid(torch.sum(torch.mul(embs, self.weights(label)), 1))
+        )
+        negs = (
+            torch.multinomial(
+                self.sample_weights, self.num_sampled * n, replacement=True
+            )
+            .view(n, self.num_sampled)
+            .to(input.device)
+        )
+        noise = torch.neg(self.weights(negs))
+        sum_log_sampled = torch.sum(
+            torch.log(torch.sigmoid(torch.bmm(noise, embs.unsqueeze(2)))), 1
+        ).squeeze()
+
+        loss = log_target + sum_log_sampled
+        return -loss.sum() / n
+
+
+def run(proc_id, n_gpus, args, devices, data):
+    dev_id = devices[proc_id]
+    if n_gpus > 1:
+        dist_init_method = "tcp://{master_ip}:{master_port}".format(
+            master_ip="127.0.0.1", master_port="12345"
+        )
+        world_size = n_gpus
+        torch.distributed.init_process_group(
+            backend="gloo",
+            init_method=dist_init_method,
+            world_size=world_size,
+            rank=proc_id,
+            timeout=datetime.timedelta(seconds=100),
+        )
+    torch.cuda.set_device(dev_id)
+
+    g, train_pairs, index2word, edge_types, num_nodes, edge_type_count = data
+
+    epochs = args.epoch
+    batch_size = args.batch_size
+    embedding_size = args.dimensions
+    embedding_u_size = args.edge_dim
+    u_num = edge_type_count
+    num_sampled = args.negative_samples
+    dim_a = args.att_dim
+    att_head = 1
+    neighbor_samples = args.neighbor_samples
+    num_workers = args.workers
+
+    train_pairs = torch.split(
+        torch.tensor(train_pairs), math.ceil(len(train_pairs) / n_gpus)
+    )[proc_id]
+    neighbor_sampler = NeighborSampler(g, [neighbor_samples])
+    train_dataloader = torch.utils.data.DataLoader(
+        train_pairs,
+        batch_size=batch_size,
+        collate_fn=neighbor_sampler.sample,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+
+    model = DGLGATNE(
+        num_nodes, embedding_size, embedding_u_size, edge_types, edge_type_count, dim_a,
+    )
+
+    nsloss = NSLoss(num_nodes, num_sampled, embedding_size)
+
+    model.to(dev_id)
+    if n_gpus > 1:
+        model = DistributedDataParallel(
+            model, device_ids=[dev_id], output_device=dev_id
+        )
+
+    nsloss.to(dev_id)
+
+    if n_gpus > 1:
+        mmodel = model.module
+    else:
+        mmodel = model
+
+    embeddings_params = list(map(id, mmodel.node_embeddings.parameters())) + list(
+        map(id, mmodel.node_type_embeddings.parameters())
+    )
+    weights_params = list(map(id, nsloss.weights.parameters()))
+
+    optimizer = torch.optim.Adam(
+        [
+            {
+                "params": filter(
+                    lambda p: id(p) not in embeddings_params, model.parameters(),
+                )
+            },
+            {
+                "params": filter(
+                    lambda p: id(p) not in weights_params, nsloss.parameters(),
+                )
+            },
+        ],
+        lr=2e-3,
+    )
+
+    sparse_optimizer = torch.optim.SparseAdam(
+        [
+            {"params": mmodel.node_embeddings.parameters()},
+            {"params": mmodel.node_type_embeddings.parameters()},
+            {"params": nsloss.weights.parameters()},
+        ],
+        lr=2e-3,
+    )
+
+    if n_gpus > 1:
+        torch.distributed.barrier()
+
+    if proc_id == 0:
+        start = time.time()
+
+    for epoch in range(epochs):
+        model.train()
+
+        data_iter = train_dataloader
+        if proc_id == 0:
+            data_iter = tqdm(
+                train_dataloader,
+                desc="epoch %d" % (epoch),
+                total=(len(train_pairs) + (batch_size - 1)) // batch_size,
+            )
+            avg_loss = 0.0
+
+        for i, (block, head_invmap, tails, block_types) in enumerate(data_iter):
+            optimizer.zero_grad()
+            sparse_optimizer.zero_grad()
+            # embs: [batch_size, edge_type_count, embedding_size]
+            block_types = block_types.to(dev_id)
+            embs = model(block[0].to(dev_id))[head_invmap]
+            embs = embs.gather(
+                1, block_types.view(-1, 1, 1).expand(embs.shape[0], 1, embs.shape[2]),
+            )[:, 0]
+            loss = nsloss(
+                block[0].dstdata[dgl.NID][head_invmap].to(dev_id),
+                embs,
+                tails.to(dev_id),
+            )
+            loss.backward()
+            optimizer.step()
+            sparse_optimizer.step()
+
+            if proc_id == 0:
+                avg_loss += loss.item()
+
+                post_fix = {
+                    "avg_loss": avg_loss / (i + 1),
+                    "loss": loss.item(),
+                }
+                data_iter.set_postfix(post_fix)
+
+        if n_gpus > 1:
+            torch.distributed.barrier()
+
+        if proc_id == 0:
+            model.eval()
+            # {'1': {}, '2': {}}
+            final_model = dict(
+                zip(edge_types, [dict() for _ in range(edge_type_count)])
+            )
+            for i in range(num_nodes):
+                train_inputs = (
+                    torch.tensor([i for _ in range(edge_type_count)])
+                    .unsqueeze(1)
+                    .to(dev_id)
+                )  # [i, i]
+                train_types = (
+                    torch.tensor(list(range(edge_type_count))).unsqueeze(1).to(dev_id)
+                )  # [0, 1]
+                pairs = torch.cat(
+                    (train_inputs, train_inputs, train_types), dim=1
+                )  # (2, 3)
+                (
+                    train_blocks,
+                    train_invmap,
+                    fake_tails,
+                    train_types,
+                ) = neighbor_sampler.sample(pairs.cpu())
+
+                node_emb = model(train_blocks[0].to(dev_id))[train_invmap]
+                node_emb = node_emb.gather(
+                    1,
+                    train_types.to(dev_id)
+                    .view(-1, 1, 1)
+                    .expand(node_emb.shape[0], 1, node_emb.shape[2]),
+                )[:, 0]
+
+                for j in range(edge_type_count):
+                    final_model[edge_types[j]][index2word[i]] = (
+                        node_emb[j].cpu().detach().numpy()
+                    )
+
+            valid_aucs, valid_f1s, valid_prs = [], [], []
+            test_aucs, test_f1s, test_prs = [], [], []
+            for i in range(edge_type_count):
+                if args.eval_type == "all" or edge_types[i] in args.eval_type.split(
+                    ","
+                ):
+                    tmp_auc, tmp_f1, tmp_pr = evaluate(
+                        final_model[edge_types[i]],
+                        valid_true_data_by_edge[edge_types[i]],
+                        valid_false_data_by_edge[edge_types[i]],
+                        num_workers,
+                    )
+                    valid_aucs.append(tmp_auc)
+                    valid_f1s.append(tmp_f1)
+                    valid_prs.append(tmp_pr)
+
+                    tmp_auc, tmp_f1, tmp_pr = evaluate(
+                        final_model[edge_types[i]],
+                        testing_true_data_by_edge[edge_types[i]],
+                        testing_false_data_by_edge[edge_types[i]],
+                        num_workers,
+                    )
+                    test_aucs.append(tmp_auc)
+                    test_f1s.append(tmp_f1)
+                    test_prs.append(tmp_pr)
+            print("valid auc:", np.mean(valid_aucs))
+            print("valid pr:", np.mean(valid_prs))
+            print("valid f1:", np.mean(valid_f1s))
+
+    if proc_id == 0:
+        end = time.time()
+        average_auc = np.mean(test_aucs)
+        average_f1 = np.mean(test_f1s)
+        average_pr = np.mean(test_prs)
+        print("Overall ROC-AUC:", average_auc)
+        print("Overall PR-AUC", average_pr)
+        print("Overall F1:", average_f1)
+        print("Training Time", end - start)
+
+
+def train_model(network_data):
+    index2word, vocab, type_nodes = generate_vocab(network_data)
+
+    edge_types = list(network_data.keys())
+    num_nodes = len(index2word)
+    edge_type_count = len(edge_types)
+
+    devices = list(map(int, args.gpu.split(",")))
+    n_gpus = len(devices)
+    neighbor_samples = args.neighbor_samples
+    num_workers = args.workers
+
+    g = get_graph(network_data, vocab)
+    all_walks = []
+    for i in range(edge_type_count):
+        nodes = torch.LongTensor(type_nodes[i] * args.num_walks)
+        traces, types = dgl.sampling.random_walk(
+            g, nodes, metapath=[edge_types[i]] * (neighbor_samples - 1)
+        )
+        all_walks.append(traces)
+
+    train_pairs = generate_pairs(all_walks, args.window_size, num_workers)
+    data = g, train_pairs, index2word, edge_types, num_nodes, edge_type_count
+
+    if n_gpus == 1:
+        run(0, n_gpus, args, devices, data)
+    else:
+        procs = []
+        for proc_id in range(n_gpus):
+            p = mp.Process(
+                target=thread_wrapped_func(run),
+                args=(proc_id, n_gpus, args, devices, data),
+            )
+            p.start()
+            procs.append(p)
+        for p in procs:
+            p.join()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    file_name = args.input
+    print(args)
+    setup_seed(1234)
+
+    training_data_by_type = load_training_data(file_name + "/train.txt")
+    valid_true_data_by_edge, valid_false_data_by_edge = load_testing_data(
+        file_name + "/valid.txt"
+    )
+    testing_true_data_by_edge, testing_false_data_by_edge = load_testing_data(
+        file_name + "/test.txt"
+    )
+
+    train_model(training_data_by_type)

examples/pytorch/GATNE-T/src/utils.py

@@ -5,69 +5,159 @@ import networkx as nx
 import numpy as np
 from gensim.models.keyedvectors import Vocab
 from six import iteritems
-from sklearn.metrics import (auc, f1_score, precision_recall_curve,
-                             roc_auc_score)
+from sklearn.metrics import auc, f1_score, precision_recall_curve, roc_auc_score
 import torch
 
-def parse_args():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument('--input', type=str, default='data/amazon',
-                        help='Input dataset path')
-    
-    parser.add_argument('--features', type=str, default=None,
-                        help='Input node features')
-
-    parser.add_argument('--epoch', type=int, default=100,
-                        help='Number of epoch. Default is 100.')
-
-    parser.add_argument('--batch-size', type=int, default=64,
-                        help='Number of batch_size. Default is 64.')
-
-    parser.add_argument('--eval-type', type=str, default='all',
-                        help='The edge type(s) for evaluation.')
+import time
+import multiprocessing
+from functools import partial, reduce, wraps
 
-    parser.add_argument('--schema', type=str, default=None,
-                        help='The metapath schema (e.g., U-I-U,I-U-I).')
+import torch.multiprocessing as mp
+from _thread import start_new_thread
+import traceback
 
-    parser.add_argument('--dimensions', type=int, default=200,
-                        help='Number of dimensions. Default is 200.')
 
-    parser.add_argument('--edge-dim', type=int, default=10,
-                        help='Number of edge embedding dimensions. Default is 10.')
+def thread_wrapped_func(func):
+    """
+    Wraps a process entry point to make it work with OpenMP.
+    """
 
-    parser.add_argument('--att-dim', type=int, default=20,
-                        help='Number of attention dimensions. Default is 20.')
+    @wraps(func)
+    def decorated_function(*args, **kwargs):
+        queue = mp.Queue()
 
-    parser.add_argument('--walk-length', type=int, default=10,
-                        help='Length of walk per source. Default is 10.')
-
-    parser.add_argument('--num-walks', type=int, default=20,
-                        help='Number of walks per source. Default is 20.')
+        def _queue_result():
+            exception, trace, res = None, None, None
+            try:
+                res = func(*args, **kwargs)
+            except Exception as e:
+                exception = e
+                trace = traceback.format_exc()
+            queue.put((res, exception, trace))
+
+        start_new_thread(_queue_result, ())
+        result, exception, trace = queue.get()
+        if exception is None:
+            return result
+        else:
+            assert isinstance(exception, Exception)
+            raise exception.__class__(trace)
 
-    parser.add_argument('--window-size', type=int, default=5,
-                        help='Context size for optimization. Default is 5.')
+    return decorated_function
 
-    parser.add_argument('--negative-samples', type=int, default=5,
-                        help='Negative samples for optimization. Default is 5.')
 
-    parser.add_argument('--neighbor-samples', type=int, default=10,
-                        help='Neighbor samples for aggregation. Default is 10.')
+def parse_args():
+    parser = argparse.ArgumentParser()
 
-    parser.add_argument('--patience', type=int, default=5,
-                        help='Early stopping patience. Default is 5.')
+    parser.add_argument(
+        "--input", type=str, default="data/amazon", help="Input dataset path"
+    )
+
+    parser.add_argument(
+        "--features", type=str, default=None, help="Input node features"
+    )
+
+    parser.add_argument(
+        "--epoch", type=int, default=100, help="Number of epoch. Default is 100."
+    )
+
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=64,
+        help="Number of batch_size. Default is 64.",
+    )
+
+    parser.add_argument(
+        "--eval-type", type=str, default="all", help="The edge type(s) for evaluation."
+    )
+
+    parser.add_argument(
+        "--schema",
+        type=str,
+        default=None,
+        help="The metapath schema (e.g., U-I-U,I-U-I).",
+    )
+
+    parser.add_argument(
+        "--dimensions",
+        type=int,
+        default=200,
+        help="Number of dimensions. Default is 200.",
+    )
+
+    parser.add_argument(
+        "--edge-dim",
+        type=int,
+        default=10,
+        help="Number of edge embedding dimensions. Default is 10.",
+    )
+
+    parser.add_argument(
+        "--att-dim",
+        type=int,
+        default=20,
+        help="Number of attention dimensions. Default is 20.",
+    )
+
+    parser.add_argument(
+        "--walk-length",
+        type=int,
+        default=10,
+        help="Length of walk per source. Default is 10.",
+    )
+
+    parser.add_argument(
+        "--num-walks",
+        type=int,
+        default=20,
+        help="Number of walks per source. Default is 20.",
+    )
+
+    parser.add_argument(
+        "--window-size",
+        type=int,
+        default=5,
+        help="Context size for optimization. Default is 5.",
+    )
+
+    parser.add_argument(
+        "--negative-samples",
+        type=int,
+        default=5,
+        help="Negative samples for optimization. Default is 5.",
+    )
+
+    parser.add_argument(
+        "--neighbor-samples",
+        type=int,
+        default=10,
+        help="Neighbor samples for aggregation. Default is 10.",
+    )
+
+    parser.add_argument(
+        "--patience", type=int, default=5, help="Early stopping patience. Default is 5."
+    )
+
+    parser.add_argument(
+        "--gpu", type=str, default=None, help="Comma separated list of GPU device IDs."
+    )
+
+    parser.add_argument(
+        "--workers", type=int, default=4, help="Number of workers.",
+    )
 
     return parser.parse_args()
 
 
 # for each line, the data is [edge_type, node, node]
 def load_training_data(f_name):
-    print('We are loading data from:', f_name)
+    print("We are loading data from:", f_name)
     edge_data_by_type = dict()
     all_nodes = list()
-    with open(f_name, 'r') as f:
+    with open(f_name, "r") as f:
         for line in f:
-            words = line[:-1].split(' ')  # line[-1] == '\n'
+            words = line[:-1].split(" ")  # line[-1] == '\n'
             if words[0] not in edge_data_by_type:
                 edge_data_by_type[words[0]] = list()
             x, y = words[1], words[2]
@@ -75,20 +165,20 @@ def load_training_data(f_name):
             all_nodes.append(x)
             all_nodes.append(y)
     all_nodes = list(set(all_nodes))
-    print('Total training nodes: ' + str(len(all_nodes)))
+    print("Total training nodes: " + str(len(all_nodes)))
     return edge_data_by_type
 
 
 # for each line, the data is [edge_type, node, node, true_or_false]
 def load_testing_data(f_name):
-    print('We are loading data from:', f_name)
+    print("We are loading data from:", f_name)
     true_edge_data_by_type = dict()
     false_edge_data_by_type = dict()
     all_edges = list()
     all_nodes = list()
-    with open(f_name, 'r') as f:
+    with open(f_name, "r") as f:
         for line in f:
-            words = line[:-1].split(' ')
+            words = line[:-1].split(" ")
             x, y = words[1], words[2]
             if int(words[3]) == 1:
                 if words[0] not in true_edge_data_by_type:
@@ -105,29 +195,61 @@ def load_testing_data(f_name):
 
 
 def load_node_type(f_name):
-    print('We are loading node type from:', f_name)
+    print("We are loading node type from:", f_name)
     node_type = {}
-    with open(f_name, 'r') as f:
+    with open(f_name, "r") as f:
         for line in f:
             items = line.strip().split()
             node_type[items[0]] = items[1]
     return node_type
 
 
-def generate_pairs(all_walks, window_size):
-    # for each node, choose the first neighbor and second neighbor of it to form pairs
+def generate_pairs_parallel(walks, skip_window=None, layer_id=None):
     pairs = []
-    skip_window = window_size // 2
-    for layer_id, walks in enumerate(all_walks):
     for walk in walks:
+        walk = walk.tolist()
         for i in range(len(walk)):
             for j in range(1, skip_window + 1):
                 if i - j >= 0:
-                        pairs.append((walk[i], walk[i-j], layer_id))
+                    pairs.append((walk[i], walk[i - j], layer_id))
                 if i + j < len(walk):
-                        pairs.append((walk[i], walk[i+j], layer_id))
+                    pairs.append((walk[i], walk[i + j], layer_id))
     return pairs
 
+
+def generate_pairs(all_walks, window_size, num_workers):
+    # for each node, choose the first neighbor and second neighbor of it to form pairs
+    # Get all worker processes
+    start_time = time.time()
+    print("We are generating pairs with {} cores.".format(num_workers))
+
+    # Start all worker processes
+    pool = multiprocessing.Pool(processes=num_workers)
+    pairs = []
+    skip_window = window_size // 2
+    for layer_id, walks in enumerate(all_walks):
+        block_num = len(walks) // num_workers
+        if block_num > 0:
+            walks_list = [
+                walks[i * block_num : min((i + 1) * block_num, len(walks))]
+                for i in range(num_workers)
+            ]
+        else:
+            walks_list = [walks]
+        tmp_result = pool.map(
+            partial(
+                generate_pairs_parallel, skip_window=skip_window, layer_id=layer_id
+            ),
+            walks_list,
+        )
+        pairs += reduce(lambda x, y: x + y, tmp_result)
+
+    pool.close()
+    end_time = time.time()
+    print("Generate pairs end, use {}s.".format(end_time - start_time))
+    return np.array([list(pair) for pair in set(pairs)])
+
+
 def generate_vocab(network_data):
     nodes, index2word = [], []
     for edge_type in network_data:
@@ -150,31 +272,39 @@ def generate_vocab(network_data):
     return index2word, vocab, nodes
 
 
-def get_score(local_model, node1, node2):
+def get_score(local_model, edge):
+    node1, node2 = str(edge[0]), str(edge[1])
     try:
         vector1 = local_model[node1]
         vector2 = local_model[node2]
-        return np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))
+        return np.dot(vector1, vector2) / (
+            np.linalg.norm(vector1) * np.linalg.norm(vector2)
+        )
     except Exception as e:
         pass
 
 
-def evaluate(model, true_edges, false_edges):
+def evaluate(model, true_edges, false_edges, num_workers):
     true_list = list()
     prediction_list = list()
     true_num = 0
-    for edge in true_edges:
-        tmp_score = get_score(model, str(edge[0]), str(edge[1]))
-        if tmp_score is not None:
-            true_list.append(1)
-            prediction_list.append(tmp_score)
-            true_num += 1
-
-    for edge in false_edges:
-        tmp_score = get_score(model, str(edge[0]), str(edge[1]))
-        if tmp_score is not None:
-            true_list.append(0)
-            prediction_list.append(tmp_score)
+
+    # Start all worker processes
+    pool = multiprocessing.Pool(processes=num_workers)
+    tmp_true_score_list = pool.map(partial(get_score, model), true_edges)
+    tmp_false_score_list = pool.map(partial(get_score, model), false_edges)
+    pool.close()
+
+    prediction_list += [
+        tmp_score for tmp_score in tmp_true_score_list if tmp_score is not None
+    ]
+    true_num = len(prediction_list)
+    true_list += [1] * true_num
+
+    prediction_list += [
+        tmp_score for tmp_score in tmp_false_score_list if tmp_score is not None
+    ]
+    true_list += [0] * (len(prediction_list) - true_num)
 
     sorted_pred = prediction_list[:]
     sorted_pred.sort()