chore: import upstream snapshot with attribution

examples/pytorch/GATNE-T/README.md

@@ -0,0 +1,67 @@
+Representation Learning for Attributed Multiplex Heterogeneous Network (GANTE)
+============
+
+- Paper link: [https://arxiv.org/abs/1905.01669](https://arxiv.org/abs/1905.01669)
+- Author's code repo: [https://github.com/THUDM/GATNE](https://github.com/THUDM/GATNE). Note that only GATNE-T is implemented here.
+
+Requirements
+------------
+- requirements
+
+```bash
+pip install -r requirements.txt
+```
+
+Also requires PyTorch 1.7.0+.
+
+Datasets
+--------
+
+To prepare the datasets:
+1. ```bash
+   mkdir data
+   cd data
+   ```
+2. Download datasets from the following links:
+    - example: https://s3.us-west-2.amazonaws.com/dgl-data/dataset/recsys/GATNE/example.zip
+    - amazon: https://s3.us-west-2.amazonaws.com/dgl-data/dataset/recsys/GATNE/amazon.zip
+    - youtube: https://s3.us-west-2.amazonaws.com/dgl-data/dataset/recsys/GATNE/youtube.zip
+    - twitter: https://s3.us-west-2.amazonaws.com/dgl-data/dataset/recsys/GATNE/twitter.zip
+3. Unzip the datasets
+
+Training
+--------
+
+Run with following (available dataset: "example", "youtube", "amazon")
+```bash
+python src/main.py --input data/example
+```
+
+To run on "twitter" dataset, use
+```bash
+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 |
+| example | 94.65 | 94.57 | 89.99 |

examples/pytorch/GATNE-T/requirements.txt

@@ -0,0 +1,7 @@
+tqdm
+numpy
+scikit-learn
+networkx
+gensim
+requests
+--pre dgl-cu101

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

@@ -0,0 +1 @@
+python src/main.py --input data/example --gpu 0

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

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

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

@@ -0,0 +1 @@
+python src/main_sparse_multi_gpus.py --input data/example

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

@@ -0,0 +1,443 @@
+import math
+import os
+import sys
+import time
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from numpy import random
+from torch.nn.parameter import Parameter
+from tqdm.auto import tqdm
+from utils import *
+
+import dgl
+import dgl.function as fn
+
+
+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
+    ------
+    DGLGraph
+        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):
+        heads, tails, types = zip(*pairs)
+        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 = 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.trans_weights_s2 = Parameter(
+            torch.FloatTensor(edge_type_count, dim_a, 1)
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.node_embeddings.data.uniform_(-1.0, 1.0)
+        self.node_type_embeddings.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_embed = self.node_embeddings
+        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, 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,
+                )
+                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 = 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]
+            )
+
+
+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
+        self.weights = Parameter(torch.FloatTensor(num_nodes, 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.reset_parameters()
+
+    def reset_parameters(self):
+        self.weights.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)
+        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)
+
+    optimizer = torch.optim.Adam(
+        [{"params": model.parameters()}, {"params": nsloss.parameters()}],
+        lr=1e-3,
+    )
+
+    best_score = 0
+    patience = 0
+    for epoch in range(epochs):
+        model.train()
+        random.shuffle(train_pairs)
+
+        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()
+            # 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()
+            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)
+
+            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.py

@@ -0,0 +1,483 @@
+import math
+import os
+import sys
+import time
+from collections import defaultdict
+
+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
+from utils import *
+
+import dgl
+import dgl.function as fn
+
+
+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
+    ------
+    DGLGraph
+        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

@@ -0,0 +1,549 @@
+import datetime
+import math
+import os
+import sys
+import time
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+import torch.nn.functional as F
+from numpy import random
+from torch.nn.parallel import DistributedDataParallel
+from torch.nn.parameter import Parameter
+from tqdm.auto import tqdm
+from utils import *
+
+import dgl
+import dgl.function as fn
+
+
+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
+    ------
+    DGLGraph
+        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
+
+    neighbor_sampler = NeighborSampler(g, [neighbor_samples])
+    if n_gpus > 1:
+        train_sampler = torch.utils.data.distributed.DistributedSampler(
+            train_pairs,
+            num_replicas=world_size,
+            rank=proc_id,
+            shuffle=True,
+            drop_last=False,
+        )
+        train_dataloader = torch.utils.data.DataLoader(
+            train_pairs,
+            batch_size=batch_size,
+            collate_fn=neighbor_sampler.sample,
+            num_workers=num_workers,
+            sampler=train_sampler,
+            pin_memory=True,
+        )
+    else:
+        train_dataloader = torch.utils.data.DataLoader(
+            train_pairs,
+            batch_size=batch_size,
+            collate_fn=neighbor_sampler.sample,
+            num_workers=num_workers,
+            shuffle=True,
+            drop_last=False,
+            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):
+        if n_gpus > 1:
+            train_sampler.set_epoch(epoch)
+        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:
+        mp.spawn(run, args=(n_gpus, args, devices, data), nprocs=n_gpus)
+
+
+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

@@ -0,0 +1,308 @@
+import argparse
+import multiprocessing
+import time
+from collections import defaultdict
+from functools import partial, reduce, wraps
+
+import networkx as nx
+import numpy as np
+import torch
+from gensim.models.keyedvectors import Vocab
+from six import iteritems
+from sklearn.metrics import auc, f1_score, precision_recall_curve, roc_auc_score
+
+
+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.",
+    )
+
+    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)
+    edge_data_by_type = dict()
+    all_nodes = list()
+    with open(f_name, "r") as f:
+        for line in f:
+            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]
+            edge_data_by_type[words[0]].append((x, y))
+            all_nodes.append(x)
+            all_nodes.append(y)
+    all_nodes = list(set(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)
+    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:
+        for line in f:
+            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:
+                    true_edge_data_by_type[words[0]] = list()
+                true_edge_data_by_type[words[0]].append((x, y))
+            else:
+                if words[0] not in false_edge_data_by_type:
+                    false_edge_data_by_type[words[0]] = list()
+                false_edge_data_by_type[words[0]].append((x, y))
+            all_nodes.append(x)
+            all_nodes.append(y)
+    all_nodes = list(set(all_nodes))
+    return true_edge_data_by_type, false_edge_data_by_type
+
+
+def load_node_type(f_name):
+    print("We are loading node type from:", f_name)
+    node_type = {}
+    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_parallel(walks, skip_window=None, layer_id=None):
+    pairs = []
+    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))
+                if i + j < len(walk):
+                    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:
+        node1, node2 = zip(*network_data[edge_type])
+        index2word = index2word + list(node1) + list(node2)
+
+    index2word = list(set(index2word))
+    vocab = {}
+    i = 0
+    for word in index2word:
+        vocab[word] = i
+        i = i + 1
+
+    for edge_type in network_data:
+        node1, node2 = zip(*network_data[edge_type])
+        tmp_nodes = list(set(list(node1) + list(node2)))
+        tmp_nodes = [vocab[word] for word in tmp_nodes]
+        nodes.append(tmp_nodes)
+
+    return index2word, vocab, nodes
+
+
+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)
+        )
+    except Exception as e:
+        pass
+
+
+def evaluate(model, true_edges, false_edges, num_workers):
+    true_list = list()
+    prediction_list = list()
+    true_num = 0
+
+    # 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()
+    threshold = sorted_pred[-true_num]
+
+    y_pred = np.zeros(len(prediction_list), dtype=np.int32)
+    for i in range(len(prediction_list)):
+        if prediction_list[i] >= threshold:
+            y_pred[i] = 1
+
+    y_true = np.array(true_list)
+    y_scores = np.array(prediction_list)
+    ps, rs, _ = precision_recall_curve(y_true, y_scores)
+    return (
+        roc_auc_score(y_true, y_scores),
+        f1_score(y_true, y_pred),
+        auc(rs, ps),
+    )

examples/pytorch/GNN-FiLM/README.md

@@ -0,0 +1,58 @@
+# DGL Implementation of the GNN-FiLM Model
+
+This DGL example implements the GNN model proposed in the paper [GNN-FiLM: Graph Neural Networks with Feature-wise Linear Modulation](https://arxiv.org/pdf/1906.12192.pdf). 
+The author's codes of implementation is in [here](https://github.com/Microsoft/tf-gnn-samples)
+
+
+Example implementor
+----------------------
+This example was implemented by [Kounianhua Du](https://github.com/KounianhuaDu) during her Software Dev Engineer Intern work at the AWS Shanghai AI Lab.
+
+
+Dependencies
+----------------------
+- numpy 1.19.4
+- scikit-learn 0.22.1
+- pytorch 1.4.0
+- dgl 0.5.3
+
+
+The graph dataset used in this example 
+---------------------------------------
+The DGL's built-in PPIDataset. This is a Protein-Protein Interaction dataset for inductive node classification. The PPIDataset is a toy Protein-Protein Interaction network dataset. The dataset contains 24 graphs. The average number of nodes per graph is 2372. Each node has 50 features and 121 labels. There are 20 graphs for training, 2 for validation, and 2 for testing.
+
+NOTE: Following the paper, in addition to the dataset-provided untyped edges, a fresh "self-loop" edge type is added.
+
+Statistics:
+- Train examples: 20
+- Valid examples: 2
+- Test examples: 2
+- AvgNodesPerGraph: 2372
+- NumFeats: 50
+- NumLabels: 121
+
+
+How to run example files
+--------------------------------
+In the GNNFiLM folder, run
+
+```bash
+python main.py 
+```
+
+If want to use a GPU, run
+
+```bash
+python main.py --gpu ${your_device_id_here}
+```
+
+
+Performance
+-------------------------
+
+NOTE: We do not perform grid search or finetune here, so there is a gap between the performance reported in the original paper and this example. Below results, mean(standard deviation), were computed over ten runs.
+
+**GNN-FiLM results on PPI task**
+| Model         | Paper (tensorflow)               | ours (dgl)                  |
+| ------------- | -------------------------------- | --------------------------- |
+| Avg. Micro-F1 | 0.992 (0.000)                    | 0.983 (0.001)               |

examples/pytorch/GNN-FiLM/data_loader.py

@@ -0,0 +1,98 @@
+import collections
+
+import dgl
+from dgl.data import PPIDataset
+
+from torch.utils.data import DataLoader, Dataset
+
+# implement the collate_fn for dgl graph data class
+PPIBatch = collections.namedtuple("PPIBatch", ["graph", "label"])
+
+
+def batcher(device):
+    def batcher_dev(batch):
+        batch_graphs = dgl.batch(batch)
+        return PPIBatch(
+            graph=batch_graphs, label=batch_graphs.ndata["label"].to(device)
+        )
+
+    return batcher_dev
+
+
+# add a fresh "self-loop" edge type to the untyped PPI dataset and prepare train, val, test loaders
+def load_PPI(batch_size=1, device="cpu"):
+    train_set = PPIDataset(mode="train")
+    valid_set = PPIDataset(mode="valid")
+    test_set = PPIDataset(mode="test")
+    # for each graph, add self-loops as a new relation type
+    # here we reconstruct the graph since the schema of a heterograph cannot be changed once constructed
+    for i in range(len(train_set)):
+        g = dgl.heterograph(
+            {
+                ("_N", "_E", "_N"): train_set[i].edges(),
+                ("_N", "self", "_N"): (
+                    train_set[i].nodes(),
+                    train_set[i].nodes(),
+                ),
+            }
+        )
+        g.ndata["label"] = train_set[i].ndata["label"]
+        g.ndata["feat"] = train_set[i].ndata["feat"]
+        g.ndata["_ID"] = train_set[i].ndata["_ID"]
+        g.edges["_E"].data["_ID"] = train_set[i].edata["_ID"]
+        train_set.graphs[i] = g
+    for i in range(len(valid_set)):
+        g = dgl.heterograph(
+            {
+                ("_N", "_E", "_N"): valid_set[i].edges(),
+                ("_N", "self", "_N"): (
+                    valid_set[i].nodes(),
+                    valid_set[i].nodes(),
+                ),
+            }
+        )
+        g.ndata["label"] = valid_set[i].ndata["label"]
+        g.ndata["feat"] = valid_set[i].ndata["feat"]
+        g.ndata["_ID"] = valid_set[i].ndata["_ID"]
+        g.edges["_E"].data["_ID"] = valid_set[i].edata["_ID"]
+        valid_set.graphs[i] = g
+    for i in range(len(test_set)):
+        g = dgl.heterograph(
+            {
+                ("_N", "_E", "_N"): test_set[i].edges(),
+                ("_N", "self", "_N"): (
+                    test_set[i].nodes(),
+                    test_set[i].nodes(),
+                ),
+            }
+        )
+        g.ndata["label"] = test_set[i].ndata["label"]
+        g.ndata["feat"] = test_set[i].ndata["feat"]
+        g.ndata["_ID"] = test_set[i].ndata["_ID"]
+        g.edges["_E"].data["_ID"] = test_set[i].edata["_ID"]
+        test_set.graphs[i] = g
+
+    etypes = train_set[0].etypes
+    in_size = train_set[0].ndata["feat"].shape[1]
+    out_size = train_set[0].ndata["label"].shape[1]
+
+    # prepare train, valid, and test dataloaders
+    train_loader = DataLoader(
+        train_set,
+        batch_size=batch_size,
+        collate_fn=batcher(device),
+        shuffle=True,
+    )
+    valid_loader = DataLoader(
+        valid_set,
+        batch_size=batch_size,
+        collate_fn=batcher(device),
+        shuffle=True,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=batch_size,
+        collate_fn=batcher(device),
+        shuffle=True,
+    )
+    return train_loader, valid_loader, test_loader, etypes, in_size, out_size

examples/pytorch/GNN-FiLM/main.py

@@ -0,0 +1,289 @@
+import argparse
+import os
+
+import dgl
+import dgl.function as fn
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from data_loader import load_PPI
+from utils import evaluate_f1_score
+
+
+class GNNFiLMLayer(nn.Module):
+    def __init__(self, in_size, out_size, etypes, dropout=0.1):
+        super(GNNFiLMLayer, self).__init__()
+        self.in_size = in_size
+        self.out_size = out_size
+
+        # weights for different types of edges
+        self.W = nn.ModuleDict(
+            {name: nn.Linear(in_size, out_size, bias=False) for name in etypes}
+        )
+
+        # hypernets to learn the affine functions for different types of edges
+        self.film = nn.ModuleDict(
+            {
+                name: nn.Linear(in_size, 2 * out_size, bias=False)
+                for name in etypes
+            }
+        )
+
+        # layernorm before each propogation
+        self.layernorm = nn.LayerNorm(out_size)
+
+        # dropout layer
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, g, feat_dict):
+        # the input graph is a multi-relational graph, so treated as hetero-graph.
+
+        funcs = {}  # message and reduce functions dict
+        # for each type of edges, compute messages and reduce them all
+        for srctype, etype, dsttype in g.canonical_etypes:
+            messages = self.W[etype](
+                feat_dict[srctype]
+            )  # apply W_l on src feature
+            film_weights = self.film[etype](
+                feat_dict[dsttype]
+            )  # use dst feature to compute affine function paras
+            gamma = film_weights[
+                :, : self.out_size
+            ]  # "gamma" for the affine function
+            beta = film_weights[
+                :, self.out_size :
+            ]  # "beta" for the affine function
+            messages = gamma * messages + beta  # compute messages
+            messages = F.relu_(messages)
+            g.nodes[srctype].data[etype] = messages  # store in ndata
+            funcs[etype] = (
+                fn.copy_u(etype, "m"),
+                fn.sum("m", "h"),
+            )  # define message and reduce functions
+        g.multi_update_all(
+            funcs, "sum"
+        )  # update all, reduce by first type-wisely then across different types
+        feat_dict = {}
+        for ntype in g.ntypes:
+            feat_dict[ntype] = self.dropout(
+                self.layernorm(g.nodes[ntype].data["h"])
+            )  # apply layernorm and dropout
+        return feat_dict
+
+
+class GNNFiLM(nn.Module):
+    def __init__(
+        self, etypes, in_size, hidden_size, out_size, num_layers, dropout=0.1
+    ):
+        super(GNNFiLM, self).__init__()
+        self.film_layers = nn.ModuleList()
+        self.film_layers.append(
+            GNNFiLMLayer(in_size, hidden_size, etypes, dropout)
+        )
+        for i in range(num_layers - 1):
+            self.film_layers.append(
+                GNNFiLMLayer(hidden_size, hidden_size, etypes, dropout)
+            )
+        self.predict = nn.Linear(hidden_size, out_size, bias=True)
+
+    def forward(self, g, out_key):
+        h_dict = {
+            ntype: g.nodes[ntype].data["feat"] for ntype in g.ntypes
+        }  # prepare input feature dict
+        for layer in self.film_layers:
+            h_dict = layer(g, h_dict)
+        h = self.predict(
+            h_dict[out_key]
+        )  # use the final embed to predict, out_size = num_classes
+        h = torch.sigmoid(h)
+        return h
+
+
+def main(args):
+    # Step 1: Prepare graph data and retrieve train/validation/test dataloader ============================= #
+    if args.gpu >= 0 and torch.cuda.is_available():
+        device = "cuda:{}".format(args.gpu)
+    else:
+        device = "cpu"
+
+    if args.dataset == "PPI":
+        train_set, valid_set, test_set, etypes, in_size, out_size = load_PPI(
+            args.batch_size, device
+        )
+
+    # Step 2: Create model and training components=========================================================== #
+    model = GNNFiLM(
+        etypes, in_size, args.hidden_size, out_size, args.num_layers
+    ).to(device)
+    criterion = nn.BCELoss()
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.wd
+    )
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer, args.step_size, gamma=args.gamma
+    )
+
+    # Step 4: training epoches ============================================================================== #
+    lastf1 = 0
+    cnt = 0
+    best_val_f1 = 0
+    for epoch in range(args.max_epoch):
+        train_loss = []
+        train_f1 = []
+        val_loss = []
+        val_f1 = []
+        model.train()
+        for batch in train_set:
+            g = batch.graph
+            g = g.to(device)
+            logits = model.forward(g, "_N")
+            labels = batch.label
+            loss = criterion(logits, labels)
+            f1 = evaluate_f1_score(
+                logits.detach().cpu().numpy(), labels.detach().cpu().numpy()
+            )
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            train_loss.append(loss.item())
+            train_f1.append(f1)
+
+        train_loss = np.mean(train_loss)
+        train_f1 = np.mean(train_f1)
+        scheduler.step()
+
+        model.eval()
+        with torch.no_grad():
+            for batch in valid_set:
+                g = batch.graph
+                g = g.to(device)
+                logits = model.forward(g, "_N")
+                labels = batch.label
+                loss = criterion(logits, labels)
+                f1 = evaluate_f1_score(
+                    logits.detach().cpu().numpy(), labels.detach().cpu().numpy()
+                )
+                val_loss.append(loss.item())
+                val_f1.append(f1)
+
+        val_loss = np.mean(val_loss)
+        val_f1 = np.mean(val_f1)
+        print(
+            "Epoch {:d} | Train Loss {:.4f} | Train F1 {:.4f} | Val Loss {:.4f} | Val F1 {:.4f} |".format(
+                epoch + 1, train_loss, train_f1, val_loss, val_f1
+            )
+        )
+        if val_f1 > best_val_f1:
+            best_val_f1 = val_f1
+            torch.save(
+                model.state_dict(), os.path.join(args.save_dir, args.name)
+            )
+
+        if val_f1 < lastf1:
+            cnt += 1
+            if cnt == args.early_stopping:
+                print("Early stop.")
+                break
+        else:
+            cnt = 0
+            lastf1 = val_f1
+
+    model.eval()
+    test_loss = []
+    test_f1 = []
+    model.load_state_dict(
+        torch.load(os.path.join(args.save_dir, args.name), weights_only=False)
+    )
+    with torch.no_grad():
+        for batch in test_set:
+            g = batch.graph
+            g = g.to(device)
+            logits = model.forward(g, "_N")
+            labels = batch.label
+            loss = criterion(logits, labels)
+            f1 = evaluate_f1_score(
+                logits.detach().cpu().numpy(), labels.detach().cpu().numpy()
+            )
+            test_loss.append(loss.item())
+            test_f1.append(f1)
+    test_loss = np.mean(test_loss)
+    test_f1 = np.mean(test_f1)
+
+    print("Test F1: {:.4f} | Test loss: {:.4f}".format(test_f1, test_loss))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="GNN-FiLM")
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="PPI",
+        help="DGL dataset for this GNN-FiLM",
+    )
+    parser.add_argument(
+        "--gpu", type=int, default=-1, help="GPU Index. Default: -1, using CPU."
+    )
+    parser.add_argument(
+        "--in_size", type=int, default=50, help="Input dimensionalities"
+    )
+    parser.add_argument(
+        "--hidden_size",
+        type=int,
+        default=320,
+        help="Hidden layer dimensionalities",
+    )
+    parser.add_argument(
+        "--out_size", type=int, default=121, help="Output dimensionalities"
+    )
+    parser.add_argument(
+        "--num_layers", type=int, default=4, help="Number of GNN layers"
+    )
+    parser.add_argument("--batch_size", type=int, default=5, help="Batch size")
+    parser.add_argument(
+        "--max_epoch",
+        type=int,
+        default=1500,
+        help="The max number of epoches. Default: 500",
+    )
+    parser.add_argument(
+        "--early_stopping",
+        type=int,
+        default=80,
+        help="Early stopping. Default: 50",
+    )
+    parser.add_argument(
+        "--lr", type=float, default=0.001, help="Learning rate. Default: 3e-1"
+    )
+    parser.add_argument(
+        "--wd", type=float, default=0.0009, help="Weight decay. Default: 3e-1"
+    )
+    parser.add_argument(
+        "--step-size",
+        type=int,
+        default=40,
+        help="Period of learning rate decay.",
+    )
+    parser.add_argument(
+        "--gamma",
+        type=float,
+        default=0.8,
+        help="Multiplicative factor of learning rate decay.",
+    )
+    parser.add_argument(
+        "--dropout", type=float, default=0.1, help="Dropout rate. Default: 0.9"
+    )
+    parser.add_argument(
+        "--save_dir", type=str, default="./out", help="Path to save the model."
+    )
+    parser.add_argument(
+        "--name", type=str, default="GNN-FiLM", help="Saved model name."
+    )
+
+    args = parser.parse_args()
+    print(args)
+    if not os.path.exists(args.save_dir):
+        os.mkdir(args.save_dir)
+    main(args)

examples/pytorch/GNN-FiLM/utils.py

@@ -0,0 +1,10 @@
+import numpy as np
+from sklearn.metrics import f1_score
+
+
+# function to compute f1 score
+def evaluate_f1_score(pred, label):
+    pred = np.round(pred, 0).astype(np.int16)
+    pred = pred.flatten()
+    label = label.flatten()
+    return f1_score(y_pred=pred, y_true=label)

examples/pytorch/NGCF/Data/load_amazon-book.sh

@@ -0,0 +1,2 @@
+wget https://s3.us-west-2.amazonaws.com/dgl-data/dataset/amazon-book.zip
+unzip amazon-book.zip

examples/pytorch/NGCF/Data/load_gowalla.sh

@@ -0,0 +1,2 @@
+wget https://s3.us-west-2.amazonaws.com/dgl-data/dataset/gowalla.zip
+unzip gowalla.zip

examples/pytorch/NGCF/NGCF/main.py

@@ -0,0 +1,147 @@
+import os
+from time import time
+
+import torch
+import torch.optim as optim
+from model import NGCF
+from utility.batch_test import *
+from utility.helper import early_stopping
+
+
+def main(args):
+    # Step 1: Prepare graph data and device ================================================================= #
+    if args.gpu >= 0 and torch.cuda.is_available():
+        device = "cuda:{}".format(args.gpu)
+    else:
+        device = "cpu"
+
+    g = data_generator.g
+    g = g.to(device)
+
+    # Step 2: Create model and training components=========================================================== #
+    model = NGCF(
+        g, args.embed_size, args.layer_size, args.mess_dropout, args.regs[0]
+    ).to(device)
+    optimizer = optim.Adam(model.parameters(), lr=args.lr)
+
+    # Step 3: training epoches ============================================================================== #
+    n_batch = data_generator.n_train // args.batch_size + 1
+    t0 = time()
+    cur_best_pre_0, stopping_step = 0, 0
+    loss_loger, pre_loger, rec_loger, ndcg_loger, hit_loger = [], [], [], [], []
+    for epoch in range(args.epoch):
+        t1 = time()
+        loss, mf_loss, emb_loss = 0.0, 0.0, 0.0
+        for idx in range(n_batch):
+            users, pos_items, neg_items = data_generator.sample()
+            u_g_embeddings, pos_i_g_embeddings, neg_i_g_embeddings = model(
+                g, "user", "item", users, pos_items, neg_items
+            )
+
+            batch_loss, batch_mf_loss, batch_emb_loss = model.create_bpr_loss(
+                u_g_embeddings, pos_i_g_embeddings, neg_i_g_embeddings
+            )
+            optimizer.zero_grad()
+            batch_loss.backward()
+            optimizer.step()
+
+            loss += batch_loss
+            mf_loss += batch_mf_loss
+            emb_loss += batch_emb_loss
+
+        if (epoch + 1) % 10 != 0:
+            if args.verbose > 0 and epoch % args.verbose == 0:
+                perf_str = "Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f]" % (
+                    epoch,
+                    time() - t1,
+                    loss,
+                    mf_loss,
+                    emb_loss,
+                )
+                print(perf_str)
+            continue  # end the current epoch and move to the next epoch, let the following evaluation run every 10 epoches
+
+        # evaluate the model every 10 epoches
+        t2 = time()
+        users_to_test = list(data_generator.test_set.keys())
+        ret = test(model, g, users_to_test)
+        t3 = time()
+
+        loss_loger.append(loss)
+        rec_loger.append(ret["recall"])
+        pre_loger.append(ret["precision"])
+        ndcg_loger.append(ret["ndcg"])
+        hit_loger.append(ret["hit_ratio"])
+
+        if args.verbose > 0:
+            perf_str = (
+                "Epoch %d [%.1fs + %.1fs]: train==[%.5f=%.5f + %.5f], recall=[%.5f, %.5f], "
+                "precision=[%.5f, %.5f], hit=[%.5f, %.5f], ndcg=[%.5f, %.5f]"
+                % (
+                    epoch,
+                    t2 - t1,
+                    t3 - t2,
+                    loss,
+                    mf_loss,
+                    emb_loss,
+                    ret["recall"][0],
+                    ret["recall"][-1],
+                    ret["precision"][0],
+                    ret["precision"][-1],
+                    ret["hit_ratio"][0],
+                    ret["hit_ratio"][-1],
+                    ret["ndcg"][0],
+                    ret["ndcg"][-1],
+                )
+            )
+            print(perf_str)
+
+        cur_best_pre_0, stopping_step, should_stop = early_stopping(
+            ret["recall"][0],
+            cur_best_pre_0,
+            stopping_step,
+            expected_order="acc",
+            flag_step=5,
+        )
+
+        # early stop
+        if should_stop == True:
+            break
+
+        if ret["recall"][0] == cur_best_pre_0 and args.save_flag == 1:
+            torch.save(model.state_dict(), args.weights_path + args.model_name)
+            print(
+                "save the weights in path: ",
+                args.weights_path + args.model_name,
+            )
+
+    recs = np.array(rec_loger)
+    pres = np.array(pre_loger)
+    ndcgs = np.array(ndcg_loger)
+    hit = np.array(hit_loger)
+
+    best_rec_0 = max(recs[:, 0])
+    idx = list(recs[:, 0]).index(best_rec_0)
+
+    final_perf = (
+        "Best Iter=[%d]@[%.1f]\trecall=[%s], precision=[%s], hit=[%s], ndcg=[%s]"
+        % (
+            idx,
+            time() - t0,
+            "\t".join(["%.5f" % r for r in recs[idx]]),
+            "\t".join(["%.5f" % r for r in pres[idx]]),
+            "\t".join(["%.5f" % r for r in hit[idx]]),
+            "\t".join(["%.5f" % r for r in ndcgs[idx]]),
+        )
+    )
+    print(final_perf)
+
+
+if __name__ == "__main__":
+    if not os.path.exists(args.weights_path):
+        os.mkdir(args.weights_path)
+    args.mess_dropout = eval(args.mess_dropout)
+    args.layer_size = eval(args.layer_size)
+    args.regs = eval(args.regs)
+    print(args)
+    main(args)

examples/pytorch/NGCF/NGCF/model.py

@@ -0,0 +1,151 @@
+import dgl.function as fn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class NGCFLayer(nn.Module):
+    def __init__(self, in_size, out_size, norm_dict, dropout):
+        super(NGCFLayer, self).__init__()
+        self.in_size = in_size
+        self.out_size = out_size
+
+        # weights for different types of messages
+        self.W1 = nn.Linear(in_size, out_size, bias=True)
+        self.W2 = nn.Linear(in_size, out_size, bias=True)
+
+        # leaky relu
+        self.leaky_relu = nn.LeakyReLU(0.2)
+
+        # dropout layer
+        self.dropout = nn.Dropout(dropout)
+
+        # initialization
+        torch.nn.init.xavier_uniform_(self.W1.weight)
+        torch.nn.init.constant_(self.W1.bias, 0)
+        torch.nn.init.xavier_uniform_(self.W2.weight)
+        torch.nn.init.constant_(self.W2.bias, 0)
+
+        # norm
+        self.norm_dict = norm_dict
+
+    def forward(self, g, feat_dict):
+        funcs = {}  # message and reduce functions dict
+        # for each type of edges, compute messages and reduce them all
+        for srctype, etype, dsttype in g.canonical_etypes:
+            if srctype == dsttype:  # for self loops
+                messages = self.W1(feat_dict[srctype])
+                g.nodes[srctype].data[etype] = messages  # store in ndata
+                funcs[(srctype, etype, dsttype)] = (
+                    fn.copy_u(etype, "m"),
+                    fn.sum("m", "h"),
+                )  # define message and reduce functions
+            else:
+                src, dst = g.edges(etype=(srctype, etype, dsttype))
+                norm = self.norm_dict[(srctype, etype, dsttype)]
+                messages = norm * (
+                    self.W1(feat_dict[srctype][src])
+                    + self.W2(feat_dict[srctype][src] * feat_dict[dsttype][dst])
+                )  # compute messages
+                g.edges[(srctype, etype, dsttype)].data[
+                    etype
+                ] = messages  # store in edata
+                funcs[(srctype, etype, dsttype)] = (
+                    fn.copy_e(etype, "m"),
+                    fn.sum("m", "h"),
+                )  # define message and reduce functions
+
+        g.multi_update_all(
+            funcs, "sum"
+        )  # update all, reduce by first type-wisely then across different types
+        feature_dict = {}
+        for ntype in g.ntypes:
+            h = self.leaky_relu(g.nodes[ntype].data["h"])  # leaky relu
+            h = self.dropout(h)  # dropout
+            h = F.normalize(h, dim=1, p=2)  # l2 normalize
+            feature_dict[ntype] = h
+        return feature_dict
+
+
+class NGCF(nn.Module):
+    def __init__(self, g, in_size, layer_size, dropout, lmbd=1e-5):
+        super(NGCF, self).__init__()
+        self.lmbd = lmbd
+        self.norm_dict = dict()
+        for srctype, etype, dsttype in g.canonical_etypes:
+            src, dst = g.edges(etype=(srctype, etype, dsttype))
+            dst_degree = g.in_degrees(
+                dst, etype=(srctype, etype, dsttype)
+            ).float()  # obtain degrees
+            src_degree = g.out_degrees(
+                src, etype=(srctype, etype, dsttype)
+            ).float()
+            norm = torch.pow(src_degree * dst_degree, -0.5).unsqueeze(
+                1
+            )  # compute norm
+            self.norm_dict[(srctype, etype, dsttype)] = norm
+
+        self.layers = nn.ModuleList()
+        self.layers.append(
+            NGCFLayer(in_size, layer_size[0], self.norm_dict, dropout[0])
+        )
+        self.num_layers = len(layer_size)
+        for i in range(self.num_layers - 1):
+            self.layers.append(
+                NGCFLayer(
+                    layer_size[i],
+                    layer_size[i + 1],
+                    self.norm_dict,
+                    dropout[i + 1],
+                )
+            )
+        self.initializer = nn.init.xavier_uniform_
+
+        # embeddings for different types of nodes
+        self.feature_dict = nn.ParameterDict(
+            {
+                ntype: nn.Parameter(
+                    self.initializer(torch.empty(g.num_nodes(ntype), in_size))
+                )
+                for ntype in g.ntypes
+            }
+        )
+
+    def create_bpr_loss(self, users, pos_items, neg_items):
+        pos_scores = (users * pos_items).sum(1)
+        neg_scores = (users * neg_items).sum(1)
+
+        mf_loss = nn.LogSigmoid()(pos_scores - neg_scores).mean()
+        mf_loss = -1 * mf_loss
+
+        regularizer = (
+            torch.norm(users) ** 2
+            + torch.norm(pos_items) ** 2
+            + torch.norm(neg_items) ** 2
+        ) / 2
+        emb_loss = self.lmbd * regularizer / users.shape[0]
+
+        return mf_loss + emb_loss, mf_loss, emb_loss
+
+    def rating(self, u_g_embeddings, pos_i_g_embeddings):
+        return torch.matmul(u_g_embeddings, pos_i_g_embeddings.t())
+
+    def forward(self, g, user_key, item_key, users, pos_items, neg_items):
+        h_dict = {ntype: self.feature_dict[ntype] for ntype in g.ntypes}
+        # obtain features of each layer and concatenate them all
+        user_embeds = []
+        item_embeds = []
+        user_embeds.append(h_dict[user_key])
+        item_embeds.append(h_dict[item_key])
+        for layer in self.layers:
+            h_dict = layer(g, h_dict)
+            user_embeds.append(h_dict[user_key])
+            item_embeds.append(h_dict[item_key])
+        user_embd = torch.cat(user_embeds, 1)
+        item_embd = torch.cat(item_embeds, 1)
+
+        u_g_embeddings = user_embd[users, :]
+        pos_i_g_embeddings = item_embd[pos_items, :]
+        neg_i_g_embeddings = item_embd[neg_items, :]
+
+        return u_g_embeddings, pos_i_g_embeddings, neg_i_g_embeddings

examples/pytorch/NGCF/NGCF/utility/batch_test.py

@@ -0,0 +1,195 @@
+# This file is based on the NGCF author's implementation
+# <https://github.com/xiangwang1223/neural_graph_collaborative_filtering/blob/master/NGCF/utility/batch_test.py>.
+# It implements the batch test.
+
+import heapq
+import multiprocessing
+
+import utility.metrics as metrics
+from utility.load_data import *
+from utility.parser import parse_args
+
+cores = multiprocessing.cpu_count()
+
+args = parse_args()
+Ks = eval(args.Ks)
+
+data_generator = Data(
+    path=args.data_path + args.dataset, batch_size=args.batch_size
+)
+USR_NUM, ITEM_NUM = data_generator.n_users, data_generator.n_items
+N_TRAIN, N_TEST = data_generator.n_train, data_generator.n_test
+BATCH_SIZE = args.batch_size
+
+
+def ranklist_by_heapq(user_pos_test, test_items, rating, Ks):
+    item_score = {}
+    for i in test_items:
+        item_score[i] = rating[i]
+
+    K_max = max(Ks)
+    K_max_item_score = heapq.nlargest(K_max, item_score, key=item_score.get)
+
+    r = []
+    for i in K_max_item_score:
+        if i in user_pos_test:
+            r.append(1)
+        else:
+            r.append(0)
+    auc = 0.0
+    return r, auc
+
+
+def get_auc(item_score, user_pos_test):
+    item_score = sorted(item_score.items(), key=lambda kv: kv[1])
+    item_score.reverse()
+    item_sort = [x[0] for x in item_score]
+    posterior = [x[1] for x in item_score]
+
+    r = []
+    for i in item_sort:
+        if i in user_pos_test:
+            r.append(1)
+        else:
+            r.append(0)
+    auc = metrics.auc(ground_truth=r, prediction=posterior)
+    return auc
+
+
+def ranklist_by_sorted(user_pos_test, test_items, rating, Ks):
+    item_score = {}
+    for i in test_items:
+        item_score[i] = rating[i]
+
+    K_max = max(Ks)
+    K_max_item_score = heapq.nlargest(K_max, item_score, key=item_score.get)
+
+    r = []
+    for i in K_max_item_score:
+        if i in user_pos_test:
+            r.append(1)
+        else:
+            r.append(0)
+    auc = get_auc(item_score, user_pos_test)
+    return r, auc
+
+
+def get_performance(user_pos_test, r, auc, Ks):
+    precision, recall, ndcg, hit_ratio = [], [], [], []
+
+    for K in Ks:
+        precision.append(metrics.precision_at_k(r, K))
+        recall.append(metrics.recall_at_k(r, K, len(user_pos_test)))
+        ndcg.append(metrics.ndcg_at_k(r, K))
+        hit_ratio.append(metrics.hit_at_k(r, K))
+
+    return {
+        "recall": np.array(recall),
+        "precision": np.array(precision),
+        "ndcg": np.array(ndcg),
+        "hit_ratio": np.array(hit_ratio),
+        "auc": auc,
+    }
+
+
+def test_one_user(x):
+    # user u's ratings for user u
+    rating = x[0]
+    # uid
+    u = x[1]
+    # user u's items in the training set
+    try:
+        training_items = data_generator.train_items[u]
+    except Exception:
+        training_items = []
+    # user u's items in the test set
+    user_pos_test = data_generator.test_set[u]
+
+    all_items = set(range(ITEM_NUM))
+
+    test_items = list(all_items - set(training_items))
+
+    if args.test_flag == "part":
+        r, auc = ranklist_by_heapq(user_pos_test, test_items, rating, Ks)
+    else:
+        r, auc = ranklist_by_sorted(user_pos_test, test_items, rating, Ks)
+
+    return get_performance(user_pos_test, r, auc, Ks)
+
+
+def test(model, g, users_to_test, batch_test_flag=False):
+    result = {
+        "precision": np.zeros(len(Ks)),
+        "recall": np.zeros(len(Ks)),
+        "ndcg": np.zeros(len(Ks)),
+        "hit_ratio": np.zeros(len(Ks)),
+        "auc": 0.0,
+    }
+
+    pool = multiprocessing.Pool(cores)
+
+    u_batch_size = 5000
+    i_batch_size = BATCH_SIZE
+
+    test_users = users_to_test
+    n_test_users = len(test_users)
+    n_user_batchs = n_test_users // u_batch_size + 1
+
+    count = 0
+
+    for u_batch_id in range(n_user_batchs):
+        start = u_batch_id * u_batch_size
+        end = (u_batch_id + 1) * u_batch_size
+
+        user_batch = test_users[start:end]
+
+        if batch_test_flag:
+            # batch-item test
+            n_item_batchs = ITEM_NUM // i_batch_size + 1
+            rate_batch = np.zeros(shape=(len(user_batch), ITEM_NUM))
+
+            i_count = 0
+            for i_batch_id in range(n_item_batchs):
+                i_start = i_batch_id * i_batch_size
+                i_end = min((i_batch_id + 1) * i_batch_size, ITEM_NUM)
+
+                item_batch = range(i_start, i_end)
+
+                u_g_embeddings, pos_i_g_embeddings, _ = model(
+                    g, "user", "item", user_batch, item_batch, []
+                )
+                i_rate_batch = (
+                    model.rating(u_g_embeddings, pos_i_g_embeddings)
+                    .detach()
+                    .cpu()
+                )
+
+                rate_batch[:, i_start:i_end] = i_rate_batch
+                i_count += i_rate_batch.shape[1]
+
+            assert i_count == ITEM_NUM
+
+        else:
+            # all-item test
+            item_batch = range(ITEM_NUM)
+            u_g_embeddings, pos_i_g_embeddings, _ = model(
+                g, "user", "item", user_batch, item_batch, []
+            )
+            rate_batch = (
+                model.rating(u_g_embeddings, pos_i_g_embeddings).detach().cpu()
+            )
+
+        user_batch_rating_uid = zip(rate_batch.numpy(), user_batch)
+        batch_result = pool.map(test_one_user, user_batch_rating_uid)
+        count += len(batch_result)
+
+        for re in batch_result:
+            result["precision"] += re["precision"] / n_test_users
+            result["recall"] += re["recall"] / n_test_users
+            result["ndcg"] += re["ndcg"] / n_test_users
+            result["hit_ratio"] += re["hit_ratio"] / n_test_users
+            result["auc"] += re["auc"] / n_test_users
+
+    assert count == n_test_users
+    pool.close()
+    return result

examples/pytorch/NGCF/NGCF/utility/helper.py

@@ -0,0 +1,68 @@
+# This file is copied from the NGCF author's implementation
+# <https://github.com/xiangwang1223/neural_graph_collaborative_filtering/blob/master/NGCF/utility/helper.py>.
+# It implements the helper functions.
+"""
+Created on Aug 19, 2016
+@author: Xiang Wang (xiangwang@u.nus.edu)
+"""
+__author__ = "xiangwang"
+import os
+import re
+
+
+def txt2list(file_src):
+    orig_file = open(file_src, "r")
+    lines = orig_file.readlines()
+    return lines
+
+
+def ensureDir(dir_path):
+    d = os.path.dirname(dir_path)
+    if not os.path.exists(d):
+        os.makedirs(d)
+
+
+def uni2str(unicode_str):
+    return str(unicode_str.encode("ascii", "ignore")).replace("\n", "").strip()
+
+
+def hasNumbers(inputString):
+    return bool(re.search(r"\d", inputString))
+
+
+def delMultiChar(inputString, chars):
+    for ch in chars:
+        inputString = inputString.replace(ch, "")
+    return inputString
+
+
+def merge_two_dicts(x, y):
+    z = x.copy()  # start with x's keys and values
+    z.update(y)  # modifies z with y's keys and values & returns None
+    return z
+
+
+def early_stopping(
+    log_value, best_value, stopping_step, expected_order="acc", flag_step=100
+):
+    # early stopping strategy:
+    assert expected_order in ["acc", "dec"]
+
+    if (expected_order == "acc" and log_value >= best_value) or (
+        expected_order == "dec" and log_value <= best_value
+    ):
+        stopping_step = 0
+        best_value = log_value
+    else:
+        stopping_step += 1
+
+    if stopping_step >= flag_step:
+        print(
+            "Early stopping is trigger at step: {} log:{}".format(
+                flag_step, log_value
+            )
+        )
+        should_stop = True
+    else:
+        should_stop = False
+    return best_value, stopping_step, should_stop

examples/pytorch/NGCF/NGCF/utility/load_data.py

@@ -0,0 +1,151 @@
+# This file is based on the NGCF author's implementation
+# <https://github.com/xiangwang1223/neural_graph_collaborative_filtering/blob/master/NGCF/utility/load_data.py>.
+# It implements the data processing and graph construction.
+import random as rd
+
+import dgl
+
+import numpy as np
+
+
+class Data(object):
+    def __init__(self, path, batch_size):
+        self.path = path
+        self.batch_size = batch_size
+
+        train_file = path + "/train.txt"
+        test_file = path + "/test.txt"
+
+        # get number of users and items
+        self.n_users, self.n_items = 0, 0
+        self.n_train, self.n_test = 0, 0
+        self.exist_users = []
+
+        user_item_src = []
+        user_item_dst = []
+
+        with open(train_file) as f:
+            for l in f.readlines():
+                if len(l) > 0:
+                    l = l.strip("\n").split(" ")
+                    items = [int(i) for i in l[1:]]
+                    uid = int(l[0])
+                    self.exist_users.append(uid)
+                    self.n_items = max(self.n_items, max(items))
+                    self.n_users = max(self.n_users, uid)
+                    self.n_train += len(items)
+                    for i in l[1:]:
+                        user_item_src.append(uid)
+                        user_item_dst.append(int(i))
+
+        with open(test_file) as f:
+            for l in f.readlines():
+                if len(l) > 0:
+                    l = l.strip("\n")
+                    try:
+                        items = [int(i) for i in l.split(" ")[1:]]
+                    except Exception:
+                        continue
+                    self.n_items = max(self.n_items, max(items))
+                    self.n_test += len(items)
+        self.n_items += 1
+        self.n_users += 1
+
+        self.print_statistics()
+
+        # training positive items corresponding to each user; testing positive items corresponding to each user
+        self.train_items, self.test_set = {}, {}
+        with open(train_file) as f_train:
+            with open(test_file) as f_test:
+                for l in f_train.readlines():
+                    if len(l) == 0:
+                        break
+                    l = l.strip("\n")
+                    items = [int(i) for i in l.split(" ")]
+                    uid, train_items = items[0], items[1:]
+                    self.train_items[uid] = train_items
+
+                for l in f_test.readlines():
+                    if len(l) == 0:
+                        break
+                    l = l.strip("\n")
+                    try:
+                        items = [int(i) for i in l.split(" ")]
+                    except Exception:
+                        continue
+
+                    uid, test_items = items[0], items[1:]
+                    self.test_set[uid] = test_items
+
+        # construct graph from the train data and add self-loops
+        user_selfs = [i for i in range(self.n_users)]
+        item_selfs = [i for i in range(self.n_items)]
+
+        data_dict = {
+            ("user", "user_self", "user"): (user_selfs, user_selfs),
+            ("item", "item_self", "item"): (item_selfs, item_selfs),
+            ("user", "ui", "item"): (user_item_src, user_item_dst),
+            ("item", "iu", "user"): (user_item_dst, user_item_src),
+        }
+        num_dict = {"user": self.n_users, "item": self.n_items}
+
+        self.g = dgl.heterograph(data_dict, num_nodes_dict=num_dict)
+
+    def sample(self):
+        if self.batch_size <= self.n_users:
+            users = rd.sample(self.exist_users, self.batch_size)
+        else:
+            users = [
+                rd.choice(self.exist_users) for _ in range(self.batch_size)
+            ]
+
+        def sample_pos_items_for_u(u, num):
+            # sample num pos items for u-th user
+            pos_items = self.train_items[u]
+            n_pos_items = len(pos_items)
+            pos_batch = []
+            while True:
+                if len(pos_batch) == num:
+                    break
+                pos_id = np.random.randint(low=0, high=n_pos_items, size=1)[0]
+                pos_i_id = pos_items[pos_id]
+
+                if pos_i_id not in pos_batch:
+                    pos_batch.append(pos_i_id)
+            return pos_batch
+
+        def sample_neg_items_for_u(u, num):
+            # sample num neg items for u-th user
+            neg_items = []
+            while True:
+                if len(neg_items) == num:
+                    break
+                neg_id = np.random.randint(low=0, high=self.n_items, size=1)[0]
+                if (
+                    neg_id not in self.train_items[u]
+                    and neg_id not in neg_items
+                ):
+                    neg_items.append(neg_id)
+            return neg_items
+
+        pos_items, neg_items = [], []
+        for u in users:
+            pos_items += sample_pos_items_for_u(u, 1)
+            neg_items += sample_neg_items_for_u(u, 1)
+
+        return users, pos_items, neg_items
+
+    def get_num_users_items(self):
+        return self.n_users, self.n_items
+
+    def print_statistics(self):
+        print("n_users=%d, n_items=%d" % (self.n_users, self.n_items))
+        print("n_interactions=%d" % (self.n_train + self.n_test))
+        print(
+            "n_train=%d, n_test=%d, sparsity=%.5f"
+            % (
+                self.n_train,
+                self.n_test,
+                (self.n_train + self.n_test) / (self.n_users * self.n_items),
+            )
+        )

examples/pytorch/NGCF/NGCF/utility/metrics.py

@@ -0,0 +1,112 @@
+# This file is copied from the NGCF author's implementation
+# <https://github.com/xiangwang1223/neural_graph_collaborative_filtering/blob/master/NGCF/utility/metrics.py>.
+# It implements the metrics.
+"""
+Created on Oct 10, 2018
+Tensorflow Implementation of Neural Graph Collaborative Filtering (NGCF) model in:
+Wang Xiang et al. Neural Graph Collaborative Filtering. In SIGIR 2019.
+@author: Xiang Wang (xiangwang@u.nus.edu)
+"""
+
+import numpy as np
+from sklearn.metrics import roc_auc_score
+
+
+def recall(rank, ground_truth, N):
+    return len(set(rank[:N]) & set(ground_truth)) / float(
+        len(set(ground_truth))
+    )
+
+
+def precision_at_k(r, k):
+    """Score is precision @ k
+    Relevance is binary (nonzero is relevant).
+    Returns:
+        Precision @ k
+    Raises:
+        ValueError: len(r) must be >= k
+    """
+    assert k >= 1
+    r = np.asarray(r)[:k]
+    return np.mean(r)
+
+
+def average_precision(r, cut):
+    """Score is average precision (area under PR curve)
+    Relevance is binary (nonzero is relevant).
+    Returns:
+        Average precision
+    """
+    r = np.asarray(r)
+    out = [precision_at_k(r, k + 1) for k in range(cut) if r[k]]
+    if not out:
+        return 0.0
+    return np.sum(out) / float(min(cut, np.sum(r)))
+
+
+def mean_average_precision(rs):
+    """Score is mean average precision
+    Relevance is binary (nonzero is relevant).
+    Returns:
+        Mean average precision
+    """
+    return np.mean([average_precision(r) for r in rs])
+
+
+def dcg_at_k(r, k, method=1):
+    """Score is discounted cumulative gain (dcg)
+    Relevance is positive real values.  Can use binary
+    as the previous methods.
+    Returns:
+        Discounted cumulative gain
+    """
+    r = np.asfarray(r)[:k]
+    if r.size:
+        if method == 0:
+            return r[0] + np.sum(r[1:] / np.log2(np.arange(2, r.size + 1)))
+        elif method == 1:
+            return np.sum(r / np.log2(np.arange(2, r.size + 2)))
+        else:
+            raise ValueError("method must be 0 or 1.")
+    return 0.0
+
+
+def ndcg_at_k(r, k, method=1):
+    """Score is normalized discounted cumulative gain (ndcg)
+    Relevance is positive real values.  Can use binary
+    as the previous methods.
+    Returns:
+        Normalized discounted cumulative gain
+    """
+    dcg_max = dcg_at_k(sorted(r, reverse=True), k, method)
+    if not dcg_max:
+        return 0.0
+    return dcg_at_k(r, k, method) / dcg_max
+
+
+def recall_at_k(r, k, all_pos_num):
+    r = np.asfarray(r)[:k]
+    return np.sum(r) / all_pos_num
+
+
+def hit_at_k(r, k):
+    r = np.array(r)[:k]
+    if np.sum(r) > 0:
+        return 1.0
+    else:
+        return 0.0
+
+
+def F1(pre, rec):
+    if pre + rec > 0:
+        return (2.0 * pre * rec) / (pre + rec)
+    else:
+        return 0.0
+
+
+def auc(ground_truth, prediction):
+    try:
+        res = roc_auc_score(y_true=ground_truth, y_score=prediction)
+    except Exception:
+        res = 0.0
+    return res

examples/pytorch/NGCF/NGCF/utility/parser.py

@@ -0,0 +1,90 @@
+# This file is based on the NGCF author's implementation
+# <https://github.com/xiangwang1223/neural_graph_collaborative_filtering/blob/master/NGCF/utility/parser.py>.
+
+import argparse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run NGCF.")
+    parser.add_argument(
+        "--weights_path", nargs="?", default="model/", help="Store model path."
+    )
+    parser.add_argument(
+        "--data_path", nargs="?", default="../Data/", help="Input data path."
+    )
+    parser.add_argument(
+        "--model_name", type=str, default="NGCF.pkl", help="Saved model name."
+    )
+
+    parser.add_argument(
+        "--dataset",
+        nargs="?",
+        default="gowalla",
+        help="Choose a dataset from {gowalla, yelp2018, amazon-book}",
+    )
+    parser.add_argument(
+        "--verbose", type=int, default=1, help="Interval of evaluation."
+    )
+    parser.add_argument(
+        "--epoch", type=int, default=400, help="Number of epoch."
+    )
+
+    parser.add_argument(
+        "--embed_size", type=int, default=64, help="Embedding size."
+    )
+    parser.add_argument(
+        "--layer_size",
+        nargs="?",
+        default="[64,64,64]",
+        help="Output sizes of every layer",
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=1024, help="Batch size."
+    )
+
+    parser.add_argument(
+        "--regs", nargs="?", default="[1e-5]", help="Regularizations."
+    )
+    parser.add_argument(
+        "--lr", type=float, default=0.0001, help="Learning rate."
+    )
+
+    parser.add_argument(
+        "--gpu", type=int, default=0, help="0 for NAIS_prod, 1 for NAIS_concat"
+    )
+
+    parser.add_argument(
+        "--mess_dropout",
+        nargs="?",
+        default="[0.1,0.1,0.1]",
+        help="Keep probability w.r.t. message dropout (i.e., 1-dropout_ratio) for each deep layer. 1: no dropout.",
+    )
+
+    parser.add_argument(
+        "--Ks",
+        nargs="?",
+        default="[20, 40]",
+        help="Output sizes of every layer",
+    )
+
+    parser.add_argument(
+        "--save_flag",
+        type=int,
+        default=1,
+        help="0: Disable model saver, 1: Activate model saver",
+    )
+
+    parser.add_argument(
+        "--test_flag",
+        nargs="?",
+        default="part",
+        help="Specify the test type from {part, full}, indicating whether the reference is done in mini-batch",
+    )
+
+    parser.add_argument(
+        "--report",
+        type=int,
+        default=0,
+        help="0: Disable performance report w.r.t. sparsity levels, 1: Show performance report w.r.t. sparsity levels",
+    )
+    return parser.parse_args()

examples/pytorch/NGCF/README.md

@@ -0,0 +1,48 @@
+# DGL Implementation of the NGCF Model
+
+This DGL example implements the GNN model proposed in the paper [Neural Graph Collaborative Filtering](https://arxiv.org/abs/1905.08108). 
+The author's codes of implementation is in [here](https://github.com/xiangwang1223/neural_graph_collaborative_filtering). A pytorch re-implementation can be found [here](https://github.com/huangtinglin/NGCF-PyTorch).
+
+Example implementor
+----------------------
+This example was implemented by [Kounianhua Du](https://github.com/KounianhuaDu) during her Software Dev Engineer Intern work at the AWS Shanghai AI Lab.
+
+
+The graph dataset used in this example 
+---------------------------------------
+Gowalla: This is the check-in dataset obtained from Gowalla, where users share their locations by checking-in. To ensure the quality of the dataset, we use the 10-core setting, i.e., retaining users and items with at least ten interactions. The dataset used can be found [here](https://github.com/xiangwang1223/neural_graph_collaborative_filtering/tree/master/Data).
+
+Statistics:
+- Users: 29858
+- Items: 40981
+- Interactions: 1027370
+- Density: 0.00084
+
+
+How to run example files
+--------------------------------
+First to get the data, in the Data folder, run
+
+```bash
+sh load_gowalla.sh
+```
+
+Then, in the NGCF folder, run
+
+```bash
+python main.py --dataset gowalla --regs [1e-5] --embed_size 64 --layer_size [64,64,64] --lr 0.0001 --save_flag 1 --batch_size 1024 --epoch 400 --verbose 1 --mess_dropout [0.1,0.1,0.1] --gpu 0 
+```
+
+NOTE: Following the paper's setting, the node dropout is disabled.
+
+
+Performance
+-------------------------
+The following results are the results in 400 epoches.
+
+**NGCF results**
+| Model         | Paper (tensorflow)               | ours (DGL)                  |
+| ------------- | -------------------------------- | --------------------------- |
+| recall@20     | 0.1569                           | 0.1552                      |
+| ndcg@20       | 0.1327                           | 0.2707                      |
+

examples/pytorch/P-GNN/README.md

@@ -0,0 +1,52 @@
+# DGL Implementations of P-GNN
+
+This DGL example implements the GNN model proposed in the paper [Position-aware Graph Neural Networks](http://proceedings.mlr.press/v97/you19b/you19b.pdf). For the original implementation, see [here](https://github.com/JiaxuanYou/P-GNN).
+
+Contributor: [RecLusIve-F](https://github.com/RecLusIve-F)
+
+## Requirements
+
+The codebase is implemented in Python 3.8. For version requirement of packages, see below.
+
+```
+dgl 0.7.2
+numpy 1.21.2
+torch 1.10.1
+networkx 2.6.3
+scikit-learn 1.0.2
+```
+
+## Instructions for experiments
+
+### Link prediction
+
+```bash
+# Communities-T
+python main.py --task link
+
+# Communities
+python main.py --task link --inductive
+```
+
+### Link pair prediction
+
+```bash
+# Communities
+python main.py --task link_pair --inductive
+```
+
+## Performance
+
+### Link prediction (Grid-T and Communities-T refer to the transductive learning setting of Grid and Communities)
+
+|             Dataset              | Communities-T |  Communities  |
+| :------------------------------: | :-----------: | :-----------: |
+| ROC AUC ( P-GNN-E-2L in Table 1) | 0.988 ± 0.003 | 0.985 ± 0.008 |
+|    ROC AUC (DGL: P-GNN-E-2L)     | 0.984 ± 0.010 | 0.991 ± 0.004 |
+
+### Link pair prediction
+
+|             Dataset              | Communities |
+| :------------------------------: | :---------: |
+| ROC AUC ( P-GNN-E-2L in Table 1) | 1.0 ± 0.001 |
+|    ROC AUC (DGL: P-GNN-E-2L)     | 1.0 ± 0.000 |

examples/pytorch/P-GNN/main.py

@@ -0,0 +1,211 @@
+import os
+import warnings
+
+import dgl
+
+import numpy as np
+import torch
+import torch.nn as nn
+from model import PGNN
+from sklearn.metrics import roc_auc_score
+from utils import get_dataset, preselect_anchor
+
+warnings.filterwarnings("ignore")
+
+
+def get_loss(p, data, out, loss_func, device, get_auc=True):
+    edge_mask = np.concatenate(
+        (
+            data["positive_edges_{}".format(p)],
+            data["negative_edges_{}".format(p)],
+        ),
+        axis=-1,
+    )
+
+    nodes_first = torch.index_select(
+        out, 0, torch.from_numpy(edge_mask[0, :]).long().to(out.device)
+    )
+    nodes_second = torch.index_select(
+        out, 0, torch.from_numpy(edge_mask[1, :]).long().to(out.device)
+    )
+
+    pred = torch.sum(nodes_first * nodes_second, dim=-1)
+
+    label_positive = torch.ones(
+        [
+            data["positive_edges_{}".format(p)].shape[1],
+        ],
+        dtype=pred.dtype,
+    )
+    label_negative = torch.zeros(
+        [
+            data["negative_edges_{}".format(p)].shape[1],
+        ],
+        dtype=pred.dtype,
+    )
+    label = torch.cat((label_positive, label_negative)).to(device)
+    loss = loss_func(pred, label)
+
+    if get_auc:
+        auc = roc_auc_score(
+            label.flatten().cpu().numpy(),
+            torch.sigmoid(pred).flatten().data.cpu().numpy(),
+        )
+        return loss, auc
+    else:
+        return loss
+
+
+def train_model(data, model, loss_func, optimizer, device, g_data):
+    model.train()
+    out = model(g_data)
+
+    loss = get_loss("train", data, out, loss_func, device, get_auc=False)
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+
+    return g_data
+
+
+def eval_model(data, g_data, model, loss_func, device):
+    model.eval()
+    out = model(g_data)
+
+    # train loss and auc
+    tmp_loss, auc_train = get_loss("train", data, out, loss_func, device)
+    loss_train = tmp_loss.cpu().data.numpy()
+
+    # val loss and auc
+    _, auc_val = get_loss("val", data, out, loss_func, device)
+
+    # test loss and auc
+    _, auc_test = get_loss("test", data, out, loss_func, device)
+
+    return loss_train, auc_train, auc_val, auc_test
+
+
+def main(args):
+    # The mean and standard deviation of the experiment results
+    # are stored in the 'results' folder
+    if not os.path.isdir("results"):
+        os.mkdir("results")
+
+    if torch.cuda.is_available():
+        device = "cuda:0"
+    else:
+        device = "cpu"
+
+    print(
+        "Learning Type: {}".format(
+            ["Transductive", "Inductive"][args.inductive]
+        ),
+        "Task: {}".format(args.task),
+    )
+
+    results = []
+
+    for repeat in range(args.repeat_num):
+        data = get_dataset(args)
+
+        # pre-sample anchor nodes and compute shortest distance values for all epochs
+        (
+            g_list,
+            anchor_eid_list,
+            dist_max_list,
+            edge_weight_list,
+        ) = preselect_anchor(data, args)
+
+        # model
+        model = PGNN(input_dim=data["feature"].shape[1]).to(device)
+
+        # loss
+        optimizer = torch.optim.Adam(
+            model.parameters(), lr=1e-2, weight_decay=5e-4
+        )
+        loss_func = nn.BCEWithLogitsLoss()
+
+        best_auc_val = -1
+        best_auc_test = -1
+
+        for epoch in range(args.epoch_num):
+            if epoch == 200:
+                for param_group in optimizer.param_groups:
+                    param_group["lr"] /= 10
+
+            g = dgl.graph(g_list[epoch])
+            g.ndata["feat"] = torch.FloatTensor(data["feature"])
+            g.edata["sp_dist"] = torch.FloatTensor(edge_weight_list[epoch])
+            g_data = {
+                "graph": g.to(device),
+                "anchor_eid": anchor_eid_list[epoch],
+                "dists_max": dist_max_list[epoch],
+            }
+
+            train_model(data, model, loss_func, optimizer, device, g_data)
+
+            loss_train, auc_train, auc_val, auc_test = eval_model(
+                data, g_data, model, loss_func, device
+            )
+            if auc_val > best_auc_val:
+                best_auc_val = auc_val
+                best_auc_test = auc_test
+
+            if epoch % args.epoch_log == 0:
+                print(
+                    repeat,
+                    epoch,
+                    "Loss {:.4f}".format(loss_train),
+                    "Train AUC: {:.4f}".format(auc_train),
+                    "Val AUC: {:.4f}".format(auc_val),
+                    "Test AUC: {:.4f}".format(auc_test),
+                    "Best Val AUC: {:.4f}".format(best_auc_val),
+                    "Best Test AUC: {:.4f}".format(best_auc_test),
+                )
+
+        results.append(best_auc_test)
+
+    results = np.array(results)
+    results_mean = np.mean(results).round(6)
+    results_std = np.std(results).round(6)
+    print("-----------------Final-------------------")
+    print(results_mean, results_std)
+
+    with open(
+        "results/{}_{}_{}.txt".format(
+            ["Transductive", "Inductive"][args.inductive],
+            args.task,
+            args.k_hop_dist,
+        ),
+        "w",
+    ) as f:
+        f.write("{}, {}\n".format(results_mean, results_std))
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--task", type=str, default="link", choices=["link", "link_pair"]
+    )
+    parser.add_argument(
+        "--inductive",
+        action="store_true",
+        help="Inductive learning or transductive learning",
+    )
+    parser.add_argument(
+        "--k_hop_dist",
+        default=-1,
+        type=int,
+        help="K-hop shortest path distance, -1 means exact shortest path.",
+    )
+
+    parser.add_argument("--epoch_num", type=int, default=2000)
+    parser.add_argument("--repeat_num", type=int, default=10)
+    parser.add_argument("--epoch_log", type=int, default=100)
+
+    args = parser.parse_args()
+    main(args)

examples/pytorch/P-GNN/model.py

@@ -0,0 +1,68 @@
+import dgl.function as fn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PGNN_layer(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super(PGNN_layer, self).__init__()
+        self.input_dim = input_dim
+
+        self.linear_hidden_u = nn.Linear(input_dim, output_dim)
+        self.linear_hidden_v = nn.Linear(input_dim, output_dim)
+        self.linear_out_position = nn.Linear(output_dim, 1)
+        self.act = nn.ReLU()
+
+    def forward(self, graph, feature, anchor_eid, dists_max):
+        with graph.local_scope():
+            u_feat = self.linear_hidden_u(feature)
+            v_feat = self.linear_hidden_v(feature)
+            graph.srcdata.update({"u_feat": u_feat})
+            graph.dstdata.update({"v_feat": v_feat})
+
+            graph.apply_edges(fn.u_mul_e("u_feat", "sp_dist", "u_message"))
+            graph.apply_edges(fn.v_add_e("v_feat", "u_message", "message"))
+
+            messages = torch.index_select(
+                graph.edata["message"],
+                0,
+                torch.LongTensor(anchor_eid).to(feature.device),
+            )
+            messages = messages.reshape(
+                dists_max.shape[0], dists_max.shape[1], messages.shape[-1]
+            )
+
+            messages = self.act(messages)  # n*m*d
+
+            out_position = self.linear_out_position(messages).squeeze(
+                -1
+            )  # n*m_out
+            out_structure = torch.mean(messages, dim=1)  # n*d
+
+            return out_position, out_structure
+
+
+class PGNN(nn.Module):
+    def __init__(self, input_dim, feature_dim=32, dropout=0.5):
+        super(PGNN, self).__init__()
+        self.dropout = nn.Dropout(dropout)
+
+        self.linear_pre = nn.Linear(input_dim, feature_dim)
+        self.conv_first = PGNN_layer(feature_dim, feature_dim)
+        self.conv_out = PGNN_layer(feature_dim, feature_dim)
+
+    def forward(self, data):
+        x = data["graph"].ndata["feat"]
+        graph = data["graph"]
+        x = self.linear_pre(x)
+        x_position, x = self.conv_first(
+            graph, x, data["anchor_eid"], data["dists_max"]
+        )
+
+        x = self.dropout(x)
+        x_position, x = self.conv_out(
+            graph, x, data["anchor_eid"], data["dists_max"]
+        )
+        x_position = F.normalize(x_position, p=2, dim=-1)
+        return x_position

examples/pytorch/P-GNN/utils.py

@@ -0,0 +1,328 @@
+import multiprocessing as mp
+import random
+from multiprocessing import get_context
+
+import networkx as nx
+import numpy as np
+import torch
+from tqdm.auto import tqdm
+
+
+def get_communities(remove_feature):
+    community_size = 20
+
+    # Create 20 cliques (communities) of size 20,
+    # then rewire a single edge in each clique to a node in an adjacent clique
+    graph = nx.connected_caveman_graph(20, community_size)
+
+    # Randomly rewire 1% edges
+    node_list = list(graph.nodes)
+    for u, v in graph.edges():
+        if random.random() < 0.01:
+            x = random.choice(node_list)
+            if graph.has_edge(u, x):
+                continue
+            graph.remove_edge(u, v)
+            graph.add_edge(u, x)
+
+    # remove self-loops
+    graph.remove_edges_from(nx.selfloop_edges(graph))
+    edge_index = np.array(list(graph.edges))
+    # Add (i, j) for an edge (j, i)
+    edge_index = np.concatenate((edge_index, edge_index[:, ::-1]), axis=0)
+    edge_index = torch.from_numpy(edge_index).long().permute(1, 0)
+
+    n = graph.number_of_nodes()
+    label = np.zeros((n, n), dtype=int)
+    for u in node_list:
+        # the node IDs are simply consecutive integers from 0
+        for v in range(u):
+            if u // community_size == v // community_size:
+                label[u, v] = 1
+
+    if remove_feature:
+        feature = torch.ones((n, 1))
+    else:
+        rand_order = np.random.permutation(n)
+        feature = np.identity(n)[:, rand_order]
+
+    data = {
+        "edge_index": edge_index,
+        "feature": feature,
+        "positive_edges": np.stack(np.nonzero(label)),
+        "num_nodes": feature.shape[0],
+    }
+
+    return data
+
+
+def to_single_directed(edges):
+    edges_new = np.zeros((2, edges.shape[1] // 2), dtype=int)
+    j = 0
+    for i in range(edges.shape[1]):
+        if edges[0, i] < edges[1, i]:
+            edges_new[:, j] = edges[:, i]
+            j += 1
+
+    return edges_new
+
+
+# each node at least remain in the new graph
+def split_edges(p, edges, data, non_train_ratio=0.2):
+    e = edges.shape[1]
+    edges = edges[:, np.random.permutation(e)]
+    split1 = int((1 - non_train_ratio) * e)
+    split2 = int((1 - non_train_ratio / 2) * e)
+
+    data.update(
+        {
+            "{}_edges_train".format(p): edges[:, :split1],  # 80%
+            "{}_edges_val".format(p): edges[:, split1:split2],  # 10%
+            "{}_edges_test".format(p): edges[:, split2:],  # 10%
+        }
+    )
+
+
+def to_bidirected(edges):
+    return np.concatenate((edges, edges[::-1, :]), axis=-1)
+
+
+def get_negative_edges(positive_edges, num_nodes, num_negative_edges):
+    positive_edge_set = []
+    positive_edges = to_bidirected(positive_edges)
+    for i in range(positive_edges.shape[1]):
+        positive_edge_set.append(tuple(positive_edges[:, i]))
+    positive_edge_set = set(positive_edge_set)
+
+    negative_edges = np.zeros(
+        (2, num_negative_edges), dtype=positive_edges.dtype
+    )
+    for i in range(num_negative_edges):
+        while True:
+            mask_temp = tuple(
+                np.random.choice(num_nodes, size=(2,), replace=False)
+            )
+            if mask_temp not in positive_edge_set:
+                negative_edges[:, i] = mask_temp
+                break
+
+    return negative_edges
+
+
+def get_pos_neg_edges(data, infer_link_positive=True):
+    if infer_link_positive:
+        data["positive_edges"] = to_single_directed(data["edge_index"].numpy())
+    split_edges("positive", data["positive_edges"], data)
+
+    # resample edge mask link negative
+    negative_edges = get_negative_edges(
+        data["positive_edges"],
+        data["num_nodes"],
+        num_negative_edges=data["positive_edges"].shape[1],
+    )
+    split_edges("negative", negative_edges, data)
+
+    return data
+
+
+def shortest_path(graph, node_range, cutoff):
+    dists_dict = {}
+    for node in tqdm(node_range, leave=False):
+        dists_dict[node] = nx.single_source_shortest_path_length(
+            graph, node, cutoff
+        )
+    return dists_dict
+
+
+def merge_dicts(dicts):
+    result = {}
+    for dictionary in dicts:
+        result.update(dictionary)
+    return result
+
+
+def all_pairs_shortest_path(graph, cutoff=None, num_workers=4):
+    nodes = list(graph.nodes)
+    random.shuffle(nodes)
+    pool = mp.Pool(processes=num_workers)
+    interval_size = len(nodes) / num_workers
+    results = [
+        pool.apply_async(
+            shortest_path,
+            args=(
+                graph,
+                nodes[int(interval_size * i) : int(interval_size * (i + 1))],
+                cutoff,
+            ),
+        )
+        for i in range(num_workers)
+    ]
+    output = [p.get() for p in results]
+    dists_dict = merge_dicts(output)
+    pool.close()
+    pool.join()
+    return dists_dict
+
+
+def precompute_dist_data(edge_index, num_nodes, approximate=0):
+    """
+    Here dist is 1/real_dist, higher actually means closer, 0 means disconnected
+    :return:
+    """
+    graph = nx.Graph()
+    edge_list = edge_index.transpose(1, 0).tolist()
+    graph.add_edges_from(edge_list)
+
+    n = num_nodes
+    dists_array = np.zeros((n, n))
+    dists_dict = all_pairs_shortest_path(
+        graph, cutoff=approximate if approximate > 0 else None
+    )
+    node_list = graph.nodes()
+    for node_i in node_list:
+        shortest_dist = dists_dict[node_i]
+        for node_j in node_list:
+            dist = shortest_dist.get(node_j, -1)
+            if dist != -1:
+                dists_array[node_i, node_j] = 1 / (dist + 1)
+    return dists_array
+
+
+def get_dataset(args):
+    # Generate graph data
+    data_info = get_communities(args.inductive)
+    # Get positive and negative edges
+    data = get_pos_neg_edges(
+        data_info, infer_link_positive=True if args.task == "link" else False
+    )
+    # Pre-compute shortest path length
+    if args.task == "link":
+        dists_removed = precompute_dist_data(
+            data["positive_edges_train"],
+            data["num_nodes"],
+            approximate=args.k_hop_dist,
+        )
+        data["dists"] = torch.from_numpy(dists_removed).float()
+        data["edge_index"] = torch.from_numpy(
+            to_bidirected(data["positive_edges_train"])
+        ).long()
+    else:
+        dists = precompute_dist_data(
+            data["edge_index"].numpy(),
+            data["num_nodes"],
+            approximate=args.k_hop_dist,
+        )
+        data["dists"] = torch.from_numpy(dists).float()
+
+    return data
+
+
+def get_anchors(n):
+    """Get a list of NumPy arrays, each of them is an anchor node set"""
+    m = int(np.log2(n))
+    anchor_set_id = []
+    for i in range(m):
+        anchor_size = int(n / np.exp2(i + 1))
+        for _ in range(m):
+            anchor_set_id.append(
+                np.random.choice(n, size=anchor_size, replace=False)
+            )
+    return anchor_set_id
+
+
+def get_dist_max(anchor_set_id, dist):
+    # N x K, N is number of nodes, K is the number of anchor sets
+    dist_max = torch.zeros((dist.shape[0], len(anchor_set_id)))
+    dist_argmax = torch.zeros((dist.shape[0], len(anchor_set_id))).long()
+    for i in range(len(anchor_set_id)):
+        temp_id = torch.as_tensor(anchor_set_id[i], dtype=torch.long)
+        # Get reciprocal of shortest distance to each node in the i-th anchor set
+        dist_temp = torch.index_select(dist, 1, temp_id)
+        # For each node in the graph, find its closest anchor node in the set
+        # and the reciprocal of shortest distance
+        dist_max_temp, dist_argmax_temp = torch.max(dist_temp, dim=-1)
+        dist_max[:, i] = dist_max_temp
+        dist_argmax[:, i] = torch.index_select(temp_id, 0, dist_argmax_temp)
+    return dist_max, dist_argmax
+
+
+def get_a_graph(dists_max, dists_argmax):
+    src = []
+    dst = []
+    real_src = []
+    real_dst = []
+    edge_weight = []
+    dists_max = dists_max.numpy()
+    for i in range(dists_max.shape[0]):
+        # Get unique closest anchor nodes for node i across all anchor sets
+        tmp_dists_argmax, tmp_dists_argmax_idx = np.unique(
+            dists_argmax[i, :], True
+        )
+        src.extend([i] * tmp_dists_argmax.shape[0])
+        real_src.extend([i] * dists_argmax[i, :].shape[0])
+        real_dst.extend(list(dists_argmax[i, :].numpy()))
+        dst.extend(list(tmp_dists_argmax))
+        edge_weight.extend(dists_max[i, tmp_dists_argmax_idx].tolist())
+    eid_dict = {(u, v): i for i, (u, v) in enumerate(list(zip(dst, src)))}
+    anchor_eid = [eid_dict.get((u, v)) for u, v in zip(real_dst, real_src)]
+    g = (dst, src)
+    return g, anchor_eid, edge_weight
+
+
+def get_graphs(data, anchor_sets):
+    graphs = []
+    anchor_eids = []
+    dists_max_list = []
+    edge_weights = []
+    for anchor_set in tqdm(anchor_sets, leave=False):
+        dists_max, dists_argmax = get_dist_max(anchor_set, data["dists"])
+        g, anchor_eid, edge_weight = get_a_graph(dists_max, dists_argmax)
+        graphs.append(g)
+        anchor_eids.append(anchor_eid)
+        dists_max_list.append(dists_max)
+        edge_weights.append(edge_weight)
+
+    return graphs, anchor_eids, dists_max_list, edge_weights
+
+
+def merge_result(outputs):
+    graphs = []
+    anchor_eids = []
+    dists_max_list = []
+    edge_weights = []
+
+    for g, anchor_eid, dists_max, edge_weight in outputs:
+        graphs.extend(g)
+        anchor_eids.extend(anchor_eid)
+        dists_max_list.extend(dists_max)
+        edge_weights.extend(edge_weight)
+
+    return graphs, anchor_eids, dists_max_list, edge_weights
+
+
+def preselect_anchor(data, args, num_workers=4):
+    pool = get_context("spawn").Pool(processes=num_workers)
+    # Pre-compute anchor sets, a collection of anchor sets per epoch
+    anchor_set_ids = [
+        get_anchors(data["num_nodes"]) for _ in range(args.epoch_num)
+    ]
+    interval_size = len(anchor_set_ids) / num_workers
+    results = [
+        pool.apply_async(
+            get_graphs,
+            args=(
+                data,
+                anchor_set_ids[
+                    int(interval_size * i) : int(interval_size * (i + 1))
+                ],
+            ),
+        )
+        for i in range(num_workers)
+    ]
+
+    output = [p.get() for p in results]
+    graphs, anchor_eids, dists_max_list, edge_weights = merge_result(output)
+    pool.close()
+    pool.join()
+
+    return graphs, anchor_eids, dists_max_list, edge_weights

examples/pytorch/TAHIN/TAHIN.py

@@ -0,0 +1,245 @@
+import dgl
+import dgl.function as fn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.nn.pytorch import GATConv
+
+
+# Semantic attention in the metapath-based aggregation (the same as that in the HAN)
+class SemanticAttention(nn.Module):
+    def __init__(self, in_size, hidden_size=128):
+        super(SemanticAttention, self).__init__()
+
+        self.project = nn.Sequential(
+            nn.Linear(in_size, hidden_size),
+            nn.Tanh(),
+            nn.Linear(hidden_size, 1, bias=False),
+        )
+
+    def forward(self, z):
+        """
+        Shape of z: (N, M , D*K)
+        N: number of nodes
+        M: number of metapath patterns
+        D: hidden_size
+        K: number of heads
+        """
+        w = self.project(z).mean(0)  # (M, 1)
+        beta = torch.softmax(w, dim=0)  # (M, 1)
+        beta = beta.expand((z.shape[0],) + beta.shape)  # (N, M, 1)
+
+        return (beta * z).sum(1)  # (N, D * K)
+
+
+# Metapath-based aggregation (the same as the HANLayer)
+class HANLayer(nn.Module):
+    def __init__(
+        self, meta_path_patterns, in_size, out_size, layer_num_heads, dropout
+    ):
+        super(HANLayer, self).__init__()
+
+        # One GAT layer for each meta path based adjacency matrix
+        self.gat_layers = nn.ModuleList()
+        for i in range(len(meta_path_patterns)):
+            self.gat_layers.append(
+                GATConv(
+                    in_size,
+                    out_size,
+                    layer_num_heads,
+                    dropout,
+                    dropout,
+                    activation=F.elu,
+                    allow_zero_in_degree=True,
+                )
+            )
+        self.semantic_attention = SemanticAttention(
+            in_size=out_size * layer_num_heads
+        )
+        self.meta_path_patterns = list(
+            tuple(meta_path_pattern) for meta_path_pattern in meta_path_patterns
+        )
+
+        self._cached_graph = None
+        self._cached_coalesced_graph = {}
+
+    def forward(self, g, h):
+        semantic_embeddings = []
+        # obtain metapath reachable graph
+        if self._cached_graph is None or self._cached_graph is not g:
+            self._cached_graph = g
+            self._cached_coalesced_graph.clear()
+            for meta_path_pattern in self.meta_path_patterns:
+                self._cached_coalesced_graph[
+                    meta_path_pattern
+                ] = dgl.metapath_reachable_graph(g, meta_path_pattern)
+
+        for i, meta_path_pattern in enumerate(self.meta_path_patterns):
+            new_g = self._cached_coalesced_graph[meta_path_pattern]
+            semantic_embeddings.append(self.gat_layers[i](new_g, h).flatten(1))
+        semantic_embeddings = torch.stack(
+            semantic_embeddings, dim=1
+        )  # (N, M, D * K)
+
+        return self.semantic_attention(semantic_embeddings)  # (N, D * K)
+
+
+# Relational neighbor aggregation
+class RelationalAGG(nn.Module):
+    def __init__(self, g, in_size, out_size, dropout=0.1):
+        super(RelationalAGG, self).__init__()
+        self.in_size = in_size
+        self.out_size = out_size
+
+        # Transform weights for different types of edges
+        self.W_T = nn.ModuleDict(
+            {
+                name: nn.Linear(in_size, out_size, bias=False)
+                for name in g.etypes
+            }
+        )
+
+        # Attention weights for different types of edges
+        self.W_A = nn.ModuleDict(
+            {name: nn.Linear(out_size, 1, bias=False) for name in g.etypes}
+        )
+
+        # layernorm
+        self.layernorm = nn.LayerNorm(out_size)
+
+        # dropout layer
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, g, feat_dict):
+        funcs = {}
+        for srctype, etype, dsttype in g.canonical_etypes:
+            g.nodes[dsttype].data["h"] = feat_dict[
+                dsttype
+            ]  # nodes' original feature
+            g.nodes[srctype].data["h"] = feat_dict[srctype]
+            g.nodes[srctype].data["t_h"] = self.W_T[etype](
+                feat_dict[srctype]
+            )  # src nodes' transformed feature
+
+            # compute the attention numerator (exp)
+            g.apply_edges(fn.u_mul_v("t_h", "h", "x"), etype=etype)
+            g.edges[etype].data["x"] = torch.exp(
+                self.W_A[etype](g.edges[etype].data["x"])
+            )
+
+            # first update to compute the attention denominator (\sum exp)
+            funcs[etype] = (fn.copy_e("x", "m"), fn.sum("m", "att"))
+        g.multi_update_all(funcs, "sum")
+
+        funcs = {}
+        for srctype, etype, dsttype in g.canonical_etypes:
+            g.apply_edges(
+                fn.e_div_v("x", "att", "att"), etype=etype
+            )  # compute attention weights (numerator/denominator)
+            funcs[etype] = (
+                fn.u_mul_e("h", "att", "m"),
+                fn.sum("m", "h"),
+            )  # \sum(h0*att) -> h1
+        # second update to obtain h1
+        g.multi_update_all(funcs, "sum")
+
+        # apply activation, layernorm, and dropout
+        feat_dict = {}
+        for ntype in g.ntypes:
+            feat_dict[ntype] = self.dropout(
+                self.layernorm(F.relu_(g.nodes[ntype].data["h"]))
+            )  # apply activation, layernorm, and dropout
+
+        return feat_dict
+
+
+class TAHIN(nn.Module):
+    def __init__(
+        self, g, meta_path_patterns, in_size, out_size, num_heads, dropout
+    ):
+        super(TAHIN, self).__init__()
+
+        # embeddings for different types of nodes, h0
+        self.initializer = nn.init.xavier_uniform_
+        self.feature_dict = nn.ParameterDict(
+            {
+                ntype: nn.Parameter(
+                    self.initializer(torch.empty(g.num_nodes(ntype), in_size))
+                )
+                for ntype in g.ntypes
+            }
+        )
+
+        # relational neighbor aggregation, this produces h1
+        self.RelationalAGG = RelationalAGG(g, in_size, out_size)
+
+        # metapath-based aggregation modules for user and item, this produces h2
+        self.meta_path_patterns = meta_path_patterns
+        # one HANLayer for user, one HANLayer for item
+        self.hans = nn.ModuleDict(
+            {
+                key: HANLayer(value, in_size, out_size, num_heads, dropout)
+                for key, value in self.meta_path_patterns.items()
+            }
+        )
+
+        # layers to combine h0, h1, and h2
+        # used to update node embeddings
+        self.user_layer1 = nn.Linear(
+            (num_heads + 1) * out_size, out_size, bias=True
+        )
+        self.user_layer2 = nn.Linear(2 * out_size, out_size, bias=True)
+        self.item_layer1 = nn.Linear(
+            (num_heads + 1) * out_size, out_size, bias=True
+        )
+        self.item_layer2 = nn.Linear(2 * out_size, out_size, bias=True)
+
+        # layernorm
+        self.layernorm = nn.LayerNorm(out_size)
+
+        # network to score the node pairs
+        self.pred = nn.Linear(out_size, out_size)
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(out_size, 1)
+
+    def forward(self, g, user_key, item_key, user_idx, item_idx):
+        # relational neighbor aggregation, h1
+        h1 = self.RelationalAGG(g, self.feature_dict)
+
+        # metapath-based aggregation, h2
+        h2 = {}
+        for key in self.meta_path_patterns.keys():
+            h2[key] = self.hans[key](g, self.feature_dict[key])
+
+        # update node embeddings
+        user_emb = torch.cat((h1[user_key], h2[user_key]), 1)
+        item_emb = torch.cat((h1[item_key], h2[item_key]), 1)
+        user_emb = self.user_layer1(user_emb)
+        item_emb = self.item_layer1(item_emb)
+        user_emb = self.user_layer2(
+            torch.cat((user_emb, self.feature_dict[user_key]), 1)
+        )
+        item_emb = self.item_layer2(
+            torch.cat((item_emb, self.feature_dict[item_key]), 1)
+        )
+
+        # Relu
+        user_emb = F.relu_(user_emb)
+        item_emb = F.relu_(item_emb)
+
+        # layer norm
+        user_emb = self.layernorm(user_emb)
+        item_emb = self.layernorm(item_emb)
+
+        # obtain users/items embeddings and their interactions
+        user_feat = user_emb[user_idx]
+        item_feat = item_emb[item_idx]
+        interaction = user_feat * item_feat
+
+        # score the node pairs
+        pred = self.pred(interaction)
+        pred = self.dropout(pred)  # dropout
+        pred = self.fc(pred)
+        pred = torch.sigmoid(pred)
+
+        return pred.squeeze(1)

examples/pytorch/TAHIN/data_loader.py

@@ -0,0 +1,399 @@
+import os
+import pickle as pkl
+import random
+
+import dgl
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+
+# Split data into train/eval/test
+def split_data(hg, etype_name):
+    src, dst = hg.edges(etype=etype_name)
+    user_item_src = src.numpy().tolist()
+    user_item_dst = dst.numpy().tolist()
+
+    num_link = len(user_item_src)
+    pos_label = [1] * num_link
+    pos_data = list(zip(user_item_src, user_item_dst, pos_label))
+
+    ui_adj = np.array(hg.adj_external(etype=etype_name).to_dense())
+    full_idx = np.where(ui_adj == 0)
+
+    sample = random.sample(range(0, len(full_idx[0])), num_link)
+    neg_label = [0] * num_link
+    neg_data = list(zip(full_idx[0][sample], full_idx[1][sample], neg_label))
+
+    full_data = pos_data + neg_data
+    random.shuffle(full_data)
+
+    train_size = int(len(full_data) * 0.6)
+    eval_size = int(len(full_data) * 0.2)
+    test_size = len(full_data) - train_size - eval_size
+    train_data = full_data[:train_size]
+    eval_data = full_data[train_size : train_size + eval_size]
+    test_data = full_data[
+        train_size + eval_size : train_size + eval_size + test_size
+    ]
+    train_data = np.array(train_data)
+    eval_data = np.array(eval_data)
+    test_data = np.array(test_data)
+
+    return train_data, eval_data, test_data
+
+
+def process_amazon(root_path):
+    # User-Item 3584 2753 50903 UIUI
+    # Item-View 2753 3857 5694 UIVI
+    # Item-Brand 2753 334 2753 UIBI
+    # Item-Category 2753 22 5508 UICI
+
+    # Construct graph from raw data.
+    # load data of amazon
+    data_path = os.path.join(root_path, "Amazon")
+    if not (os.path.exists(data_path)):
+        print(
+            "Can not find amazon in {}, please download the dataset first.".format(
+                data_path
+            )
+        )
+
+    # item_view
+    item_view_src = []
+    item_view_dst = []
+    with open(os.path.join(data_path, "item_view.dat")) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split(",")
+            item, view = int(_line[0]), int(_line[1])
+            item_view_src.append(item)
+            item_view_dst.append(view)
+
+    # user_item
+    user_item_src = []
+    user_item_dst = []
+    with open(os.path.join(data_path, "user_item.dat")) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split("\t")
+            user, item, rate = int(_line[0]), int(_line[1]), int(_line[2])
+            if rate > 3:
+                user_item_src.append(user)
+                user_item_dst.append(item)
+
+    # item_brand
+    item_brand_src = []
+    item_brand_dst = []
+    with open(os.path.join(data_path, "item_brand.dat")) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split(",")
+            item, brand = int(_line[0]), int(_line[1])
+            item_brand_src.append(item)
+            item_brand_dst.append(brand)
+
+    # item_category
+    item_category_src = []
+    item_category_dst = []
+    with open(os.path.join(data_path, "item_category.dat")) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split(",")
+            item, category = int(_line[0]), int(_line[1])
+            item_category_src.append(item)
+            item_category_dst.append(category)
+
+    # build graph
+    hg = dgl.heterograph(
+        {
+            ("item", "iv", "view"): (item_view_src, item_view_dst),
+            ("view", "vi", "item"): (item_view_dst, item_view_src),
+            ("user", "ui", "item"): (user_item_src, user_item_dst),
+            ("item", "iu", "user"): (user_item_dst, user_item_src),
+            ("item", "ib", "brand"): (item_brand_src, item_brand_dst),
+            ("brand", "bi", "item"): (item_brand_dst, item_brand_src),
+            ("item", "ic", "category"): (item_category_src, item_category_dst),
+            ("category", "ci", "item"): (item_category_dst, item_category_src),
+        }
+    )
+
+    print("Graph constructed.")
+
+    # Split data into train/eval/test
+    train_data, eval_data, test_data = split_data(hg, "ui")
+
+    # delete the positive edges in eval/test data in the original graph
+    train_pos = np.nonzero(train_data[:, 2])
+    train_pos_idx = train_pos[0]
+    user_item_src_processed = train_data[train_pos_idx, 0]
+    user_item_dst_processed = train_data[train_pos_idx, 1]
+    edges_dict = {
+        ("item", "iv", "view"): (item_view_src, item_view_dst),
+        ("view", "vi", "item"): (item_view_dst, item_view_src),
+        ("user", "ui", "item"): (
+            user_item_src_processed,
+            user_item_dst_processed,
+        ),
+        ("item", "iu", "user"): (
+            user_item_dst_processed,
+            user_item_src_processed,
+        ),
+        ("item", "ib", "brand"): (item_brand_src, item_brand_dst),
+        ("brand", "bi", "item"): (item_brand_dst, item_brand_src),
+        ("item", "ic", "category"): (item_category_src, item_category_dst),
+        ("category", "ci", "item"): (item_category_dst, item_category_src),
+    }
+    nodes_dict = {
+        "user": hg.num_nodes("user"),
+        "item": hg.num_nodes("item"),
+        "view": hg.num_nodes("view"),
+        "brand": hg.num_nodes("brand"),
+        "category": hg.num_nodes("category"),
+    }
+    hg_processed = dgl.heterograph(
+        data_dict=edges_dict, num_nodes_dict=nodes_dict
+    )
+    print("Graph processed.")
+
+    # save the processed data
+    with open(os.path.join(root_path, "amazon_hg.pkl"), "wb") as file:
+        pkl.dump(hg_processed, file)
+    with open(os.path.join(root_path, "amazon_train.pkl"), "wb") as file:
+        pkl.dump(train_data, file)
+    with open(os.path.join(root_path, "amazon_test.pkl"), "wb") as file:
+        pkl.dump(test_data, file)
+    with open(os.path.join(root_path, "amazon_eval.pkl"), "wb") as file:
+        pkl.dump(eval_data, file)
+
+    return hg_processed, train_data, eval_data, test_data
+
+
+def process_movielens(root_path):
+    # User-Movie 943 1682 100000 UMUM
+    # User-Age 943 8 943 UAUM
+    # User-Occupation 943 21 943 UOUM
+    # Movie-Genre 1682 18 2861 UMGM
+
+    data_path = os.path.join(root_path, "Movielens")
+    if not (os.path.exists(data_path)):
+        print(
+            "Can not find movielens in {}, please download the dataset first.".format(
+                data_path
+            )
+        )
+
+    # Construct graph from raw data.
+    # movie_genre
+    movie_genre_src = []
+    movie_genre_dst = []
+    with open(os.path.join(data_path, "movie_genre.dat")) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split("\t")
+            movie, genre = int(_line[0]), int(_line[1])
+            movie_genre_src.append(movie)
+            movie_genre_dst.append(genre)
+
+    # user_movie
+    user_movie_src = []
+    user_movie_dst = []
+    with open(os.path.join(data_path, "user_movie.dat")) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split("\t")
+            user, item, rate = int(_line[0]), int(_line[1]), int(_line[2])
+            if rate > 3:
+                user_movie_src.append(user)
+                user_movie_dst.append(item)
+
+    # user_occupation
+    user_occupation_src = []
+    user_occupation_dst = []
+    with open(os.path.join(data_path, "user_occupation.dat")) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split("\t")
+            user, occupation = int(_line[0]), int(_line[1])
+            user_occupation_src.append(user)
+            user_occupation_dst.append(occupation)
+
+    # user_age
+    user_age_src = []
+    user_age_dst = []
+    with open(os.path.join(data_path, "user_age.dat")) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split("\t")
+            user, age = int(_line[0]), int(_line[1])
+            user_age_src.append(user)
+            user_age_dst.append(age)
+
+    # build graph
+    hg = dgl.heterograph(
+        {
+            ("movie", "mg", "genre"): (movie_genre_src, movie_genre_dst),
+            ("genre", "gm", "movie"): (movie_genre_dst, movie_genre_src),
+            ("user", "um", "movie"): (user_movie_src, user_movie_dst),
+            ("movie", "mu", "user"): (user_movie_dst, user_movie_src),
+            ("user", "uo", "occupation"): (
+                user_occupation_src,
+                user_occupation_dst,
+            ),
+            ("occupation", "ou", "user"): (
+                user_occupation_dst,
+                user_occupation_src,
+            ),
+            ("user", "ua", "age"): (user_age_src, user_age_dst),
+            ("age", "au", "user"): (user_age_dst, user_age_src),
+        }
+    )
+
+    print("Graph constructed.")
+
+    # Split data into train/eval/test
+    train_data, eval_data, test_data = split_data(hg, "um")
+
+    # delete the positive edges in eval/test data in the original graph
+    train_pos = np.nonzero(train_data[:, 2])
+    train_pos_idx = train_pos[0]
+    user_movie_src_processed = train_data[train_pos_idx, 0]
+    user_movie_dst_processed = train_data[train_pos_idx, 1]
+    edges_dict = {
+        ("movie", "mg", "genre"): (movie_genre_src, movie_genre_dst),
+        ("genre", "gm", "movie"): (movie_genre_dst, movie_genre_src),
+        ("user", "um", "movie"): (
+            user_movie_src_processed,
+            user_movie_dst_processed,
+        ),
+        ("movie", "mu", "user"): (
+            user_movie_dst_processed,
+            user_movie_src_processed,
+        ),
+        ("user", "uo", "occupation"): (
+            user_occupation_src,
+            user_occupation_dst,
+        ),
+        ("occupation", "ou", "user"): (
+            user_occupation_dst,
+            user_occupation_src,
+        ),
+        ("user", "ua", "age"): (user_age_src, user_age_dst),
+        ("age", "au", "user"): (user_age_dst, user_age_src),
+    }
+    nodes_dict = {
+        "user": hg.num_nodes("user"),
+        "movie": hg.num_nodes("movie"),
+        "genre": hg.num_nodes("genre"),
+        "occupation": hg.num_nodes("occupation"),
+        "age": hg.num_nodes("age"),
+    }
+    hg_processed = dgl.heterograph(
+        data_dict=edges_dict, num_nodes_dict=nodes_dict
+    )
+    print("Graph processed.")
+
+    # save the processed data
+    with open(os.path.join(root_path, "movielens_hg.pkl"), "wb") as file:
+        pkl.dump(hg_processed, file)
+    with open(os.path.join(root_path, "movielens_train.pkl"), "wb") as file:
+        pkl.dump(train_data, file)
+    with open(os.path.join(root_path, "movielens_test.pkl"), "wb") as file:
+        pkl.dump(test_data, file)
+    with open(os.path.join(root_path, "movielens_eval.pkl"), "wb") as file:
+        pkl.dump(eval_data, file)
+
+    return hg_processed, train_data, eval_data, test_data
+
+
+class MyDataset(Dataset):
+    def __init__(self, triple):
+        self.triple = triple
+        self.len = self.triple.shape[0]
+
+    def __getitem__(self, index):
+        return (
+            self.triple[index, 0],
+            self.triple[index, 1],
+            self.triple[index, 2].float(),
+        )
+
+    def __len__(self):
+        return self.len
+
+
+def load_data(dataset, batch_size=128, num_workers=10, root_path="./data"):
+    if os.path.exists(os.path.join(root_path, dataset + "_train.pkl")):
+        g_file = open(os.path.join(root_path, dataset + "_hg.pkl"), "rb")
+        hg = pkl.load(g_file)
+        g_file.close()
+        train_set_file = open(
+            os.path.join(root_path, dataset + "_train.pkl"), "rb"
+        )
+        train_set = pkl.load(train_set_file)
+        train_set_file.close()
+        test_set_file = open(
+            os.path.join(root_path, dataset + "_test.pkl"), "rb"
+        )
+        test_set = pkl.load(test_set_file)
+        test_set_file.close()
+        eval_set_file = open(
+            os.path.join(root_path, dataset + "_eval.pkl"), "rb"
+        )
+        eval_set = pkl.load(eval_set_file)
+        eval_set_file.close()
+    else:
+        if dataset == "movielens":
+            hg, train_set, eval_set, test_set = process_movielens(root_path)
+        elif dataset == "amazon":
+            hg, train_set, eval_set, test_set = process_amazon(root_path)
+        else:
+            print("Available datasets: movielens, amazon.")
+            raise NotImplementedError
+
+    if dataset == "movielens":
+        meta_paths = {
+            "user": [["um", "mu"]],
+            "movie": [["mu", "um"], ["mg", "gm"]],
+        }
+        user_key = "user"
+        item_key = "movie"
+    elif dataset == "amazon":
+        meta_paths = {
+            "user": [["ui", "iu"]],
+            "item": [["iu", "ui"], ["ic", "ci"], ["ib", "bi"], ["iv", "vi"]],
+        }
+        user_key = "user"
+        item_key = "item"
+    else:
+        print("Available datasets: movielens, amazon.")
+        raise NotImplementedError
+
+    train_set = torch.Tensor(train_set).long()
+    eval_set = torch.Tensor(eval_set).long()
+    test_set = torch.Tensor(test_set).long()
+
+    train_set = MyDataset(train_set)
+    train_loader = DataLoader(
+        dataset=train_set,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+    )
+    eval_set = MyDataset(eval_set)
+    eval_loader = DataLoader(
+        dataset=eval_set,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+    )
+    test_set = MyDataset(test_set)
+    test_loader = DataLoader(
+        dataset=test_set,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+    )
+
+    return (
+        hg,
+        train_loader,
+        eval_loader,
+        test_loader,
+        meta_paths,
+        user_key,
+        item_key,
+    )

examples/pytorch/TAHIN/main.py

@@ -0,0 +1,221 @@
+import argparse
+import pickle as pkl
+
+import dgl
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from data_loader import load_data
+from TAHIN import TAHIN
+from utils import (
+    evaluate_acc,
+    evaluate_auc,
+    evaluate_f1_score,
+    evaluate_logloss,
+)
+
+
+def main(args):
+    # step 1: Check device
+    if args.gpu >= 0 and torch.cuda.is_available():
+        device = "cuda:{}".format(args.gpu)
+    else:
+        device = "cpu"
+
+    # step 2: Load data
+    (
+        g,
+        train_loader,
+        eval_loader,
+        test_loader,
+        meta_paths,
+        user_key,
+        item_key,
+    ) = load_data(args.dataset, args.batch, args.num_workers, args.path)
+    g = g.to(device)
+    print("Data loaded.")
+
+    # step 3: Create model and training components
+    model = TAHIN(
+        g, meta_paths, args.in_size, args.out_size, args.num_heads, args.dropout
+    )
+    model = model.to(device)
+    criterion = nn.BCELoss()
+    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
+    print("Model created.")
+
+    # step 4: Training
+    print("Start training.")
+    best_acc = 0.0
+    kill_cnt = 0
+    for epoch in range(args.epochs):
+        # Training and validation using a full graph
+        model.train()
+        train_loss = []
+        for step, batch in enumerate(train_loader):
+            user, item, label = [_.to(device) for _ in batch]
+            logits = model.forward(g, user_key, item_key, user, item)
+
+            # compute loss
+            tr_loss = criterion(logits, label)
+            train_loss.append(tr_loss)
+
+            # backward
+            optimizer.zero_grad()
+            tr_loss.backward()
+            optimizer.step()
+
+        train_loss = torch.stack(train_loss).sum().cpu().item()
+
+        model.eval()
+        with torch.no_grad():
+            validate_loss = []
+            validate_acc = []
+            for step, batch in enumerate(eval_loader):
+                user, item, label = [_.to(device) for _ in batch]
+                logits = model.forward(g, user_key, item_key, user, item)
+
+                # compute loss
+                val_loss = criterion(logits, label)
+                val_acc = evaluate_acc(
+                    logits.detach().cpu().numpy(), label.detach().cpu().numpy()
+                )
+                validate_loss.append(val_loss)
+                validate_acc.append(val_acc)
+
+            validate_loss = torch.stack(validate_loss).sum().cpu().item()
+            validate_acc = np.mean(validate_acc)
+
+            # validate
+            if validate_acc > best_acc:
+                best_acc = validate_acc
+                best_epoch = epoch
+                torch.save(model.state_dict(), "TAHIN" + "_" + args.dataset)
+                kill_cnt = 0
+                print("saving model...")
+            else:
+                kill_cnt += 1
+                if kill_cnt > args.early_stop:
+                    print("early stop.")
+                    print("best epoch:{}".format(best_epoch))
+                    break
+
+            print(
+                "In epoch {}, Train Loss: {:.4f}, Valid Loss: {:.5}\n, Valid ACC: {:.5}".format(
+                    epoch, train_loss, validate_loss, validate_acc
+                )
+            )
+
+    # test use the best model
+    model.eval()
+    with torch.no_grad():
+        model.load_state_dict(
+            torch.load("TAHIN" + "_" + args.dataset, weights_only=False)
+        )
+        test_loss = []
+        test_acc = []
+        test_auc = []
+        test_f1 = []
+        test_logloss = []
+        for step, batch in enumerate(test_loader):
+            user, item, label = [_.to(device) for _ in batch]
+            logits = model.forward(g, user_key, item_key, user, item)
+
+            # compute loss
+            loss = criterion(logits, label)
+            acc = evaluate_acc(
+                logits.detach().cpu().numpy(), label.detach().cpu().numpy()
+            )
+            auc = evaluate_auc(
+                logits.detach().cpu().numpy(), label.detach().cpu().numpy()
+            )
+            f1 = evaluate_f1_score(
+                logits.detach().cpu().numpy(), label.detach().cpu().numpy()
+            )
+            log_loss = evaluate_logloss(
+                logits.detach().cpu().numpy(), label.detach().cpu().numpy()
+            )
+
+            test_loss.append(loss)
+            test_acc.append(acc)
+            test_auc.append(auc)
+            test_f1.append(f1)
+            test_logloss.append(log_loss)
+
+        test_loss = torch.stack(test_loss).sum().cpu().item()
+        test_acc = np.mean(test_acc)
+        test_auc = np.mean(test_auc)
+        test_f1 = np.mean(test_f1)
+        test_logloss = np.mean(test_logloss)
+        print(
+            "Test Loss: {:.5}\n, Test ACC: {:.5}\n, AUC: {:.5}\n, F1: {:.5}\n, Logloss: {:.5}\n".format(
+                test_loss, test_acc, test_auc, test_f1, test_logloss
+            )
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Parser For Arguments",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--dataset",
+        default="movielens",
+        help="Dataset to use, default: movielens",
+    )
+    parser.add_argument(
+        "--path", default="./data", help="Path to save the data"
+    )
+    parser.add_argument("--model", default="TAHIN", help="Model Name")
+
+    parser.add_argument("--batch", default=128, type=int, help="Batch size")
+    parser.add_argument(
+        "--gpu",
+        type=int,
+        default="0",
+        help="Set GPU Ids : Eg: For CPU = -1, For Single GPU = 0",
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=500, help="Maximum number of epochs"
+    )
+    parser.add_argument(
+        "--wd", type=float, default=0, help="L2 Regularization for Optimizer"
+    )
+    parser.add_argument("--lr", type=float, default=0.001, help="Learning Rate")
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=10,
+        help="Number of processes to construct batches",
+    )
+    parser.add_argument(
+        "--early_stop", default=15, type=int, help="Patience for early stop."
+    )
+
+    parser.add_argument(
+        "--in_size",
+        default=128,
+        type=int,
+        help="Initial dimension size for entities.",
+    )
+    parser.add_argument(
+        "--out_size",
+        default=128,
+        type=int,
+        help="Output dimension size for entities.",
+    )
+
+    parser.add_argument(
+        "--num_heads", default=1, type=int, help="Number of attention heads"
+    )
+    parser.add_argument("--dropout", default=0.1, type=float, help="Dropout.")
+
+    args = parser.parse_args()
+
+    print(args)
+
+    main(args)

examples/pytorch/TAHIN/readme.md

@@ -0,0 +1,84 @@
+# DGL Implementation of the TAHIN
+
+This DGL example implements the TAHIN module proposed in the paper [HCDIR](https://arxiv.org/pdf/2007.15293.pdf). Since the code and dataset have not been published yet, we implement its main idea and experiment on two other datasets.
+
+Example implementor
+----------------------
+This example was implemented by [KounianhuaDu](https://github.com/KounianhuaDu) during her software development intern time at the AWS Shanghai AI Lab.
+
+Dependencies
+----------------------
+- pytorch 1.7.1
+- dgl 0.6.0
+- scikit-learn 0.22.1
+
+Datasets
+---------------------------------------
+The datasets used can be downloaded from [here](https://github.com/librahu/HIN-Datasets-for-Recommendation-and-Network-Embedding). For the experiments, all the positive edges are fetched and the same number of negative edges are randomly sampled. The edges are then shuffled and splitted into train/validate/test at a ratio of 6:2:2. The positive edges that appear in the validation and test sets are then removed from the original graph.
+
+The original graph statistics:
+
+**Movielens** 
+
+(Source : https://grouplens.org/datasets/movielens/)
+
+| Entity         |#Entity        |
+| :-------------:|:-------------:|
+| User           | 943           |
+| Age            | 8             |
+| Occupation     | 21            |
+| Movie          | 1,682         |
+| Genre          | 18            |
+
+| Relation            |#Relation      |
+| :-------------:     |:-------------:|
+| User - Movie        | 100,000       |
+| User - User (KNN)   | 47,150        |
+| User - Age          | 943           |
+| User - Occupation   | 943           |
+| Movie - Movie (KNN) | 82,798        |
+| Movie - Genre       | 2,861         |
+
+**Amazon** 
+
+(Source : http://jmcauley.ucsd.edu/data/amazon/)
+
+| Entity         |#Entity        |
+| :-------------:|:-------------:|
+| User           | 6,170         |
+| Item           | 2,753         |
+| View           | 3,857         |           
+| Category       | 22            |
+| Brand          | 334           |
+
+| Relation          |#Relation      |
+| :-------------:   |:-------------:|
+| User - Item       | 195,791       |
+| Item - View       | 5,694         |
+| Item - Category   | 5,508         | 
+| Item - Brand      | 2,753         |
+
+How to run
+--------------------------------
+
+```python
+python main.py --dataset amazon --gpu 0
+```
+
+
+```python
+python main.py --dataset movielens --gpu 0
+```
+
+
+Performance
+-------------------------
+**Results**
+
+| Dataset |         Movielens        |          Amazon          |
+|---------| ------------------------ | ------------------------ |
+|  Metric |    HAN     /    TAHIN    |    HAN     /    TAHIN    |
+|   AUC   |   0.9297   /   0.9392    |   0.8470   /   0.8442    |
+|   ACC   |   0.8627   /   0.8683    |   0.7672   /   0.7619    |
+|    F1   |   0.8631   /   0.8707    |   0.7628   /   0.7499    |
+| Logloss |   0.3689   /   0.3266    |   0.5311   /   0.5150    |

examples/pytorch/TAHIN/utils.py

@@ -0,0 +1,25 @@
+from sklearn.metrics import accuracy_score, f1_score, log_loss, roc_auc_score
+
+
+def evaluate_auc(pred, label):
+    res = roc_auc_score(y_score=pred, y_true=label)
+    return res
+
+
+def evaluate_acc(pred, label):
+    res = []
+    for _value in pred:
+        res.append(1 if _value >= 0.5 else 0)
+    return accuracy_score(y_pred=res, y_true=label)
+
+
+def evaluate_f1_score(pred, label):
+    res = []
+    for _value in pred:
+        res.append(1 if _value >= 0.5 else 0)
+    return f1_score(y_pred=res, y_true=label)
+
+
+def evaluate_logloss(pred, label):
+    res = log_loss(y_true=label, y_pred=pred, normalize=True)
+    return res

examples/pytorch/appnp/README.md

@@ -0,0 +1,32 @@
+Predict then Propagate: Graph Neural Networks meet Personalized PageRank (APPNP)
+============
+
+- Paper link: [Predict then Propagate: Graph Neural Networks meet Personalized PageRank](https://arxiv.org/abs/1810.05997)
+- Author's code repo: [https://github.com/klicperajo/ppnp](https://github.com/klicperajo/ppnp). 
+
+Dependencies
+------------
+- PyTorch 0.4.1+
+- requests
+
+``bash
+pip install torch requests
+``
+
+Code
+-----
+The folder contains an implementation of APPNP (`appnp.py`).
+
+Results
+-------
+
+Run with following (available dataset: "cora", "citeseer", "pubmed")
+```bash
+python3 train.py --dataset cora --gpu 0
+```
+
+* cora: 0.8370 (paper: 0.850)
+* citeseer: 0.715 (paper: 0.757)
+* pubmed: 0.793 (paper: 0.797)
+
+Experiments were done on dgl datasets (GCN settings) which are different from those used in the original implementation. (discrepancies are detailed in experimental section of the original paper)

examples/pytorch/appnp/appnp.py

@@ -0,0 +1,58 @@
+"""
+APPNP implementation in DGL.
+References
+----------
+Paper: https://arxiv.org/abs/1810.05997
+Author's code: https://github.com/klicperajo/ppnp
+"""
+import torch.nn as nn
+
+from dgl.nn.pytorch.conv import APPNPConv
+
+
+class APPNP(nn.Module):
+    def __init__(
+        self,
+        g,
+        in_feats,
+        hiddens,
+        n_classes,
+        activation,
+        feat_drop,
+        edge_drop,
+        alpha,
+        k,
+    ):
+        super(APPNP, self).__init__()
+        self.g = g
+        self.layers = nn.ModuleList()
+        # input layer
+        self.layers.append(nn.Linear(in_feats, hiddens[0]))
+        # hidden layers
+        for i in range(1, len(hiddens)):
+            self.layers.append(nn.Linear(hiddens[i - 1], hiddens[i]))
+        # output layer
+        self.layers.append(nn.Linear(hiddens[-1], n_classes))
+        self.activation = activation
+        if feat_drop:
+            self.feat_drop = nn.Dropout(feat_drop)
+        else:
+            self.feat_drop = lambda x: x
+        self.propagate = APPNPConv(k, alpha, edge_drop)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for layer in self.layers:
+            layer.reset_parameters()
+
+    def forward(self, features):
+        # prediction step
+        h = features
+        h = self.feat_drop(h)
+        h = self.activation(self.layers[0](h))
+        for layer in self.layers[1:-1]:
+            h = self.activation(layer(h))
+        h = self.layers[-1](self.feat_drop(h))
+        # propagation step
+        h = self.propagate(self.g, h)
+        return h

examples/pytorch/appnp/train.py

@@ -0,0 +1,165 @@
+import argparse
+import time
+
+import dgl
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from appnp import APPNP
+from dgl.data import (
+    CiteseerGraphDataset,
+    CoraGraphDataset,
+    PubmedGraphDataset,
+    register_data_args,
+)
+
+
+def evaluate(model, features, labels, mask):
+    model.eval()
+    with torch.no_grad():
+        logits = model(features)
+        logits = logits[mask]
+        labels = labels[mask]
+        _, indices = torch.max(logits, dim=1)
+        correct = torch.sum(indices == labels)
+        return correct.item() * 1.0 / len(labels)
+
+
+def main(args):
+    # load and preprocess dataset
+    if args.dataset == "cora":
+        data = CoraGraphDataset()
+    elif args.dataset == "citeseer":
+        data = CiteseerGraphDataset()
+    elif args.dataset == "pubmed":
+        data = PubmedGraphDataset()
+    else:
+        raise ValueError("Unknown dataset: {}".format(args.dataset))
+
+    g = data[0]
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        g = g.to(args.gpu)
+
+    features = g.ndata["feat"]
+    labels = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+    in_feats = features.shape[1]
+    n_classes = data.num_classes
+    n_edges = g.num_edges()
+    print(
+        """----Data statistics------'
+      #Edges %d
+      #Classes %d
+      #Train samples %d
+      #Val samples %d
+      #Test samples %d"""
+        % (
+            n_edges,
+            n_classes,
+            train_mask.int().sum().item(),
+            val_mask.int().sum().item(),
+            test_mask.int().sum().item(),
+        )
+    )
+
+    n_edges = g.num_edges()
+    # add self loop
+    g = dgl.remove_self_loop(g)
+    g = dgl.add_self_loop(g)
+
+    # create APPNP model
+    model = APPNP(
+        g,
+        in_feats,
+        args.hidden_sizes,
+        n_classes,
+        F.relu,
+        args.in_drop,
+        args.edge_drop,
+        args.alpha,
+        args.k,
+    )
+
+    if cuda:
+        model.cuda()
+    loss_fcn = torch.nn.CrossEntropyLoss()
+
+    # use optimizer
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+    )
+
+    # initialize graph
+    mean = 0
+    for epoch in range(args.n_epochs):
+        model.train()
+        if epoch >= 3:
+            t0 = time.time()
+        # forward
+        logits = model(features)
+        loss = loss_fcn(logits[train_mask], labels[train_mask])
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if epoch >= 3:
+            mean = (mean * (epoch - 3) + (time.time() - t0)) / (epoch - 2)
+            acc = evaluate(model, features, labels, val_mask)
+            print(
+                "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
+                "ETputs(KTEPS) {:.2f}".format(
+                    epoch,
+                    mean,
+                    loss.item(),
+                    acc,
+                    n_edges / mean / 1000,
+                )
+            )
+
+    print()
+    acc = evaluate(model, features, labels, test_mask)
+    print("Test Accuracy {:.4f}".format(acc))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="APPNP")
+    register_data_args(parser)
+    parser.add_argument(
+        "--in-drop", type=float, default=0.5, help="input feature dropout"
+    )
+    parser.add_argument(
+        "--edge-drop", type=float, default=0.5, help="edge propagation dropout"
+    )
+    parser.add_argument("--gpu", type=int, default=-1, help="gpu")
+    parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
+    parser.add_argument(
+        "--n-epochs", type=int, default=200, help="number of training epochs"
+    )
+    parser.add_argument(
+        "--hidden_sizes",
+        type=int,
+        nargs="+",
+        default=[64],
+        help="hidden unit sizes for appnp",
+    )
+    parser.add_argument(
+        "--k", type=int, default=10, help="Number of propagation steps"
+    )
+    parser.add_argument(
+        "--alpha", type=float, default=0.1, help="Teleport Probability"
+    )
+    parser.add_argument(
+        "--weight-decay", type=float, default=5e-4, help="Weight for L2 loss"
+    )
+    args = parser.parse_args()
+    print(args)
+
+    main(args)

examples/pytorch/argo/README.md

@@ -0,0 +1,158 @@
+# ARGO: An Auto-Tuning Runtime System for Scalable GNN Training on Multi-Core Processor
+
+## Overview
+
+Graph Neural Network (GNN) training suffers from low scalability on multi-core processors. 
+ARGO is a runtime system that offers scalable performance. 
+The figure below shows an example of GNN training on a Xeon 8380H platform with 112 cores. 
+Without ARGO, there is no performance improvement after applying more than 16 cores; we observe a similar scalability limit on a Xeon 6430L platform with 64 cores as well.
+However, with ARGO enabled, we are able to scale over 64 cores, allowing ARGO to speedup GNN training (in terms of epoch time) by up to 4.30x and 3.32x on a Xeon 8380H and a Xeon 6430L, respectively.
+![ARGO](https://github.com/dmlc/dgl/blob/master/examples/pytorch/argo/argo_scale.png)
+
+This README includes how to:
+1. [Installation](#1-installation)
+2. [Run the example code](#2-running-the-example-GNN-program)
+3. [Modify your own GNN program to enable ARGO.](#3-enabling-ARGO-on-your-own-GNN-program)
+
+## 1. Installation
+
+1. ARGO utilizes the scikit-optimize library for auto-tuning. Please install scikit-optimize to run ARGO:
+
+   ```shell
+   conda install -c conda-forge "scikit-optimize>=0.9.0" 
+   ```
+   or
+   ```shell
+   pip install scikit-optimize>=0.9
+   ```
+
+## 2. Running the example GNN program
+### Usage
+  ```shell
+  python main.py --dataset ogbn-products --sampler shadow --model sage
+  ``` 
+  Important Arguments: 
+  - `--dataset`: the training datasets. Available choices [ogbn-products, ogbn-papers100M, reddit, flickr, yelp]
+  - `--sampler`: the mini-batch sampling algorithm. Available choices [shadow, neighbor]
+  - `--model`: GNN model. Available choices [gcn, sage]
+  - `--layer`: number of GNN layers.
+  - `--fan_out`: number of fanout neighbors for each layer.
+  - `--hidden`: hidden feature dimension.
+  - `--batch_size`: the size of the mini-batch.
+
+
+
+## 3. Enabling ARGO on your own GNN program
+
+In this section, we provide a step-by-step tutorial on how to enable ARGO on a DGL program. We use the ```ogb_example.py``` file in this repo as an example.  
+
+>  Note: we also provide the complete example file ```ogb_example_ARGO.py``` which followed the steps below to enable ARGO on ```ogb_example.py```.
+
+1. First, include all necessary packages on top of the file. Please place your file and ```argo.py``` in the same directory.
+
+   ```python
+   import os
+   import torch.distributed as dist
+   from torch.nn.parallel import DistributedDataParallel
+   import torch.multiprocessing as mp
+   from argo import ARGO
+   ```
+
+2. Setup PyTorch Distributed Data Parallel (DDP). 
+    1. Add the initialization function on top of the training program, and wrap the ```model``` with the DDP wrapper
+     ```python
+     def train(...):
+       dist.init_process_group('gloo', rank=rank, world_size=world_size) # newly added
+       model = SAGE(...) # original code
+       model = DistributedDataParallel(model) # newly added
+       ...
+     ```
+    2. In the main program, add the following before launching the training function
+    
+     ```python
+     os.environ['MASTER_ADDR'] = '127.0.0.1'
+     os.environ['MASTER_PORT'] = '29501'
+     mp.set_start_method('fork', force=True)
+     train(args, device, data) # original code for launching the training function
+     ```
+
+3. Enable ARGO by initializing the runtime system, and wrapping the training function
+   ```python
+   runtime = ARGO(n_search = 15, epoch = args.num_epochs, batch_size = args.batch_size) #initialization
+   runtime.run(train, args=(args, device, data)) # wrap the training function
+   ```
+   >  ARGO takes three input paramters: number of searches ```n_search```, number of epochs, and the mini-batch size. Increasing ```n_search``` potentially leads to a better configuration with less epoch time; however, searching itself also causes extra overhead. We recommend setting ```n_search``` from 15 to 45 for an optimal overall performance. Details of ```n_search``` can be found in the paper.
+
+4. Modify the input of the training function, by directly adding ARGO parameters after the original inputs.
+   This is the original function:
+   ```python
+   def train(args, device, data):
+   ```
+   Add ```rank, world_size, comp_core, load_core, counter, b_size, ep``` like this:
+   ```python
+   def train(args, device, data, rank, world_size, comp_core, load_core, counter, b_size, ep):
+   ```
+
+5. Modify the ```dataloader``` function in the training function
+   ```python
+   dataloader = dgl.dataloading.DataLoader(
+           g,
+           train_nid,
+           sampler,
+           batch_size=b_size, # modified
+           shuffle=True,
+           drop_last=False,
+           num_workers=len(load_core), # modified
+           use_ddp = True) # newly added
+   ```
+
+6. Enable core-binding by adding ```enable_cpu_affinity()``` before the training for-loop, and also change the number of epochs into the variable ```ep```: 
+   ```python
+   with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core): 
+     for epoch in range(ep): # change num_epochs to ep
+   ```
+
+7. Last step! Load the model before training and save it afterward.  
+   Original Program:
+   ```python
+   with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core): 
+     for epoch in range(ep): 
+       ... # training operations
+   ```
+   Modified:
+   ```python
+   PATH = "model.pt"
+   if counter[0] != 0:
+     checkpoint = th.load(PATH)
+     model.load_state_dict(checkpoint['model_state_dict'])
+     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+     epoch = checkpoint['epoch']
+     loss = checkpoint['loss']
+   
+   with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core): 
+     for epoch in range(ep): 
+       ... # training operations
+   
+   dist.barrier()
+   if rank == 0:
+     th.save({'epoch': counter[0],
+                 'model_state_dict': model.state_dict(),
+                 'optimizer_state_dict': optimizer.state_dict(),
+                 'loss': loss,
+                 }, PATH)
+   
+   ```
+8. Done! You can now run your GNN program with ARGO enabled.
+      ```shell
+      python <your_code>.py
+      ```
+
+## Citation & Acknowledgement
+This work has been supported by the U.S. National Science Foundation (NSF) under grants CCF-1919289/SPX-2333009, CNS-2009057 and OAC-2209563, and the Semiconductor Research Corporation (SRC).
+```
+@INPROCEEDINGS{argo-ipdps24,
+  author={Yi-Chien Lin and Yuyang Chen and Sameh Gobriel and Nilesh Jain and Gopi Krishna Jhaand and Viktor Prasanna},
+  booktitle={IEEE International Parallel and Distributed Processing Symposium (IPDPS)}, 
+  title={ARGO: An Auto-Tuning Runtime System for Scalable GNN Training on Multi-Core Processor}, 
+  year={2024}}
+```

examples/pytorch/argo/argo.py

@@ -0,0 +1,270 @@
+"""
+ARGO: An Auto-Tuning Runtime System for Scalable GNN Training on Multi-Core Processor
+--------------------------------------------
+Graph Neural Network (GNN) training suffers from low scalability on multi-core CPUs. 
+Specificially, the performance often caps at 16 cores, and no improvement is observed when applying more than 16 cores.
+ARGO is a runtime system that offers scalable performance by overlapping the computation and communication during GNN training.
+With ARGO enabled, we are able to scale over 64 cores, allowing ARGO to speedup GNN training (in terms of epoch time) by up to 4.30x and 3.32x on a Xeon 8380H and a Xeon 6430L, respectively.
+--------------------------------------------
+Paper Link: https://arxiv.org/abs/2402.03671
+"""
+
+import time
+from typing import Callable, List, Tuple
+
+import dgl.multiprocessing as dmp
+import numpy as np
+import psutil
+from skopt import gp_minimize
+from skopt.space import Normalize
+
+
+def transform(self, X):
+    X = np.asarray(X)
+    if self.is_int:
+        if np.any(np.round(X) > self.high):
+            raise ValueError(
+                "All integer values should" "be less than %f" % self.high
+            )
+        if np.any(np.round(X) < self.low):
+            raise ValueError(
+                "All integer values should" "be greater than %f" % self.low
+            )
+    else:
+        if np.any(X > self.high + self._eps):
+            raise ValueError("All values should" "be less than %f" % self.high)
+        if np.any(X < self.low - self._eps):
+            raise ValueError(
+                "All values should" "be greater than %f" % self.low
+            )
+    if (self.high - self.low) == 0.0:
+        return X * 0.0
+    if self.is_int:
+        return (np.round(X).astype(int) - self.low) / (self.high - self.low)
+    else:
+        return (X - self.low) / (self.high - self.low)
+
+
+def inverse_transform(self, X):
+    X = np.asarray(X)
+    if np.any(X > 1.0 + self._eps):
+        raise ValueError("All values should be less than 1.0")
+    if np.any(X < 0.0 - self._eps):
+        raise ValueError("All values should be greater than 0.0")
+    X_orig = X * (self.high - self.low) + self.low
+    if self.is_int:
+        return np.round(X_orig).astype(int)
+    return X_orig
+
+
+# This is a workaround for scikit-optimize's incompatibility with NumPy, which results in an error::
+# AttributeError: module 'numpy' has no attribute 'int'
+Normalize.transform = transform
+Normalize.inverse_transform = inverse_transform
+
+
+class ARGO:
+    def __init__(
+        self,
+        n_search=10,
+        epoch=200,
+        batch_size=4096,
+        space=[(2, 8), (1, 4), (1, 32)],
+        random_state=1,
+    ):
+        """
+        Initialization
+
+        Parameters
+        ----------
+        n_search: int
+            Number of configuration searches the auto-tuner will conduct
+
+        epoch: int
+            Number of epochs of GNN training
+
+        batch_size: int
+            Size of the mini-batch
+
+        space: list[Tuple(int,int)]
+            Range of the search space; [range of processes, range of samplers for each process, range of trainers for each process]
+
+        random_state: int
+            Number of random initializations before searching
+
+        """
+        self.n_search = n_search
+        self.epoch = epoch
+        self.batch_size = batch_size
+        self.space = space
+        self.random_state = random_state
+        self.acq_func = "EI"
+        self.counter = [0]
+
+    def core_binder(
+        self, num_cpu_proc: int, n_samp: int, n_train: int, rank: int
+    ) -> Tuple[List[int], List[int]]:
+        """
+        Core Binder
+
+        The Core Binder binds CPU cores to perform sampling (i.e., sampling cores) and model propagation (i.e., training cores).
+        The actual binding is done using the CPU affinity function in the data_loader.
+        The core_binder function here is used to produce the list of CPU IDs for the CPU affinity function.
+
+        Parameters
+        ----------
+        num_cpu_proc: int
+            Number of processes instantiated
+
+        n_samp: int
+            Number of sampling cores for each process
+
+        n_train: int
+            Number of training cores for each process
+
+        rank: int
+            The rank of the current process
+
+        Returns: Tuple[list[int], list[int]]
+        -------
+        load_core: list[int]
+            For a given process rank, the load_core specifies a list of CPU core IDs to be used for sampling, the length of load_core = n_samp.
+
+        comp_core: list[int]
+            For a given process rank, the comp_core specifies a list of CPU core IDs to be used for training, the length of comp_core = n_comp.
+
+        .. note:: Each process is assigned with a unique list of sampling cores and training cores, and no CPU core will appear in two lists or more.
+
+        """
+        load_core, comp_core = [], []
+        n = psutil.cpu_count(logical=False)
+        size = num_cpu_proc
+        num_of_samplers = n_samp
+        load_core = list(
+            range(n // size * rank, n // size * rank + num_of_samplers)
+        )
+        comp_core = list(
+            range(
+                n // size * rank + num_of_samplers,
+                n // size * rank + num_of_samplers + n_train,
+            )
+        )
+        return load_core, comp_core
+
+    def auto_tuning(self, train: Callable, args) -> List[int]:
+        """
+        Auto-tuner
+
+        The auto-tuner runs Bayesian Optimization (BO) to search for the optimal configuration (number of processes, samplers, trainers).
+        During the search, the auto-tuner explores the design space by collecting the epoch time of various configurations.
+        Specifically, the exploration is done by feeding the Multi-Process Engine with various configurations, and record the epoch time.
+        After the searching is done, the optimal configuration will be used repeatedly until the end of model training.
+
+        Parameters
+        ----------
+        train: Callable
+            The GNN training function.
+
+        args:
+            The inputs of the GNN training function.
+
+        Returns
+        -------
+        result: list[int]
+            The optimal configurations (which leads to the shortest epoch time) found by running BO.
+            - result[0]: number of processes to instantiate
+            - result[1]: number of sampling cores for each process
+            - result[2]: number of training cores for each process
+
+        """
+        ep = 1
+        result = gp_minimize(
+            lambda x: self.mp_engine(x, train, args, ep),
+            dimensions=self.space,
+            n_calls=self.n_search,
+            random_state=self.random_state,
+            acq_func=self.acq_func,
+        )
+        return result
+
+    def mp_engine(self, x: List[int], train: Callable, args, ep: int) -> float:
+        """
+        Multi-Process Engine (MP Engine)
+
+        The MP Engine launches multiple GNN training processes in parallel to overlap computation with communication.
+        Such an approach effectively improves the utilization of the memory bandwidth and the CPU cores.
+        The MP Engine also adjust the batch size according to the number of processes instantiated, so that the effective batch size remains the same as the original program without ARGO.
+
+        Parameters
+        ----------
+        x: list[int]
+            Optimal configurations provided by the auto-tuner.
+            - x[0]: number of processes to instantiate
+            - x[1]: number of sampling cores for each process
+            - x[2]: number of training cores for each process
+
+        train: Callable
+            The GNN training function.
+
+        args:
+            The inputs of the GNN training function.
+
+        ep: int
+            number of epochs.
+
+        Returns
+        -------
+        t: float
+            The epoch time using the current configuration `x`.
+        """
+        n_proc = x[0]
+        n_samp = x[1]
+        n_train = x[2]
+        n_total = psutil.cpu_count(logical=False)
+
+        if n_proc * (n_samp + n_train) > n_total:  # handling corner cases
+            n_proc = 2
+            n_samp = 2
+            n_train = (n_total // n_proc) - n_samp
+
+        processes = []
+        cnt = self.counter
+        b_size = self.batch_size // n_proc  # adjust batch size
+
+        tik = time.time()
+        for i in range(n_proc):
+            load_core, comp_core = self.core_binder(n_proc, n_samp, n_train, i)
+            p = dmp.Process(
+                target=train,
+                args=(*args, i, n_proc, comp_core, load_core, cnt, b_size, ep),
+            )
+            p.start()
+            processes.append(p)
+        for p in processes:
+            p.join()
+        t = time.time() - tik
+
+        self.counter[0] = self.counter[0] + 1
+
+        return t
+
+    def run(self, train, args):
+        """
+        The "run" function launches ARGO to traing GNN model
+        Step 1: run the auto-tuner to search for the optimal configuration
+        Step 2: record the optimal configuration
+        Step 3: use the optimal configuration repeatedly until the end of the model training
+
+        Parameters
+        ----------
+        train: Callable
+            The GNN training function.
+
+        args:
+            The inputs of the GNN training function.
+        """
+        result = self.auto_tuning(train, args)  # Step 1
+        x = result.x  # Step 2
+        self.mp_engine(
+            x, train, args, ep=(self.epoch - self.n_search)
+        )  # Step 3

examples/pytorch/argo/main.py

@@ -0,0 +1,272 @@
+import argparse
+import os
+
+import dgl
+import dgl.nn as dglnn
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+import torch.nn.functional as F
+from argo import ARGO
+from dgl.data import (
+    AsNodePredDataset,
+    FlickrDataset,
+    RedditDataset,
+    YelpDataset,
+)
+from dgl.dataloading import DataLoader, NeighborSampler, ShaDowKHopSampler
+from ogb.nodeproppred import DglNodePropPredDataset
+from torch.nn.parallel import DistributedDataParallel
+
+
+class GNN(nn.Module):
+    def __init__(
+        self, in_size, hid_size, out_size, num_layers=3, model_name="sage"
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList()
+
+        # GraphSAGE-mean
+        if model_name.lower() == "sage":
+            self.layers.append(dglnn.SAGEConv(in_size, hid_size, "mean"))
+            for i in range(num_layers - 2):
+                self.layers.append(dglnn.SAGEConv(hid_size, hid_size, "mean"))
+            self.layers.append(dglnn.SAGEConv(hid_size, out_size, "mean"))
+        # GCN
+        elif model_name.lower() == "gcn":
+            kwargs = {
+                "norm": "both",
+                "weight": True,
+                "bias": True,
+                "allow_zero_in_degree": True,
+            }
+            self.layers.append(dglnn.GraphConv(in_size, hid_size, **kwargs))
+            for i in range(num_layers - 2):
+                self.layers.append(
+                    dglnn.GraphConv(hid_size, hid_size, **kwargs)
+                )
+            self.layers.append(dglnn.GraphConv(hid_size, out_size, **kwargs))
+        else:
+            raise NotImplementedError
+
+        self.dropout = nn.Dropout(0.5)
+        self.hid_size = hid_size
+        self.out_size = out_size
+
+    def forward(self, blocks, x):
+        h = x
+        if hasattr(blocks, "__len__"):
+            for l, (layer, block) in enumerate(zip(self.layers, blocks)):
+                h = layer(block, h)
+                if l != len(self.layers) - 1:
+                    h = F.relu(h)
+                    h = self.dropout(h)
+        else:
+            for l, layer in enumerate(self.layers):
+                h = layer(blocks, h)
+                if l != len(self.layers) - 1:
+                    h = F.relu(h)
+                    h = self.dropout(h)
+        return h
+
+
+def _train(**kwargs):
+    total_loss = 0
+    loader = kwargs["loader"]
+    model = kwargs["model"]
+    opt = kwargs["opt"]
+    load_core = kwargs["load_core"]
+    comp_core = kwargs["comp_core"]
+
+    device = torch.device("cpu")
+    with loader.enable_cpu_affinity(
+        loader_cores=load_core, compute_cores=comp_core
+    ):
+        for it, (input_nodes, output_nodes, blocks) in enumerate(loader):
+            if hasattr(blocks, "__len__"):
+                x = blocks[0].srcdata["feat"].to(torch.float32)
+                y = blocks[-1].dstdata["label"]
+            else:
+                x = blocks.srcdata["feat"].to(torch.float32)
+                y = blocks.dstdata["label"]
+            if kwargs["device"] == "cpu":  # for papers100M
+                y = y.type(torch.LongTensor)
+                y_hat = model(blocks, x)
+            else:
+                y = y.type(torch.LongTensor).to(device)
+                y_hat = model(blocks, x).to(device)
+            try:
+                loss = F.cross_entropy(
+                    y_hat[: output_nodes.shape[0]], y[: output_nodes.shape[0]]
+                )
+            except:
+                loss = F.binary_cross_entropy_with_logits(
+                    y_hat[: output_nodes.shape[0]].float(),
+                    y[: output_nodes.shape[0]].float(),
+                    reduction="sum",
+                )
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            del input_nodes, output_nodes, blocks
+            total_loss += loss.item()
+    return total_loss
+
+
+def train(
+    args, g, data, rank, world_size, comp_core, load_core, counter, b_size, ep
+):
+
+    num_classes, train_idx = data
+    dist.init_process_group("gloo", rank=rank, world_size=world_size)
+    device = torch.device("cpu")
+    hidden = args.hidden
+    # create GraphSAGE model
+    in_size = g.ndata["feat"].shape[1]
+    model = GNN(
+        in_size,
+        hidden,
+        num_classes,
+        num_layers=args.layer,
+        model_name=args.model,
+    ).to(device)
+    model = DistributedDataParallel(model)
+    num_of_samplers = len(load_core)
+    # create loader
+    drop_last, shuffle = True, True
+    if args.sampler.lower() == "neighbor":
+        sampler = NeighborSampler(
+            [int(fanout) for fanout in args.fan_out.split(",")],
+            prefetch_node_feats=["feat"],
+            prefetch_labels=["label"],
+        )
+        assert len(sampler.fanouts) == args.layer
+    elif args.sampler.lower() == "shadow":
+        sampler = ShaDowKHopSampler(
+            [10, 5],
+            output_device=device,
+            prefetch_node_feats=["feat"],
+        )
+    else:
+        raise NotImplementedError
+
+    train_dataloader = DataLoader(
+        g,
+        train_idx.to(device),
+        sampler,
+        device=device,
+        batch_size=b_size,
+        drop_last=drop_last,
+        shuffle=shuffle,
+        num_workers=num_of_samplers,
+        use_ddp=True,
+    )
+
+    # training loop
+    opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)
+    params = {
+        # training
+        "loader": train_dataloader,
+        "model": model,
+        "opt": opt,
+        # logging
+        "rank": rank,
+        "train_size": len(train_idx),
+        "batch_size": b_size,
+        "device": device,
+        "process": world_size,
+    }
+
+    PATH = "model.pt"
+    if counter[0] != 0:
+        checkpoint = torch.load(PATH, weights_only=False)
+        model.load_state_dict(checkpoint["model_state_dict"])
+        opt.load_state_dict(checkpoint["optimizer_state_dict"])
+        epoch = checkpoint["epoch"]
+        loss = checkpoint["loss"]
+
+    for epoch in range(ep):
+        params["epoch"] = epoch
+        model.train()
+        params["load_core"] = load_core
+        params["comp_core"] = comp_core
+        loss = _train(**params)
+        if rank == 0:
+            print("loss:", loss)
+
+    dist.barrier()
+    EPOCH = counter[0]
+    LOSS = loss
+    if rank == 0:
+        torch.save(
+            {
+                "epoch": EPOCH,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": opt.state_dict(),
+                "loss": LOSS,
+            },
+            PATH,
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="ogbn-products",
+        choices=[
+            "ogbn-papers100M",
+            "ogbn-products",
+            "reddit",
+            "yelp",
+            "flickr",
+        ],
+    )
+    parser.add_argument("--batch_size", type=int, default=1024 * 4)
+    parser.add_argument("--layer", type=int, default=3)
+    parser.add_argument("--fan_out", type=str, default="15,10,5")
+    parser.add_argument(
+        "--sampler",
+        type=str,
+        default="neighbor",
+        choices=["neighbor", "shadow"],
+    )
+    parser.add_argument(
+        "--model", type=str, default="sage", choices=["sage", "gcn"]
+    )
+    parser.add_argument("--hidden", type=int, default=128)
+    arguments = parser.parse_args()
+
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
+
+    if arguments.dataset in ["reddit", "flickr", "yelp"]:
+        if arguments.dataset == "reddit":
+            dataset = RedditDataset()
+        elif arguments.dataset == "flickr":
+            dataset = FlickrDataset()
+        else:
+            dataset = YelpDataset()
+        g = dataset[0]
+        train_mask = g.ndata["train_mask"]
+        idx = []
+        for i in range(len(train_mask)):
+            if train_mask[i]:
+                idx.append(i)
+        dataset.train_idx = torch.tensor(idx)
+    else:
+        dataset = AsNodePredDataset(DglNodePropPredDataset(arguments.dataset))
+        g = dataset[0]
+
+    data = (dataset.num_classes, dataset.train_idx)
+
+    in_size = g.ndata["feat"].shape[1]
+    out_size = dataset.num_classes
+    hidden_size = int(arguments.hidden)
+
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = "29501"
+    mp.set_start_method("fork", force=True)
+    runtime = ARGO(n_search=10, epoch=20, batch_size=arguments.batch_size)
+    runtime.run(train, args=(arguments, g, data))

examples/pytorch/argo/ogb_example.py

@@ -0,0 +1,302 @@
+"""
+This is modified version of: https://github.com/dmlc/dgl/blob/master/examples/pytorch/ogb/ogbn-products/graphsage/main.py
+"""
+
+import argparse
+import time
+
+import dgl
+import dgl.nn.pytorch as dglnn
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import tqdm
+from ogb.nodeproppred import DglNodePropPredDataset
+
+
+class SAGE(nn.Module):
+    def __init__(
+        self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
+    ):
+        super().__init__()
+        self.n_layers = n_layers
+        self.n_hidden = n_hidden
+        self.n_classes = n_classes
+        self.layers = nn.ModuleList()
+        self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, "mean"))
+        for i in range(1, n_layers - 1):
+            self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, "mean"))
+        self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
+        self.dropout = nn.Dropout(dropout)
+        self.activation = activation
+
+    def forward(self, blocks, x):
+        h = x
+        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
+            # We need to first copy the representation of nodes on the RHS from the
+            # appropriate nodes on the LHS.
+            # Note that the shape of h is (num_nodes_LHS, D) and the shape of h_dst
+            # would be (num_nodes_RHS, D)
+            h_dst = h[: block.num_dst_nodes()]
+            # Then we compute the updated representation on the RHS.
+            # The shape of h now becomes (num_nodes_RHS, D)
+            h = layer(block, (h, h_dst))
+            if l != len(self.layers) - 1:
+                h = self.activation(h)
+                h = self.dropout(h)
+        return h
+
+    def inference(self, g, x, device):
+        """
+        Inference with the GraphSAGE model on full neighbors (i.e. without neighbor sampling).
+        g : the entire graph.
+        x : the input of entire node set.
+        The inference code is written in a fashion that it could handle any number of nodes and
+        layers.
+        """
+        # During inference with sampling, multi-layer blocks are very inefficient because
+        # lots of computations in the first few layers are repeated.
+        # Therefore, we compute the representation of all nodes layer by layer.  The nodes
+        # on each layer are of course splitted in batches.
+        # TODO: can we standardize this?
+        for l, layer in enumerate(self.layers):
+            y = th.zeros(
+                g.num_nodes(),
+                self.n_hidden if l != len(self.layers) - 1 else self.n_classes,
+            ).to(device)
+
+            sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
+            dataloader = dgl.dataloading.DataLoader(
+                g,
+                th.arange(g.num_nodes()),
+                sampler,
+                batch_size=args.batch_size,
+                shuffle=True,
+                drop_last=False,
+                num_workers=args.num_workers,
+            )
+
+            for input_nodes, output_nodes, blocks in tqdm.tqdm(
+                dataloader, disable=None
+            ):
+                block = blocks[0].int().to(device)
+
+                h = x[input_nodes]
+                h_dst = h[: block.num_dst_nodes()]
+                h = layer(block, (h, h_dst))
+                if l != len(self.layers) - 1:
+                    h = self.activation(h)
+                    h = self.dropout(h)
+
+                y[output_nodes] = h
+
+            x = y
+        return y
+
+
+def compute_acc(pred, labels):
+    """
+    Compute the accuracy of prediction given the labels.
+    """
+    return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred)
+
+
+def evaluate(model, g, nfeat, labels, val_nid, test_nid, device):
+    """
+    Evaluate the model on the validation set specified by ``val_mask``.
+    g : The entire graph.
+    inputs : The features of all the nodes.
+    labels : The labels of all the nodes.
+    val_mask : A 0-1 mask indicating which nodes do we actually compute the accuracy for.
+    device : The GPU device to evaluate on.
+    """
+    model.eval()
+    with th.no_grad():
+        pred = model.inference(g, nfeat, device)
+    model.train()
+    return (
+        compute_acc(pred[val_nid], labels[val_nid]),
+        compute_acc(pred[test_nid], labels[test_nid]),
+        pred,
+    )
+
+
+def load_subtensor(nfeat, labels, seeds, input_nodes):
+    """
+    Extracts features and labels for a set of nodes.
+    """
+    batch_inputs = nfeat[input_nodes]
+    batch_labels = labels[seeds]
+    return batch_inputs, batch_labels
+
+
+#### Entry point
+def train(args, device, data):
+    # Unpack data
+    train_nid, val_nid, test_nid, in_feats, labels, n_classes, nfeat, g = data
+
+    # Create PyTorch DataLoader for constructing blocks
+    sampler = dgl.dataloading.MultiLayerNeighborSampler(
+        [int(fanout) for fanout in args.fan_out.split(",")]
+    )
+    dataloader = dgl.dataloading.DataLoader(
+        g,
+        train_nid,
+        sampler,
+        batch_size=args.batch_size,
+        shuffle=True,
+        drop_last=False,
+        num_workers=args.num_workers,
+    )
+
+    # Define model and optimizer
+    model = SAGE(
+        in_feats,
+        args.num_hidden,
+        n_classes,
+        args.num_layers,
+        F.relu,
+        args.dropout,
+    )
+    model = model.to(device)
+    loss_fcn = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
+
+    # Training loop
+    avg = 0
+    iter_tput = []
+    best_eval_acc = 0
+    best_test_acc = 0
+    with dataloader.enable_cpu_affinity():
+        for epoch in range(args.num_epochs):
+            tic = time.time()
+
+            # Loop over the dataloader to sample the computation dependency graph as a list of
+            # blocks.
+            for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
+                tic_step = time.time()
+
+                # copy block to gpu
+                blocks = [blk.int().to(device) for blk in blocks]
+
+                # Load the input features as well as output labels
+                batch_inputs, batch_labels = load_subtensor(
+                    nfeat, labels, seeds, input_nodes
+                )
+
+                # Compute loss and prediction
+                batch_pred = model(blocks, batch_inputs)
+                loss = loss_fcn(batch_pred, batch_labels)
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                iter_tput.append(len(seeds) / (time.time() - tic_step))
+                if step % args.log_every == 0 and step != 0:
+                    acc = compute_acc(batch_pred, batch_labels)
+                    print(
+                        "Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f}".format(
+                            step,
+                            loss.item(),
+                            acc.item(),
+                            np.mean(iter_tput[3:]),
+                        )
+                    )
+
+            toc = time.time()
+            print("Epoch Time(s): {:.4f}".format(toc - tic))
+            avg += toc - tic
+            if epoch % args.eval_every == 0 and epoch != 0:
+                eval_acc, test_acc, pred = evaluate(
+                    model, g, nfeat, labels, val_nid, test_nid, device
+                )
+                if args.save_pred:
+                    np.savetxt(
+                        args.save_pred + "%02d" % epoch,
+                        pred.argmax(1).cpu().numpy(),
+                        "%d",
+                    )
+                print("Eval Acc {:.4f}".format(eval_acc))
+                if eval_acc > best_eval_acc:
+                    best_eval_acc = eval_acc
+                    best_test_acc = test_acc
+                print(
+                    "Best Eval Acc {:.4f} Test Acc {:.4f}".format(
+                        best_eval_acc, best_test_acc
+                    )
+                )
+
+    print("Avg epoch time: {}".format(avg / args.num_epochs))
+    return best_test_acc
+
+
+if __name__ == "__main__":
+    argparser = argparse.ArgumentParser("multi-gpu training")
+    argparser.add_argument(
+        "--gpu",
+        type=int,
+        default=0,
+        help="GPU device ID. Use -1 for CPU training",
+    )
+    argparser.add_argument("--num-epochs", type=int, default=20)
+    argparser.add_argument("--num-hidden", type=int, default=256)
+    argparser.add_argument("--num-layers", type=int, default=3)
+    argparser.add_argument("--fan-out", type=str, default="5,10,15")
+    argparser.add_argument("--batch-size", type=int, default=1000)
+    argparser.add_argument("--val-batch-size", type=int, default=10000)
+    argparser.add_argument("--log-every", type=int, default=20)
+    argparser.add_argument("--eval-every", type=int, default=1)
+    argparser.add_argument("--lr", type=float, default=0.003)
+    argparser.add_argument("--dropout", type=float, default=0.5)
+    argparser.add_argument(
+        "--dataset",
+        type=str,
+        default="ogbn-products",
+        choices=["ogbn-papers100M", "ogbn-products"],
+    )
+    argparser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of sampling processes. Use 0 for no extra process.",
+    )
+    argparser.add_argument("--save-pred", type=str, default="")
+    argparser.add_argument("--wd", type=float, default=0)
+    args = argparser.parse_args()
+
+    device = th.device("cpu")
+
+    # load ogbn-products data
+    data = DglNodePropPredDataset(args.dataset)
+    splitted_idx = data.get_idx_split()
+    train_idx, val_idx, test_idx = (
+        splitted_idx["train"],
+        splitted_idx["valid"],
+        splitted_idx["test"],
+    )
+    graph, labels = data[0]
+    nfeat = graph.ndata.pop("feat").to(device)
+    labels = labels[:, 0].to(device)
+
+    in_feats = nfeat.shape[1]
+    n_classes = (labels.max() + 1).item()
+    # Create csr/coo/csc formats before launching sampling processes
+    # This avoids creating certain formats in each data loader process, which saves momory and CPU.
+    graph.create_formats_()
+    # Pack data
+    data = (
+        train_idx,
+        val_idx,
+        test_idx,
+        in_feats,
+        labels,
+        n_classes,
+        nfeat,
+        graph,
+    )
+
+    test_acc = train(args, device, data).cpu().numpy()
+    print("Test accuracy:", test_acc)

examples/pytorch/argo/ogb_example_ARGO.py

@@ -0,0 +1,364 @@
+"""
+This is a modified version of: https://github.com/dmlc/dgl/blob/master/examples/pytorch/ogb/ogbn-products/graphsage/main.py
+This example shows how to enable ARGO to automatically instantiate multi-processing and adjust CPU core assignment to achieve better performance.
+"""
+
+import argparse
+
+import ctypes
+import os
+import time
+from multiprocessing import RawValue
+
+import dgl
+import dgl.nn.pytorch as dglnn
+
+import numpy as np
+import torch as th
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import tqdm
+from argo import ARGO
+from ogb.nodeproppred import DglNodePropPredDataset
+from torch.nn.parallel import DistributedDataParallel
+
+avg_total = RawValue(ctypes.c_float, 0.0)
+
+
+class SAGE(nn.Module):
+    def __init__(
+        self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
+    ):
+        super().__init__()
+        self.n_layers = n_layers
+        self.n_hidden = n_hidden
+        self.n_classes = n_classes
+        self.layers = nn.ModuleList()
+        self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, "mean"))
+        for i in range(1, n_layers - 1):
+            self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, "mean"))
+        self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
+        self.dropout = nn.Dropout(dropout)
+        self.activation = activation
+
+    def forward(self, blocks, x):
+        h = x
+        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
+            # We need to first copy the representation of nodes on the RHS from the
+            # appropriate nodes on the LHS.
+            # Note that the shape of h is (num_nodes_LHS, D) and the shape of h_dst
+            # would be (num_nodes_RHS, D)
+            h_dst = h[: block.num_dst_nodes()]
+            # Then we compute the updated representation on the RHS.
+            # The shape of h now becomes (num_nodes_RHS, D)
+            h = layer(block, (h, h_dst))
+            if l != len(self.layers) - 1:
+                h = self.activation(h)
+                h = self.dropout(h)
+        return h
+
+    def inference(self, g, x, device):
+        """
+        Inference with the GraphSAGE model on full neighbors (i.e. without neighbor sampling).
+        g : the entire graph.
+        x : the input of entire node set.
+        The inference code is written in a fashion that it could handle any number of nodes and
+        layers.
+        """
+        # During inference with sampling, multi-layer blocks are very inefficient because
+        # lots of computations in the first few layers are repeated.
+        # Therefore, we compute the representation of all nodes layer by layer.  The nodes
+        # on each layer are of course splitted in batches.
+        # TODO: can we standardize this?
+        for l, layer in enumerate(self.layers):
+            y = th.zeros(
+                g.num_nodes(),
+                self.n_hidden if l != len(self.layers) - 1 else self.n_classes,
+            ).to(device)
+
+            sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
+            dataloader = dgl.dataloading.DataLoader(
+                g,
+                th.arange(g.num_nodes()),
+                sampler,
+                batch_size=args.batch_size,
+                shuffle=True,
+                drop_last=False,
+                num_workers=args.num_workers,
+            )
+
+            for input_nodes, output_nodes, blocks in tqdm.tqdm(
+                dataloader, disable=None
+            ):
+                block = blocks[0].int().to(device)
+
+                h = x[input_nodes]
+                h_dst = h[: block.num_dst_nodes()]
+                h = layer(block, (h, h_dst))
+                if l != len(self.layers) - 1:
+                    h = self.activation(h)
+                    h = self.dropout(h)
+
+                y[output_nodes] = h
+
+            x = y
+        return y
+
+
+def compute_acc(pred, labels):
+    """
+    Compute the accuracy of prediction given the labels.
+    """
+    return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred)
+
+
+def evaluate(model, g, nfeat, labels, val_nid, test_nid, device):
+    """
+    Evaluate the model on the validation set specified by ``val_mask``.
+    g : The entire graph.
+    inputs : The features of all the nodes.
+    labels : The labels of all the nodes.
+    val_mask : A 0-1 mask indicating which nodes do we actually compute the accuracy for.
+    device : The GPU device to evaluate on.
+    """
+    model.eval()
+    with th.no_grad():
+        pred = model.module.inference(g, nfeat, device)
+    model.train()
+    return (
+        compute_acc(pred[val_nid], labels[val_nid]),
+        compute_acc(pred[test_nid], labels[test_nid]),
+        pred,
+    )
+
+
+def load_subtensor(nfeat, labels, seeds, input_nodes):
+    """
+    Extracts features and labels for a set of nodes.
+    """
+    batch_inputs = nfeat[input_nodes]
+    batch_labels = labels[seeds]
+    return batch_inputs, batch_labels
+
+
+#### Entry point
+def train(
+    args,
+    device,
+    data,
+    rank,
+    world_size,
+    comp_core,
+    load_core,
+    counter,
+    b_size,
+    ep,
+):
+    dist.init_process_group("gloo", rank=rank, world_size=world_size)
+    # Unpack data
+    train_nid, val_nid, test_nid, in_feats, labels, n_classes, nfeat, g = data
+
+    # Create PyTorch DataLoader for constructing blocks
+    sampler = dgl.dataloading.MultiLayerNeighborSampler(
+        [int(fanout) for fanout in args.fan_out.split(",")]
+    )
+    dataloader = dgl.dataloading.DataLoader(
+        g,
+        train_nid,
+        sampler,
+        batch_size=b_size,
+        shuffle=True,
+        drop_last=False,
+        num_workers=len(load_core),
+        use_ddp=True,
+    )
+
+    # Define model and optimizer
+    model = SAGE(
+        in_feats,
+        args.num_hidden,
+        n_classes,
+        args.num_layers,
+        F.relu,
+        args.dropout,
+    )
+    model = model.to(device)
+    model = DistributedDataParallel(model)
+    loss_fcn = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
+
+    # Training loop
+    avg = 0
+    iter_tput = []
+    best_eval_acc = 0
+    best_test_acc = 0
+    PATH = "model.pt"
+    if counter[0] != 0:
+        checkpoint = th.load(PATH)
+        model.load_state_dict(checkpoint["model_state_dict"])
+        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+        epoch = checkpoint["epoch"]
+        loss = checkpoint["loss"]
+
+    with dataloader.enable_cpu_affinity(
+        loader_cores=load_core, compute_cores=comp_core
+    ):
+        for epoch in range(ep):
+            tic = time.time()
+
+            # Loop over the dataloader to sample the computation dependency graph as a list of
+            # blocks.
+            for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
+                tic_step = time.time()
+
+                # copy block to gpu
+                blocks = [blk.int().to(device) for blk in blocks]
+
+                # Load the input features as well as output labels
+                batch_inputs, batch_labels = load_subtensor(
+                    nfeat, labels, seeds, input_nodes
+                )
+
+                # Compute loss and prediction
+                batch_pred = model(blocks, batch_inputs)
+                loss = loss_fcn(batch_pred, batch_labels)
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                iter_tput.append(len(seeds) / (time.time() - tic_step))
+                if step % args.log_every == 0 and step != 0:
+                    acc = compute_acc(batch_pred, batch_labels)
+                    print(
+                        "Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f}".format(
+                            step,
+                            loss.item(),
+                            acc.item(),
+                            np.mean(iter_tput[3:]),
+                        )
+                    )
+
+            toc = time.time()
+            print("Epoch Time(s): {:.4f}".format(toc - tic))
+
+            if rank == 0:
+                global avg_total
+                avg_total.value += toc - tic
+                avg += toc - tic
+
+                if epoch % args.eval_every == 0 and epoch != 0:
+                    eval_acc, test_acc, pred = evaluate(
+                        model, g, nfeat, labels, val_nid, test_nid, device
+                    )
+                    if args.save_pred:
+                        np.savetxt(
+                            args.save_pred + "%02d" % epoch,
+                            pred.argmax(1).cpu().numpy(),
+                            "%d",
+                        )
+                    print("Eval Acc {:.4f}".format(eval_acc))
+                    if eval_acc > best_eval_acc:
+                        best_eval_acc = eval_acc
+                        best_test_acc = test_acc
+                    print(
+                        "Best Eval Acc {:.4f} Test Acc {:.4f}".format(
+                            best_eval_acc, best_test_acc
+                        )
+                    )
+
+    dist.barrier()
+    if rank == 0:
+        th.save(
+            {
+                "epoch": counter[0],
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "loss": loss,
+            },
+            PATH,
+        )
+        if args.num_epochs == counter[0] + epoch + 1:
+            print(
+                "Avg epoch time: {}".format(avg_total.value / args.num_epochs)
+            )
+            print(
+                "Avg epoch time after auto-tuning: {}".format(avg / (epoch + 1))
+            )
+
+    return best_test_acc
+
+
+if __name__ == "__main__":
+    argparser = argparse.ArgumentParser("multi-gpu training")
+    argparser.add_argument(
+        "--gpu",
+        type=int,
+        default=0,
+        help="GPU device ID. Use -1 for CPU training",
+    )
+    argparser.add_argument("--num-epochs", type=int, default=20)
+    argparser.add_argument("--num-hidden", type=int, default=256)
+    argparser.add_argument("--num-layers", type=int, default=3)
+    argparser.add_argument("--fan-out", type=str, default="5,10,15")
+    argparser.add_argument("--batch-size", type=int, default=1000)
+    argparser.add_argument("--val-batch-size", type=int, default=10000)
+    argparser.add_argument("--log-every", type=int, default=20)
+    argparser.add_argument("--eval-every", type=int, default=1)
+    argparser.add_argument("--lr", type=float, default=0.003)
+    argparser.add_argument("--dropout", type=float, default=0.5)
+    argparser.add_argument(
+        "--dataset",
+        type=str,
+        default="ogbn-products",
+        choices=["ogbn-papers100M", "ogbn-products"],
+    )
+    argparser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of sampling processes. Use 0 for no extra process.",
+    )
+    argparser.add_argument("--save-pred", type=str, default="")
+    argparser.add_argument("--wd", type=float, default=0)
+    args = argparser.parse_args()
+
+    device = th.device("cpu")
+
+    # load ogbn-products data
+    data = DglNodePropPredDataset(args.dataset)
+    splitted_idx = data.get_idx_split()
+    train_idx, val_idx, test_idx = (
+        splitted_idx["train"],
+        splitted_idx["valid"],
+        splitted_idx["test"],
+    )
+    graph, labels = data[0]
+    nfeat = graph.ndata.pop("feat").to(device)
+    labels = labels[:, 0].to(device)
+
+    in_feats = nfeat.shape[1]
+    n_classes = (labels.max() + 1).item()
+    # Create csr/coo/csc formats before launching sampling processes
+    # This avoids creating certain formats in each data loader process, which saves momory and CPU.
+    graph.create_formats_()
+    # Pack data
+    data = (
+        train_idx,
+        val_idx,
+        test_idx,
+        in_feats,
+        labels,
+        n_classes,
+        nfeat,
+        graph,
+    )
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = "29501"
+    mp.set_start_method("fork", force=True)
+    runtime = ARGO(
+        n_search=15, epoch=args.num_epochs, batch_size=args.batch_size
+    )  # initialization
+    runtime.run(train, args=(args, device, data))  # wrap the training function

examples/pytorch/arma/README.md

@@ -0,0 +1,78 @@
+# DGL Implementation of ARMA
+
+This DGL example implements the GNN model proposed in the paper [Graph Neural Networks with convolutional ARMA filters](https://arxiv.org/abs/1901.01343).
+
+Contributor: [xnuohz](https://github.com/xnuohz)
+
+### Requirements
+The codebase is implemented in Python 3.6. For version requirement of packages, see below.
+
+```
+dgl
+numpy 1.19.5
+networkx 2.5
+scikit-learn 0.24.1
+tqdm 4.56.0
+torch 1.7.0
+```
+
+### The graph datasets used in this example
+
+###### Node Classification
+
+The DGL's built-in Cora, Pubmed, Citeseer datasets. Dataset summary:
+
+| Dataset | #Nodes | #Edges | #Feats | #Classes | #Train Nodes | #Val Nodes | #Test Nodes |
+| :-: | :-: | :-: | :-: | :-: | :-: | :-: | :-: |
+| Cora | 2,708 | 10,556 | 1,433 | 7(single label) | 140 | 500 | 1000 |
+| Citeseer | 3,327 | 9,228 | 3,703 | 6(single label) | 120 | 500 | 1000 |
+| Pubmed | 19,717 | 88,651 | 500 | 3(single label) | 60 | 500 | 1000 |
+
+### Usage
+
+###### Dataset options
+```
+--dataset          str     The graph dataset name.             Default is 'Cora'.
+```
+
+###### GPU options
+```
+--gpu              int     GPU index.                          Default is -1, using CPU.
+```
+
+###### Model options
+```
+--epochs           int     Number of training epochs.          Default is 2000.
+--early-stopping   int     Early stopping rounds.              Default is 100.
+--lr               float   Adam optimizer learning rate.       Default is 0.01.
+--lamb             float   L2 regularization coefficient.      Default is 0.0005.
+--hid-dim          int     Hidden layer dimensionalities.      Default is 16.
+--num-stacks       int     Number of K.                        Default is 2.
+--num-layers       int     Number of T.                        Default is 1.
+--dropout          float   Dropout applied at all layers.      Default is 0.75.
+```
+
+###### Examples
+
+The following commands learn a neural network and predict on the test set.
+Train an ARMA model which follows the original hyperparameters on different datasets.
+```bash
+# Cora:
+python citation.py --gpu 0
+
+# Citeseer:
+python citation.py --gpu 0 --dataset Citeseer --num-stacks 3
+
+# Pubmed:
+python citation.py --gpu 0 --dataset Pubmed --dropout 0.25 --num-stacks 1
+```
+
+### Performance
+
+###### Node Classification
+
+| Dataset | Cora | Citeseer | Pubmed |
+| :-: | :-: | :-: | :-: |
+| Metrics(Table 1.Node classification accuracy) | 83.4±0.6 | 72.5±0.4 | 78.9±0.3 |
+| Metrics(PyG) | 82.3±0.5 | 70.9±1.1 | 78.3±0.8 |
+| Metrics(DGL) | 80.9±0.6 | 71.6±0.8 | 75.0±4.2 |

examples/pytorch/arma/citation.py

@@ -0,0 +1,188 @@
+""" The main file to train an ARMA model using a full graph """
+
+import argparse
+import copy
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
+from model import ARMA4NC
+from tqdm import trange
+
+
+def main(args):
+    # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
+    # Load from DGL dataset
+    if args.dataset == "Cora":
+        dataset = CoraGraphDataset()
+    elif args.dataset == "Citeseer":
+        dataset = CiteseerGraphDataset()
+    elif args.dataset == "Pubmed":
+        dataset = PubmedGraphDataset()
+    else:
+        raise ValueError("Dataset {} is invalid.".format(args.dataset))
+
+    graph = dataset[0]
+
+    # check cuda
+    device = (
+        f"cuda:{args.gpu}"
+        if args.gpu >= 0 and torch.cuda.is_available()
+        else "cpu"
+    )
+
+    # retrieve the number of classes
+    n_classes = dataset.num_classes
+
+    # retrieve labels of ground truth
+    labels = graph.ndata.pop("label").to(device).long()
+
+    # Extract node features
+    feats = graph.ndata.pop("feat").to(device)
+    n_features = feats.shape[-1]
+
+    # retrieve masks for train/validation/test
+    train_mask = graph.ndata.pop("train_mask")
+    val_mask = graph.ndata.pop("val_mask")
+    test_mask = graph.ndata.pop("test_mask")
+
+    train_idx = torch.nonzero(train_mask, as_tuple=False).squeeze().to(device)
+    val_idx = torch.nonzero(val_mask, as_tuple=False).squeeze().to(device)
+    test_idx = torch.nonzero(test_mask, as_tuple=False).squeeze().to(device)
+
+    graph = graph.to(device)
+
+    # Step 2: Create model =================================================================== #
+    model = ARMA4NC(
+        in_dim=n_features,
+        hid_dim=args.hid_dim,
+        out_dim=n_classes,
+        num_stacks=args.num_stacks,
+        num_layers=args.num_layers,
+        activation=nn.ReLU(),
+        dropout=args.dropout,
+    ).to(device)
+
+    best_model = copy.deepcopy(model)
+
+    # Step 3: Create training components ===================================================== #
+    loss_fn = nn.CrossEntropyLoss()
+    opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.lamb)
+
+    # Step 4: training epoches =============================================================== #
+    acc = 0
+    no_improvement = 0
+    epochs = trange(args.epochs, desc="Accuracy & Loss")
+
+    for _ in epochs:
+        # Training using a full graph
+        model.train()
+
+        logits = model(graph, feats)
+
+        # compute loss
+        train_loss = loss_fn(logits[train_idx], labels[train_idx])
+        train_acc = torch.sum(
+            logits[train_idx].argmax(dim=1) == labels[train_idx]
+        ).item() / len(train_idx)
+
+        # backward
+        opt.zero_grad()
+        train_loss.backward()
+        opt.step()
+
+        # Validation using a full graph
+        model.eval()
+
+        with torch.no_grad():
+            valid_loss = loss_fn(logits[val_idx], labels[val_idx])
+            valid_acc = torch.sum(
+                logits[val_idx].argmax(dim=1) == labels[val_idx]
+            ).item() / len(val_idx)
+
+        # Print out performance
+        epochs.set_description(
+            "Train Acc {:.4f} | Train Loss {:.4f} | Val Acc {:.4f} | Val loss {:.4f}".format(
+                train_acc, train_loss.item(), valid_acc, valid_loss.item()
+            )
+        )
+
+        if valid_acc < acc:
+            no_improvement += 1
+            if no_improvement == args.early_stopping:
+                print("Early stop.")
+                break
+        else:
+            no_improvement = 0
+            acc = valid_acc
+            best_model = copy.deepcopy(model)
+
+    best_model.eval()
+    logits = best_model(graph, feats)
+    test_acc = torch.sum(
+        logits[test_idx].argmax(dim=1) == labels[test_idx]
+    ).item() / len(test_idx)
+
+    print("Test Acc {:.4f}".format(test_acc))
+    return test_acc
+
+
+if __name__ == "__main__":
+    """
+    ARMA Model Hyperparameters
+    """
+    parser = argparse.ArgumentParser(description="ARMA GCN")
+
+    # data source params
+    parser.add_argument(
+        "--dataset", type=str, default="Cora", help="Name of dataset."
+    )
+    # cuda params
+    parser.add_argument(
+        "--gpu", type=int, default=-1, help="GPU index. Default: -1, using CPU."
+    )
+    # training params
+    parser.add_argument(
+        "--epochs", type=int, default=2000, help="Training epochs."
+    )
+    parser.add_argument(
+        "--early-stopping",
+        type=int,
+        default=100,
+        help="Patient epochs to wait before early stopping.",
+    )
+    parser.add_argument("--lr", type=float, default=0.01, help="Learning rate.")
+    parser.add_argument("--lamb", type=float, default=5e-4, help="L2 reg.")
+    # model params
+    parser.add_argument(
+        "--hid-dim", type=int, default=16, help="Hidden layer dimensionalities."
+    )
+    parser.add_argument(
+        "--num-stacks", type=int, default=2, help="Number of K."
+    )
+    parser.add_argument(
+        "--num-layers", type=int, default=1, help="Number of T."
+    )
+    parser.add_argument(
+        "--dropout",
+        type=float,
+        default=0.75,
+        help="Dropout applied at all layers.",
+    )
+
+    args = parser.parse_args()
+    print(args)
+
+    acc_lists = []
+
+    for _ in range(100):
+        acc_lists.append(main(args))
+
+    mean = np.around(np.mean(acc_lists, axis=0), decimals=3)
+    std = np.around(np.std(acc_lists, axis=0), decimals=3)
+    print("Total acc: ", acc_lists)
+    print("mean", mean)
+    print("std", std)

examples/pytorch/arma/model.py

@@ -0,0 +1,151 @@
+import math
+
+import dgl.function as fn
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def glorot(tensor):
+    if tensor is not None:
+        stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1)))
+        tensor.data.uniform_(-stdv, stdv)
+
+
+def zeros(tensor):
+    if tensor is not None:
+        tensor.data.fill_(0)
+
+
+class ARMAConv(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        num_stacks,
+        num_layers,
+        activation=None,
+        dropout=0.0,
+        bias=True,
+    ):
+        super(ARMAConv, self).__init__()
+
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.K = num_stacks
+        self.T = num_layers
+        self.activation = activation
+        self.dropout = nn.Dropout(p=dropout)
+
+        # init weight
+        self.w_0 = nn.ModuleDict(
+            {
+                str(k): nn.Linear(in_dim, out_dim, bias=False)
+                for k in range(self.K)
+            }
+        )
+        # deeper weight
+        self.w = nn.ModuleDict(
+            {
+                str(k): nn.Linear(out_dim, out_dim, bias=False)
+                for k in range(self.K)
+            }
+        )
+        # v
+        self.v = nn.ModuleDict(
+            {
+                str(k): nn.Linear(in_dim, out_dim, bias=False)
+                for k in range(self.K)
+            }
+        )
+        # bias
+        if bias:
+            self.bias = nn.Parameter(
+                torch.Tensor(self.K, self.T, 1, self.out_dim)
+            )
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for k in range(self.K):
+            glorot(self.w_0[str(k)].weight)
+            glorot(self.w[str(k)].weight)
+            glorot(self.v[str(k)].weight)
+        zeros(self.bias)
+
+    def forward(self, g, feats):
+        with g.local_scope():
+            init_feats = feats
+            # assume that the graphs are undirected and graph.in_degrees() is the same as graph.out_degrees()
+            degs = g.in_degrees().float().clamp(min=1)
+            norm = torch.pow(degs, -0.5).to(feats.device).unsqueeze(1)
+            output = []
+
+            for k in range(self.K):
+                feats = init_feats
+                for t in range(self.T):
+                    feats = feats * norm
+                    g.ndata["h"] = feats
+                    g.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
+                    feats = g.ndata.pop("h")
+                    feats = feats * norm
+
+                    if t == 0:
+                        feats = self.w_0[str(k)](feats)
+                    else:
+                        feats = self.w[str(k)](feats)
+
+                    feats += self.dropout(self.v[str(k)](init_feats))
+                    feats += self.v[str(k)](self.dropout(init_feats))
+
+                    if self.bias is not None:
+                        feats += self.bias[k][t]
+
+                    if self.activation is not None:
+                        feats = self.activation(feats)
+                output.append(feats)
+
+            return torch.stack(output).mean(dim=0)
+
+
+class ARMA4NC(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        hid_dim,
+        out_dim,
+        num_stacks,
+        num_layers,
+        activation=None,
+        dropout=0.0,
+    ):
+        super(ARMA4NC, self).__init__()
+
+        self.conv1 = ARMAConv(
+            in_dim=in_dim,
+            out_dim=hid_dim,
+            num_stacks=num_stacks,
+            num_layers=num_layers,
+            activation=activation,
+            dropout=dropout,
+        )
+
+        self.conv2 = ARMAConv(
+            in_dim=hid_dim,
+            out_dim=out_dim,
+            num_stacks=num_stacks,
+            num_layers=num_layers,
+            activation=activation,
+            dropout=dropout,
+        )
+
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, g, feats):
+        feats = F.relu(self.conv1(g, feats))
+        feats = self.dropout(feats)
+        feats = self.conv2(g, feats)
+        return feats

examples/pytorch/bgnn/BGNN.py

@@ -0,0 +1,539 @@
+import itertools
+import time
+from collections import defaultdict as ddict
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from catboost import CatBoostClassifier, CatBoostRegressor, Pool, sum_models
+from sklearn import preprocessing
+from sklearn.metrics import r2_score
+from tqdm import tqdm
+
+
+class BGNNPredictor:
+    """
+    Description
+    -----------
+    Boost GNN predictor for semi-supervised node classification or regression problems.
+    Publication: https://arxiv.org/abs/2101.08543
+
+    Parameters
+    ----------
+    gnn_model : nn.Module
+        DGL implementation of GNN model.
+    task: str, optional
+        Regression or classification task.
+    loss_fn : callable, optional
+        Function that takes torch tensors, pred and true, and returns a scalar.
+    trees_per_epoch : int, optional
+        Number of GBDT trees to build each epoch.
+    backprop_per_epoch : int, optional
+        Number of backpropagation steps to make each epoch.
+    lr : float, optional
+        Learning rate of gradient descent optimizer.
+    append_gbdt_pred : bool, optional
+        Append GBDT predictions or replace original input node features.
+    train_input_features : bool, optional
+        Train original input node features.
+    gbdt_depth : int, optional
+        Depth of each tree in GBDT model.
+    gbdt_lr : float, optional
+        Learning rate of GBDT model.
+    gbdt_alpha : int, optional
+        Weight to combine previous and new GBDT trees.
+    random_seed : int, optional
+        random seed for GNN and GBDT models.
+
+    Examples
+    ----------
+    gnn_model = GAT(10, 20, num_heads=5),
+    bgnn = BGNNPredictor(gnn_model)
+    metrics = bgnn.fit(graph, X, y, train_mask, val_mask, test_mask, cat_features)
+    """
+
+    def __init__(
+        self,
+        gnn_model,
+        task="regression",
+        loss_fn=None,
+        trees_per_epoch=10,
+        backprop_per_epoch=10,
+        lr=0.01,
+        append_gbdt_pred=True,
+        train_input_features=False,
+        gbdt_depth=6,
+        gbdt_lr=0.1,
+        gbdt_alpha=1,
+        random_seed=0,
+    ):
+        self.device = torch.device(
+            "cuda:0" if torch.cuda.is_available() else "cpu"
+        )
+
+        self.model = gnn_model.to(self.device)
+        self.task = task
+        self.loss_fn = loss_fn
+        self.trees_per_epoch = trees_per_epoch
+        self.backprop_per_epoch = backprop_per_epoch
+        self.lr = lr
+        self.append_gbdt_pred = append_gbdt_pred
+        self.train_input_features = train_input_features
+        self.gbdt_depth = gbdt_depth
+        self.gbdt_lr = gbdt_lr
+        self.gbdt_alpha = gbdt_alpha
+        self.random_seed = random_seed
+        torch.manual_seed(random_seed)
+        np.random.seed(random_seed)
+
+    def init_gbdt_model(self, num_epochs, epoch):
+        if self.task == "regression":
+            catboost_model_obj = CatBoostRegressor
+            catboost_loss_fn = "RMSE"
+        else:
+            if epoch == 0:  # we predict multiclass probs at first epoch
+                catboost_model_obj = CatBoostClassifier
+                catboost_loss_fn = "MultiClass"
+            else:  # we predict the gradients for each class at epochs > 0
+                catboost_model_obj = CatBoostRegressor
+                catboost_loss_fn = "MultiRMSE"
+
+        return catboost_model_obj(
+            iterations=num_epochs,
+            depth=self.gbdt_depth,
+            learning_rate=self.gbdt_lr,
+            loss_function=catboost_loss_fn,
+            random_seed=self.random_seed,
+            nan_mode="Min",
+        )
+
+    def fit_gbdt(self, pool, trees_per_epoch, epoch):
+        gbdt_model = self.init_gbdt_model(trees_per_epoch, epoch)
+        gbdt_model.fit(pool, verbose=False)
+        return gbdt_model
+
+    def append_gbdt_model(self, new_gbdt_model, weights):
+        if self.gbdt_model is None:
+            return new_gbdt_model
+        return sum_models([self.gbdt_model, new_gbdt_model], weights=weights)
+
+    def train_gbdt(
+        self,
+        gbdt_X_train,
+        gbdt_y_train,
+        cat_features,
+        epoch,
+        gbdt_trees_per_epoch,
+        gbdt_alpha,
+    ):
+        pool = Pool(gbdt_X_train, gbdt_y_train, cat_features=cat_features)
+        epoch_gbdt_model = self.fit_gbdt(pool, gbdt_trees_per_epoch, epoch)
+        if epoch == 0 and self.task == "classification":
+            self.base_gbdt = epoch_gbdt_model
+        else:
+            self.gbdt_model = self.append_gbdt_model(
+                epoch_gbdt_model, weights=[1, gbdt_alpha]
+            )
+
+    def update_node_features(self, node_features, X, original_X):
+        # get predictions from gbdt model
+        if self.task == "regression":
+            predictions = np.expand_dims(
+                self.gbdt_model.predict(original_X), axis=1
+            )
+        else:
+            predictions = self.base_gbdt.predict_proba(original_X)
+            if self.gbdt_model is not None:
+                predictions_after_one = self.gbdt_model.predict(original_X)
+                predictions += predictions_after_one
+
+        # update node features with predictions
+        if self.append_gbdt_pred:
+            if self.train_input_features:
+                predictions = np.append(
+                    node_features.detach().cpu().data[:, : -self.out_dim],
+                    predictions,
+                    axis=1,
+                )  # replace old predictions with new predictions
+            else:
+                predictions = np.append(
+                    X, predictions, axis=1
+                )  # append original features with new predictions
+
+        predictions = torch.from_numpy(predictions).to(self.device)
+
+        node_features.data = predictions.float().data
+
+    def update_gbdt_targets(
+        self, node_features, node_features_before, train_mask
+    ):
+        return (
+            (node_features - node_features_before)
+            .detach()
+            .cpu()
+            .numpy()[train_mask, -self.out_dim :]
+        )
+
+    def init_node_features(self, X):
+        node_features = torch.empty(
+            X.shape[0], self.in_dim, requires_grad=True, device=self.device
+        )
+        if self.append_gbdt_pred:
+            node_features.data[:, : -self.out_dim] = torch.from_numpy(
+                X.to_numpy(copy=True)
+            )
+        return node_features
+
+    def init_optimizer(
+        self, node_features, optimize_node_features, learning_rate
+    ):
+        params = [self.model.parameters()]
+        if optimize_node_features:
+            params.append([node_features])
+        optimizer = torch.optim.Adam(itertools.chain(*params), lr=learning_rate)
+        return optimizer
+
+    def train_model(self, model_in, target_labels, train_mask, optimizer):
+        y = target_labels[train_mask]
+
+        self.model.train()
+        logits = self.model(*model_in).squeeze()
+        pred = logits[train_mask]
+
+        if self.loss_fn is not None:
+            loss = self.loss_fn(pred, y)
+        else:
+            if self.task == "regression":
+                loss = torch.sqrt(F.mse_loss(pred, y))
+            elif self.task == "classification":
+                loss = F.cross_entropy(pred, y.long())
+            else:
+                raise NotImplemented(
+                    "Unknown task. Supported tasks: classification, regression."
+                )
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        return loss
+
+    def evaluate_model(self, logits, target_labels, mask):
+        metrics = {}
+        y = target_labels[mask]
+        with torch.no_grad():
+            pred = logits[mask]
+            if self.task == "regression":
+                metrics["loss"] = torch.sqrt(
+                    F.mse_loss(pred, y).squeeze() + 1e-8
+                )
+                metrics["rmsle"] = torch.sqrt(
+                    F.mse_loss(torch.log(pred + 1), torch.log(y + 1)).squeeze()
+                    + 1e-8
+                )
+                metrics["mae"] = F.l1_loss(pred, y)
+                metrics["r2"] = torch.Tensor(
+                    [r2_score(y.cpu().numpy(), pred.cpu().numpy())]
+                )
+            elif self.task == "classification":
+                metrics["loss"] = F.cross_entropy(pred, y.long())
+                metrics["accuracy"] = torch.Tensor(
+                    [(y == pred.max(1)[1]).sum().item() / y.shape[0]]
+                )
+
+            return metrics
+
+    def train_and_evaluate(
+        self,
+        model_in,
+        target_labels,
+        train_mask,
+        val_mask,
+        test_mask,
+        optimizer,
+        metrics,
+        gnn_passes_per_epoch,
+    ):
+        loss = None
+
+        for _ in range(gnn_passes_per_epoch):
+            loss = self.train_model(
+                model_in, target_labels, train_mask, optimizer
+            )
+
+        self.model.eval()
+        logits = self.model(*model_in).squeeze()
+        train_results = self.evaluate_model(logits, target_labels, train_mask)
+        val_results = self.evaluate_model(logits, target_labels, val_mask)
+        test_results = self.evaluate_model(logits, target_labels, test_mask)
+        for metric_name in train_results:
+            metrics[metric_name].append(
+                (
+                    train_results[metric_name].detach().item(),
+                    val_results[metric_name].detach().item(),
+                    test_results[metric_name].detach().item(),
+                )
+            )
+        return loss
+
+    def update_early_stopping(
+        self,
+        metrics,
+        epoch,
+        best_metric,
+        best_val_epoch,
+        epochs_since_last_best_metric,
+        metric_name,
+        lower_better=False,
+    ):
+        train_metric, val_metric, test_metric = metrics[metric_name][-1]
+        if (lower_better and val_metric < best_metric[1]) or (
+            not lower_better and val_metric > best_metric[1]
+        ):
+            best_metric = metrics[metric_name][-1]
+            best_val_epoch = epoch
+            epochs_since_last_best_metric = 0
+        else:
+            epochs_since_last_best_metric += 1
+        return best_metric, best_val_epoch, epochs_since_last_best_metric
+
+    def log_epoch(
+        self,
+        pbar,
+        metrics,
+        epoch,
+        loss,
+        epoch_time,
+        logging_epochs,
+        metric_name="loss",
+    ):
+        train_metric, val_metric, test_metric = metrics[metric_name][-1]
+        if epoch and epoch % logging_epochs == 0:
+            pbar.set_description(
+                "Epoch {:05d} | Loss {:.3f} | Loss {:.3f}/{:.3f}/{:.3f} | Time {:.4f}".format(
+                    epoch,
+                    loss,
+                    train_metric,
+                    val_metric,
+                    test_metric,
+                    epoch_time,
+                )
+            )
+
+    def fit(
+        self,
+        graph,
+        X,
+        y,
+        train_mask,
+        val_mask,
+        test_mask,
+        original_X=None,
+        cat_features=None,
+        num_epochs=100,
+        patience=10,
+        logging_epochs=1,
+        metric_name="loss",
+    ):
+        """
+
+        :param graph : dgl.DGLGraph
+            Input graph
+        :param X : pd.DataFrame
+            Input node features. Each column represents one input feature. Each row is a node.
+            Values in dataframe are numerical, after preprocessing.
+        :param y : pd.DataFrame
+            Input node targets. Each column represents one target. Each row is a node
+            (order of nodes should be the same as in X).
+        :param train_mask : list[int]
+            Node indexes (rows) that belong to train set.
+        :param val_mask : list[int]
+            Node indexes (rows) that belong to validation set.
+        :param test_mask : list[int]
+            Node indexes (rows) that belong to test set.
+        :param original_X : pd.DataFrame, optional
+            Input node features before preprocessing. Each column represents one input feature. Each row is a node.
+            Values in dataframe can be of any type, including categorical (e.g. string, bool) or
+            missing values (None). This is useful if you want to preprocess X with GBDT model.
+        :param cat_features: list[int]
+            Feature indexes (columns) which are categorical features.
+        :param num_epochs : int
+            Number of epochs to run.
+        :param patience : int
+            Number of epochs to wait until early stopping.
+        :param logging_epochs : int
+            Log every n epoch.
+        :param metric_name : str
+            Metric to use for early stopping.
+        :param normalize_features : bool
+            If to normalize original input features X (column wise).
+        :param replace_na: bool
+            If to replace missing values (None) in X.
+        :return: metrics evaluated during training
+        """
+
+        # initialize for early stopping and metrics
+        if metric_name in ["r2", "accuracy"]:
+            best_metric = [np.cfloat("-inf")] * 3  # for train/val/test
+        else:
+            best_metric = [np.cfloat("inf")] * 3  # for train/val/test
+
+        best_val_epoch = 0
+        epochs_since_last_best_metric = 0
+        metrics = ddict(list)
+        if cat_features is None:
+            cat_features = []
+
+        if self.task == "regression":
+            self.out_dim = y.shape[1]
+        elif self.task == "classification":
+            self.out_dim = len(set(y.iloc[test_mask, 0]))
+        self.in_dim = (
+            self.out_dim + X.shape[1] if self.append_gbdt_pred else self.out_dim
+        )
+
+        if original_X is None:
+            original_X = X.copy()
+            cat_features = []
+
+        gbdt_X_train = original_X.iloc[train_mask]
+        gbdt_y_train = y.iloc[train_mask]
+        gbdt_alpha = self.gbdt_alpha
+        self.gbdt_model = None
+
+        node_features = self.init_node_features(X)
+        optimizer = self.init_optimizer(
+            node_features, optimize_node_features=True, learning_rate=self.lr
+        )
+
+        y = (
+            torch.from_numpy(y.to_numpy(copy=True))
+            .float()
+            .squeeze()
+            .to(self.device)
+        )
+        graph = graph.to(self.device)
+
+        pbar = tqdm(range(num_epochs))
+        for epoch in pbar:
+            start2epoch = time.time()
+
+            # gbdt part
+            self.train_gbdt(
+                gbdt_X_train,
+                gbdt_y_train,
+                cat_features,
+                epoch,
+                self.trees_per_epoch,
+                gbdt_alpha,
+            )
+
+            self.update_node_features(node_features, X, original_X)
+            node_features_before = node_features.clone()
+            model_in = (graph, node_features)
+            loss = self.train_and_evaluate(
+                model_in,
+                y,
+                train_mask,
+                val_mask,
+                test_mask,
+                optimizer,
+                metrics,
+                self.backprop_per_epoch,
+            )
+            gbdt_y_train = self.update_gbdt_targets(
+                node_features, node_features_before, train_mask
+            )
+
+            self.log_epoch(
+                pbar,
+                metrics,
+                epoch,
+                loss,
+                time.time() - start2epoch,
+                logging_epochs,
+                metric_name=metric_name,
+            )
+
+            # check early stopping
+            (
+                best_metric,
+                best_val_epoch,
+                epochs_since_last_best_metric,
+            ) = self.update_early_stopping(
+                metrics,
+                epoch,
+                best_metric,
+                best_val_epoch,
+                epochs_since_last_best_metric,
+                metric_name,
+                lower_better=(metric_name not in ["r2", "accuracy"]),
+            )
+            if patience and epochs_since_last_best_metric > patience:
+                break
+
+            if np.isclose(gbdt_y_train.sum(), 0.0):
+                print("Node embeddings do not change anymore. Stopping...")
+                break
+
+        print(
+            "Best {} at iteration {}: {:.3f}/{:.3f}/{:.3f}".format(
+                metric_name, best_val_epoch, *best_metric
+            )
+        )
+        return metrics
+
+    def predict(self, graph, X, test_mask):
+        graph = graph.to(self.device)
+        node_features = torch.empty(X.shape[0], self.in_dim).to(self.device)
+        self.update_node_features(node_features, X, X)
+        logits = self.model(graph, node_features).squeeze()
+        if self.task == "regression":
+            return logits[test_mask]
+        else:
+            return logits[test_mask].max(1)[1]
+
+    def plot_interactive(
+        self,
+        metrics,
+        legend,
+        title,
+        logx=False,
+        logy=False,
+        metric_name="loss",
+        start_from=0,
+    ):
+        import plotly.graph_objects as go
+
+        metric_results = metrics[metric_name]
+        xs = [list(range(len(metric_results)))] * len(metric_results[0])
+        ys = list(zip(*metric_results))
+
+        fig = go.Figure()
+        for i in range(len(ys)):
+            fig.add_trace(
+                go.Scatter(
+                    x=xs[i][start_from:],
+                    y=ys[i][start_from:],
+                    mode="lines+markers",
+                    name=legend[i],
+                )
+            )
+
+        fig.update_layout(
+            title=title,
+            title_x=0.5,
+            xaxis_title="Epoch",
+            yaxis_title=metric_name,
+            font=dict(
+                size=40,
+            ),
+            height=600,
+        )
+
+        if logx:
+            fig.update_layout(xaxis_type="log")
+        if logy:
+            fig.update_layout(yaxis_type="log")
+
+        fig.show()

examples/pytorch/bgnn/Readme.md

@@ -0,0 +1,16 @@
+# Instructions to download datasets:
+
+1. Download datasets from here: https://www.dropbox.com/s/verx1evkykzli88/datasets.zip
+2. Extract zip folder in this directory
+3. Choose the dataset you wish in `run.py` file. 
+
+# Details about BGNN model
+`run.py` implements a class for GNN model. You can select GAT, GCN, ChebNet, AGNN, or APPNP gnn models.
+Or you can provide your favorite GNN model. You can also pretrain your model or setup the hyperparameters you like. 
+
+Hyperparameters of BGNN model. 
+* `append_gbdt_pred` -- this decides whether to append GBDT predictions from GNN to original input features or to replace original input features with predictions of GBDT. This can be important for performance, so try both values, True and False. 
+* `trees_per_epoch` and `backprop_per_epoch`. Values in the range 5-15 usually gives good results. The more, the longer training is. 
+* `lr` is learning rate for GNN. 0.01-0.1 are good values to try.
+* `gbdt_lr` is learning rate for GBDT. Should be that important. 
+* `gbdt_depth` number of levels in GBDT tree. 4-8 are good values. The more, the longer it trains. 

examples/pytorch/bgnn/run.py

@@ -0,0 +1,254 @@
+import json
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from BGNN import BGNNPredictor
+from category_encoders import CatBoostEncoder
+
+from dgl.data.utils import load_graphs
+from dgl.nn.pytorch import (
+    AGNNConv as AGNNConvDGL,
+    APPNPConv,
+    ChebConv as ChebConvDGL,
+    GATConv as GATConvDGL,
+    GraphConv,
+)
+from sklearn import preprocessing
+from torch.nn import Dropout, ELU, Linear, ReLU, Sequential
+
+
+class GNNModelDGL(torch.nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        hidden_dim,
+        out_dim,
+        dropout=0.0,
+        name="gat",
+        residual=True,
+        use_mlp=False,
+        join_with_mlp=False,
+    ):
+        super(GNNModelDGL, self).__init__()
+        self.name = name
+        self.use_mlp = use_mlp
+        self.join_with_mlp = join_with_mlp
+        self.normalize_input_columns = True
+        if name == "gat":
+            self.l1 = GATConvDGL(
+                in_dim,
+                hidden_dim // 8,
+                8,
+                feat_drop=dropout,
+                attn_drop=dropout,
+                residual=False,
+                activation=F.elu,
+            )
+            self.l2 = GATConvDGL(
+                hidden_dim,
+                out_dim,
+                1,
+                feat_drop=dropout,
+                attn_drop=dropout,
+                residual=residual,
+                activation=None,
+            )
+        elif name == "gcn":
+            self.l1 = GraphConv(in_dim, hidden_dim, activation=F.elu)
+            self.l2 = GraphConv(hidden_dim, out_dim, activation=F.elu)
+            self.drop = Dropout(p=dropout)
+        elif name == "cheb":
+            self.l1 = ChebConvDGL(in_dim, hidden_dim, k=3)
+            self.l2 = ChebConvDGL(hidden_dim, out_dim, k=3)
+            self.drop = Dropout(p=dropout)
+        elif name == "agnn":
+            self.lin1 = Sequential(
+                Dropout(p=dropout), Linear(in_dim, hidden_dim), ELU()
+            )
+            self.l1 = AGNNConvDGL(learn_beta=False)
+            self.l2 = AGNNConvDGL(learn_beta=True)
+            self.lin2 = Sequential(
+                Dropout(p=dropout), Linear(hidden_dim, out_dim), ELU()
+            )
+        elif name == "appnp":
+            self.lin1 = Sequential(
+                Dropout(p=dropout),
+                Linear(in_dim, hidden_dim),
+                ReLU(),
+                Dropout(p=dropout),
+                Linear(hidden_dim, out_dim),
+            )
+            self.l1 = APPNPConv(k=10, alpha=0.1, edge_drop=0.0)
+
+    def forward(self, graph, features):
+        h = features
+        if self.use_mlp:
+            if self.join_with_mlp:
+                h = torch.cat((h, self.mlp(features)), 1)
+            else:
+                h = self.mlp(features)
+        if self.name == "gat":
+            h = self.l1(graph, h).flatten(1)
+            logits = self.l2(graph, h).mean(1)
+        elif self.name in ["appnp"]:
+            h = self.lin1(h)
+            logits = self.l1(graph, h)
+        elif self.name == "agnn":
+            h = self.lin1(h)
+            h = self.l1(graph, h)
+            h = self.l2(graph, h)
+            logits = self.lin2(h)
+        elif self.name == "che3b":
+            lambda_max = dgl.laplacian_lambda_max(graph)
+            h = self.drop(h)
+            h = self.l1(graph, h, lambda_max)
+            logits = self.l2(graph, h, lambda_max)
+        elif self.name == "gcn":
+            h = self.drop(h)
+            h = self.l1(graph, h)
+            logits = self.l2(graph, h)
+
+        return logits
+
+
+def read_input(input_folder):
+    X = pd.read_csv(f"{input_folder}/X.csv")
+    y = pd.read_csv(f"{input_folder}/y.csv")
+
+    categorical_columns = []
+    if os.path.exists(f"{input_folder}/cat_features.txt"):
+        with open(f"{input_folder}/cat_features.txt") as f:
+            for line in f:
+                if line.strip():
+                    categorical_columns.append(line.strip())
+
+    cat_features = None
+    if categorical_columns:
+        columns = X.columns
+        cat_features = np.where(columns.isin(categorical_columns))[0]
+
+        for col in list(columns[cat_features]):
+            X[col] = X[col].astype(str)
+
+    gs, _ = load_graphs(f"{input_folder}/graph.dgl")
+    graph = gs[0]
+
+    with open(f"{input_folder}/masks.json") as f:
+        masks = json.load(f)
+
+    return graph, X, y, cat_features, masks
+
+
+def normalize_features(X, train_mask, val_mask, test_mask):
+    min_max_scaler = preprocessing.MinMaxScaler()
+    A = X.to_numpy(copy=True)
+    A[train_mask] = min_max_scaler.fit_transform(A[train_mask])
+    A[val_mask + test_mask] = min_max_scaler.transform(A[val_mask + test_mask])
+    return pd.DataFrame(A, columns=X.columns).astype(float)
+
+
+def replace_na(X, train_mask):
+    if X.isna().any().any():
+        return X.fillna(X.iloc[train_mask].min() - 1)
+    return X
+
+
+def encode_cat_features(X, y, cat_features, train_mask, val_mask, test_mask):
+    enc = CatBoostEncoder()
+    A = X.to_numpy(copy=True)
+    b = y.to_numpy(copy=True)
+    A[np.ix_(train_mask, cat_features)] = enc.fit_transform(
+        A[np.ix_(train_mask, cat_features)], b[train_mask]
+    )
+    A[np.ix_(val_mask + test_mask, cat_features)] = enc.transform(
+        A[np.ix_(val_mask + test_mask, cat_features)]
+    )
+    A = A.astype(float)
+    return pd.DataFrame(A, columns=X.columns)
+
+
+if __name__ == "__main__":
+    # datasets can be found here: https://www.dropbox.com/s/verx1evkykzli88/datasets.zip
+    # Read dataset
+    input_folder = "datasets/avazu"
+    graph, X, y, cat_features, masks = read_input(input_folder)
+    train_mask, val_mask, test_mask = (
+        masks["0"]["train"],
+        masks["0"]["val"],
+        masks["0"]["test"],
+    )
+
+    encoded_X = X.copy()
+    normalizeFeatures = False
+    replaceNa = True
+
+    if len(cat_features):
+        encoded_X = encode_cat_features(
+            encoded_X, y, cat_features, train_mask, val_mask, test_mask
+        )
+    if normalizeFeatures:
+        encoded_X = normalize_features(
+            encoded_X, train_mask, val_mask, test_mask
+        )
+    if replaceNa:
+        encoded_X = replace_na(encoded_X, train_mask)
+
+    # specify parameters
+    task = "regression"
+    hidden_dim = 128
+    trees_per_epoch = 5  # 5-10 are good values to try
+    backprop_per_epoch = 5  # 5-10 are good values to try
+    lr = 0.1  # 0.01-0.1 are good values to try
+    append_gbdt_pred = (
+        False  # this can be important for performance (try True and False)
+    )
+    train_input_features = False
+    gbdt_depth = 6
+    gbdt_lr = 0.1
+
+    out_dim = (
+        y.shape[1] if task == "regression" else len(set(y.iloc[test_mask, 0]))
+    )
+    in_dim = out_dim + X.shape[1] if append_gbdt_pred else out_dim
+
+    # specify GNN model
+    gnn_model = GNNModelDGL(in_dim, hidden_dim, out_dim)
+
+    # initialize BGNN model
+    bgnn = BGNNPredictor(
+        gnn_model,
+        task=task,
+        loss_fn=None,
+        trees_per_epoch=trees_per_epoch,
+        backprop_per_epoch=backprop_per_epoch,
+        lr=lr,
+        append_gbdt_pred=append_gbdt_pred,
+        train_input_features=train_input_features,
+        gbdt_depth=gbdt_depth,
+        gbdt_lr=gbdt_lr,
+    )
+
+    # train
+    metrics = bgnn.fit(
+        graph,
+        encoded_X,
+        y,
+        train_mask,
+        val_mask,
+        test_mask,
+        original_X=X,
+        cat_features=cat_features,
+        num_epochs=100,
+        patience=10,
+        metric_name="loss",
+    )
+
+    bgnn.plot_interactive(
+        metrics,
+        legend=["train", "valid", "test"],
+        title="Avazu",
+        metric_name="loss",
+    )

examples/pytorch/bgrl/README.md

@@ -0,0 +1,111 @@
+# DGL Implementation of BGRL
+
+This DGL example implements the GNN experiment proposed in the paper [Large-Scale Representation Learning on Graphs via Bootstrapping](https://arxiv.org/abs/2102.06514). For the original implementation, see [here](https://github.com/nerdslab/bgrl).
+
+Contributor: [RecLusIve-F](https://github.com/RecLusIve-F)
+
+### Requirements
+
+The codebase is implemented in Python 3.8. For version requirement of packages, see below.
+
+```
+dgl 0.8.3
+numpy 1.21.2
+torch 1.10.2
+scikit-learn 1.0.2
+```
+
+### Dataset
+Dataset summary:
+
+|     Dataset      |     Task     | Nodes  |  Edges  | Features |     Classes     |
+|:----------------:|:------------:|:------:|:-------:|:--------:|:---------------:|
+|      WikiCS      | Transductive | 11,701 | 216,123 |   300    |       10        |
+| Amazon Computers | Transductive | 13,752 | 245,861 |   767    |       10        |
+|  Amazon Photos   | Transductive | 7,650  | 119,081 |   745    |        8        |
+|   Coauthor CS    | Transductive | 18,333 | 81,894  |  6,805   |       15        |
+| Coauthor Physics | Transductive | 34,493 | 247,962 |  8,415   |        5        |
+|  PPI(24 graphs)  |  Inductive   | 56,944 | 818,716 |    50    | 121(multilabel) |
+
+### Usage
+
+##### Dataset options
+```
+--dataset                     str         The graph dataset name.                         Default is 'amazon_photos'.
+```
+
+##### Model options
+```
+--graph_encoder_layer         list        Convolutional layer hidden sizes.               Default is [256, 128].
+--predictor_hidden_size       int         Hidden size of predictor.                       Default is 512.
+```
+
+##### Training options
+```
+--epochs                      int         The number of training epochs.                  Default is 10000.
+--lr                          float       The learning rate.                              Default is 0.00001.
+--weight_decay                float       The weight decay.                               Default is 0.00001.
+--mm                          float       The momentum for moving average.                Default is 0.99.
+--lr_warmup_epochs            int         Warmup period for learning rate scheduling.     Default is 1000.    
+--weights_dir                 str         Where to save the weights.                      Default is '../weights'.
+```
+
+##### Augmentation options
+```
+--drop_edge_p                 float      Probability of edge dropout.                     Default is [0., 0.].
+--feat_mask_p                 float      Probability of node feature masking.             Default is [0., 0.].
+```
+
+##### Evaluation options
+```
+--eval_epochs                 int        Evaluate every eval_epochs.                      Default is 250.
+--num_eval_splits             int        Number of evaluation splits.                     Default is 20.
+--data_seed                   int        Data split seed for evaluation.                  Default is 1.
+```
+
+### Instructions for experiments
+
+##### Transductive task
+```
+# Coauthor CS
+python main.py --dataset coauthor_cs --graph_encoder_layer 512 256 --drop_edge_p 0.3 0.2 --feat_mask_p 0.3 0.4
+
+# Coauthor Physics
+python main.py --dataset coauthor_physics --graph_encoder_layer 256 128 --drop_edge_p 0.4 0.1 --feat_mask_p 0.1 0.4
+
+# WikiCS
+python main.py --dataset wiki_cs --graph_encoder_layer 512 256 --drop_edge_p 0.2 0.3 --feat_mask_p 0.2 0.1 --lr 5e-4
+
+# Amazon Photos
+python main.py --dataset amazon_photos --graph_encoder_layer 256 128 --drop_edge_p 0.4 0.1 --feat_mask_p 0.1 0.2 --lr 1e-4
+
+# Amazon Computers
+python main.py --dataset amazon_computers --graph_encoder_layer 256 128 --drop_edge_p 0.5 0.4 --feat_mask_p 0.2 0.1 --lr 5e-4
+```
+
+##### Inductive task
+```
+# PPI
+python main.py --dataset ppi --graph_encoder_layer 512 512 --drop_edge_p 0.3 0.25 --feat_mask_p 0.25 0. --lr 5e-3
+```
+
+### Performance
+
+##### Transductive Task
+|        Dataset         |    WikiCS    |  Am. Comp.   |  Am. Photos  |    Co. CS    |   Co. Phy    |
+|:----------------------:|:------------:|:------------:|:------------:|:------------:|:------------:|
+|   Accuracy Reported    | 79.98 ± 0.10 | 90.34 ± 0.19 | 93.17 ± 0.30 | 93.31 ± 0.13 | 95.73 ± 0.05 |
+| Accuracy Official Code |    79.94     |    90.62     |    93.45     |    93.42     |    95.74     |
+|      Accuracy DGL      |    80.00     |    90.64     |    93.34     |    93.76     |    95.79     |
+
+##### Inductive Task
+|        Dataset         |     PPI      |
+|:----------------------:|:------------:|
+|   Micro-F1 Reported    | 69.41 ± 0.15 |
+| Accuracy Official Code |    68.83     |
+|      Micro-F1 DGL      |    68.65     |
+
+
+##### Accuracy reported is over 20 random dataset splits and model initializations. Micro-F1 reported is over 20 random model initializations.
+
+##### Accuracy official code and Accuracy DGL is only over 1 random dataset splits and model initialization. Micro-F1 official code and Micro-F1 DGL is only over 1 random model initialization.

examples/pytorch/bgrl/eval_function.py

@@ -0,0 +1,174 @@
+import numpy as np
+import torch
+from sklearn import metrics
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import GridSearchCV, ShuffleSplit, train_test_split
+from sklearn.multiclass import OneVsRestClassifier
+from sklearn.preprocessing import normalize, OneHotEncoder
+
+
+def fit_logistic_regression(X, y, data_random_seed=1, repeat=1):
+    # transform targets to one-hot vector
+    one_hot_encoder = OneHotEncoder(categories="auto", sparse=False)
+
+    y = one_hot_encoder.fit_transform(y.reshape(-1, 1)).astype(np.bool_)
+
+    # normalize x
+    X = normalize(X, norm="l2")
+
+    # set random state, this will ensure the dataset will be split exactly the same throughout training
+    rng = np.random.RandomState(data_random_seed)
+
+    accuracies = []
+    for _ in range(repeat):
+        # different random split after each repeat
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, y, test_size=0.8, random_state=rng
+        )
+
+        # grid search with one-vs-rest classifiers
+        logreg = LogisticRegression(solver="liblinear")
+        c = 2.0 ** np.arange(-10, 11)
+        cv = ShuffleSplit(n_splits=5, test_size=0.5)
+        clf = GridSearchCV(
+            estimator=OneVsRestClassifier(logreg),
+            param_grid=dict(estimator__C=c),
+            n_jobs=5,
+            cv=cv,
+            verbose=0,
+        )
+        clf.fit(X_train, y_train)
+
+        y_pred = clf.predict_proba(X_test)
+        y_pred = np.argmax(y_pred, axis=1)
+        y_pred = one_hot_encoder.transform(y_pred.reshape(-1, 1)).astype(
+            np.bool_
+        )
+
+        test_acc = metrics.accuracy_score(y_test, y_pred)
+        accuracies.append(test_acc)
+    return accuracies
+
+
+def fit_logistic_regression_preset_splits(
+    X, y, train_mask, val_mask, test_mask
+):
+    # transform targets to one-hot vector
+    one_hot_encoder = OneHotEncoder(categories="auto", sparse=False)
+    y = one_hot_encoder.fit_transform(y.reshape(-1, 1)).astype(np.bool_)
+
+    # normalize x
+    X = normalize(X, norm="l2")
+
+    accuracies = []
+    for split_id in range(train_mask.shape[1]):
+        # get train/val/test masks
+        tmp_train_mask, tmp_val_mask = (
+            train_mask[:, split_id],
+            val_mask[:, split_id],
+        )
+
+        # make custom cv
+        X_train, y_train = X[tmp_train_mask], y[tmp_train_mask]
+        X_val, y_val = X[tmp_val_mask], y[tmp_val_mask]
+        X_test, y_test = X[test_mask], y[test_mask]
+
+        # grid search with one-vs-rest classifiers
+        best_test_acc, best_acc = 0, 0
+        for c in 2.0 ** np.arange(-10, 11):
+            clf = OneVsRestClassifier(
+                LogisticRegression(solver="liblinear", C=c)
+            )
+            clf.fit(X_train, y_train)
+
+            y_pred = clf.predict_proba(X_val)
+            y_pred = np.argmax(y_pred, axis=1)
+            y_pred = one_hot_encoder.transform(y_pred.reshape(-1, 1)).astype(
+                np.bool_
+            )
+            val_acc = metrics.accuracy_score(y_val, y_pred)
+            if val_acc > best_acc:
+                best_acc = val_acc
+                y_pred = clf.predict_proba(X_test)
+                y_pred = np.argmax(y_pred, axis=1)
+                y_pred = one_hot_encoder.transform(
+                    y_pred.reshape(-1, 1)
+                ).astype(np.bool_)
+                best_test_acc = metrics.accuracy_score(y_test, y_pred)
+
+        accuracies.append(best_test_acc)
+    return accuracies
+
+
+def fit_ppi_linear(
+    num_classes, train_data, val_data, test_data, device, repeat=1
+):
+    r"""
+    Trains a linear layer on top of the representations. This function is specific to the PPI dataset,
+    which has multiple labels.
+    """
+
+    def train(classifier, train_data, optimizer):
+        classifier.train()
+
+        x, label = train_data
+        x, label = x.to(device), label.to(device)
+        for step in range(100):
+            # forward
+            optimizer.zero_grad()
+            pred_logits = classifier(x)
+
+            # loss and backprop
+            loss = criterion(pred_logits, label)
+            loss.backward()
+            optimizer.step()
+
+    def test(classifier, data):
+        classifier.eval()
+        x, label = data
+        label = label.cpu().numpy().squeeze()
+
+        # feed to network and classifier
+        with torch.no_grad():
+            pred_logits = classifier(x.to(device))
+            pred_class = (pred_logits > 0).float().cpu().numpy()
+
+        return (
+            metrics.f1_score(label, pred_class, average="micro")
+            if pred_class.sum() > 0
+            else 0
+        )
+
+    num_feats = train_data[0].size(1)
+    criterion = torch.nn.BCEWithLogitsLoss()
+
+    # normalization
+    mean, std = train_data[0].mean(0, keepdim=True), train_data[0].std(
+        0, unbiased=False, keepdim=True
+    )
+    train_data[0] = (train_data[0] - mean) / std
+    val_data[0] = (val_data[0] - mean) / std
+    test_data[0] = (test_data[0] - mean) / std
+
+    best_val_f1 = []
+    test_f1 = []
+    for _ in range(repeat):
+        tmp_best_val_f1 = 0
+        tmp_test_f1 = 0
+        for weight_decay in 2.0 ** np.arange(-10, 11, 2):
+            classifier = torch.nn.Linear(num_feats, num_classes).to(device)
+            optimizer = torch.optim.AdamW(
+                params=classifier.parameters(),
+                lr=0.01,
+                weight_decay=weight_decay,
+            )
+
+            train(classifier, train_data, optimizer)
+            val_f1 = test(classifier, val_data)
+            if val_f1 > tmp_best_val_f1:
+                tmp_best_val_f1 = val_f1
+                tmp_test_f1 = test(classifier, test_data)
+        best_val_f1.append(tmp_best_val_f1)
+        test_f1.append(tmp_test_f1)
+
+    return [best_val_f1], [test_f1]

examples/pytorch/bgrl/main.py

@@ -0,0 +1,264 @@
+import copy
+import os
+import warnings
+
+import dgl
+
+import numpy as np
+import torch
+from eval_function import (
+    fit_logistic_regression,
+    fit_logistic_regression_preset_splits,
+    fit_ppi_linear,
+)
+from model import (
+    BGRL,
+    compute_representations,
+    GCN,
+    GraphSAGE_GCN,
+    MLP_Predictor,
+)
+from torch.nn.functional import cosine_similarity
+from torch.optim import AdamW
+from tqdm import tqdm
+from utils import CosineDecayScheduler, get_dataset, get_graph_drop_transform
+
+warnings.filterwarnings("ignore")
+
+
+def train(
+    step,
+    model,
+    optimizer,
+    lr_scheduler,
+    mm_scheduler,
+    transform_1,
+    transform_2,
+    data,
+    args,
+):
+    model.train()
+
+    # update learning rate
+    lr = lr_scheduler.get(step)
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+    # update momentum
+    mm = 1 - mm_scheduler.get(step)
+
+    # forward
+    optimizer.zero_grad()
+
+    x1, x2 = transform_1(data), transform_2(data)
+
+    if args.dataset != "ppi":
+        x1, x2 = dgl.add_self_loop(x1), dgl.add_self_loop(x2)
+
+    q1, y2 = model(x1, x2)
+    q2, y1 = model(x2, x1)
+
+    loss = (
+        2
+        - cosine_similarity(q1, y2.detach(), dim=-1).mean()
+        - cosine_similarity(q2, y1.detach(), dim=-1).mean()
+    )
+    loss.backward()
+
+    # update online network
+    optimizer.step()
+    # update target network
+    model.update_target_network(mm)
+
+    return loss.item()
+
+
+def eval(model, dataset, device, args, train_data, val_data, test_data):
+    # make temporary copy of encoder
+    tmp_encoder = copy.deepcopy(model.online_encoder).eval()
+    val_scores = None
+
+    if args.dataset == "ppi":
+        train_data = compute_representations(tmp_encoder, train_data, device)
+        val_data = compute_representations(tmp_encoder, val_data, device)
+        test_data = compute_representations(tmp_encoder, test_data, device)
+        num_classes = train_data[1].shape[1]
+        val_scores, test_scores = fit_ppi_linear(
+            num_classes,
+            train_data,
+            val_data,
+            test_data,
+            device,
+            args.num_eval_splits,
+        )
+    elif args.dataset != "wiki_cs":
+        representations, labels = compute_representations(
+            tmp_encoder, dataset, device
+        )
+        test_scores = fit_logistic_regression(
+            representations.cpu().numpy(),
+            labels.cpu().numpy(),
+            data_random_seed=args.data_seed,
+            repeat=args.num_eval_splits,
+        )
+    else:
+        g = dataset[0]
+        train_mask = g.ndata["train_mask"]
+        val_mask = g.ndata["val_mask"]
+        test_mask = g.ndata["test_mask"]
+        representations, labels = compute_representations(
+            tmp_encoder, dataset, device
+        )
+        test_scores = fit_logistic_regression_preset_splits(
+            representations.cpu().numpy(),
+            labels.cpu().numpy(),
+            train_mask,
+            val_mask,
+            test_mask,
+        )
+
+    return val_scores, test_scores
+
+
+def main(args):
+    # use CUDA_VISIBLE_DEVICES to select gpu
+    device = (
+        torch.device("cuda")
+        if torch.cuda.is_available()
+        else torch.device("cpu")
+    )
+    print("Using device:", device)
+
+    dataset, train_data, val_data, test_data = get_dataset(args.dataset)
+
+    g = dataset[0]
+    g = g.to(device)
+
+    input_size, representation_size = (
+        g.ndata["feat"].size(1),
+        args.graph_encoder_layer[-1],
+    )
+
+    # prepare transforms
+    transform_1 = get_graph_drop_transform(
+        drop_edge_p=args.drop_edge_p[0], feat_mask_p=args.feat_mask_p[0]
+    )
+    transform_2 = get_graph_drop_transform(
+        drop_edge_p=args.drop_edge_p[1], feat_mask_p=args.feat_mask_p[1]
+    )
+
+    # scheduler
+    lr_scheduler = CosineDecayScheduler(
+        args.lr, args.lr_warmup_epochs, args.epochs
+    )
+    mm_scheduler = CosineDecayScheduler(1 - args.mm, 0, args.epochs)
+
+    # build networks
+    if args.dataset == "ppi":
+        encoder = GraphSAGE_GCN([input_size] + args.graph_encoder_layer)
+    else:
+        encoder = GCN([input_size] + args.graph_encoder_layer)
+    predictor = MLP_Predictor(
+        representation_size,
+        representation_size,
+        hidden_size=args.predictor_hidden_size,
+    )
+    model = BGRL(encoder, predictor).to(device)
+
+    # optimizer
+    optimizer = AdamW(
+        model.trainable_parameters(), lr=args.lr, weight_decay=args.weight_decay
+    )
+
+    # train
+    for epoch in tqdm(range(1, args.epochs + 1), desc="  - (Training)  "):
+        train(
+            epoch - 1,
+            model,
+            optimizer,
+            lr_scheduler,
+            mm_scheduler,
+            transform_1,
+            transform_2,
+            g,
+            args,
+        )
+        if epoch % args.eval_epochs == 0:
+            val_scores, test_scores = eval(
+                model, dataset, device, args, train_data, val_data, test_data
+            )
+            if args.dataset == "ppi":
+                print(
+                    "Epoch: {:04d} | Best Val F1: {:.4f} | Test F1: {:.4f}".format(
+                        epoch, np.mean(val_scores), np.mean(test_scores)
+                    )
+                )
+            else:
+                print(
+                    "Epoch: {:04d} | Test Accuracy: {:.4f}".format(
+                        epoch, np.mean(test_scores)
+                    )
+                )
+
+    # save encoder weights
+    if not os.path.isdir(args.weights_dir):
+        os.mkdir(args.weights_dir)
+    torch.save(
+        {"model": model.online_encoder.state_dict()},
+        os.path.join(args.weights_dir, "bgrl-{}.pt".format(args.dataset)),
+    )
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+
+    # Dataset options.
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="amazon_photos",
+        choices=[
+            "coauthor_cs",
+            "coauthor_physics",
+            "amazon_photos",
+            "amazon_computers",
+            "wiki_cs",
+            "ppi",
+        ],
+    )
+
+    # Model options.
+    parser.add_argument(
+        "--graph_encoder_layer", type=int, nargs="+", default=[256, 128]
+    )
+    parser.add_argument("--predictor_hidden_size", type=int, default=512)
+
+    # Training options.
+    parser.add_argument("--epochs", type=int, default=10000)
+    parser.add_argument("--lr", type=float, default=1e-5)
+    parser.add_argument("--weight_decay", type=float, default=1e-5)
+    parser.add_argument("--mm", type=float, default=0.99)
+    parser.add_argument("--lr_warmup_epochs", type=int, default=1000)
+    parser.add_argument("--weights_dir", type=str, default="../weights")
+
+    # Augmentations options.
+    parser.add_argument(
+        "--drop_edge_p", type=float, nargs="+", default=[0.0, 0.0]
+    )
+    parser.add_argument(
+        "--feat_mask_p", type=float, nargs="+", default=[0.0, 0.0]
+    )
+
+    # Evaluation options.
+    parser.add_argument("--eval_epochs", type=int, default=250)
+    parser.add_argument("--num_eval_splits", type=int, default=20)
+    parser.add_argument("--data_seed", type=int, default=1)
+
+    # Experiment options.
+    parser.add_argument("--num_experiments", type=int, default=20)
+
+    args = parser.parse_args()
+
+    main(args)

examples/pytorch/bgrl/model.py

@@ -0,0 +1,277 @@
+import copy
+
+import dgl
+
+import torch
+from dgl.nn.pytorch.conv import GraphConv, SAGEConv
+from torch import nn
+from torch.nn import BatchNorm1d, Parameter
+from torch.nn.init import ones_, zeros_
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, in_channels, eps=1e-5, affine=True):
+        super().__init__()
+        self.in_channels = in_channels
+        self.eps = eps
+
+        if affine:
+            self.weight = Parameter(torch.Tensor(in_channels))
+            self.bias = Parameter(torch.Tensor(in_channels))
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        ones_(self.weight)
+        zeros_(self.bias)
+
+    def forward(self, x, batch=None):
+        device = x.device
+        if batch is None:
+            x = x - x.mean()
+            out = x / (x.std(unbiased=False) + self.eps)
+        else:
+            batch_size = int(batch.max()) + 1
+            batch_idx = [batch == i for i in range(batch_size)]
+            norm = (
+                torch.tensor([i.sum() for i in batch_idx], dtype=x.dtype)
+                .clamp_(min=1)
+                .to(device)
+            )
+            norm = norm.mul_(x.size(-1)).view(-1, 1)
+            tmp_list = [x[i] for i in batch_idx]
+            mean = (
+                torch.concat([i.sum(0).unsqueeze(0) for i in tmp_list], dim=0)
+                .sum(dim=-1, keepdim=True)
+                .to(device)
+            )
+            mean = mean / norm
+            x = x - mean.index_select(0, batch.long())
+            var = (
+                torch.concat(
+                    [(i * i).sum(0).unsqueeze(0) for i in tmp_list], dim=0
+                )
+                .sum(dim=-1, keepdim=True)
+                .to(device)
+            )
+            var = var / norm
+            out = x / (var + self.eps).sqrt().index_select(0, batch.long())
+
+        if self.weight is not None and self.bias is not None:
+            out = out * self.weight + self.bias
+
+        return out
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.in_channels})"
+
+
+class MLP_Predictor(nn.Module):
+    r"""MLP used for predictor. The MLP has one hidden layer.
+    Args:
+        input_size (int): Size of input features.
+        output_size (int): Size of output features.
+        hidden_size (int, optional): Size of hidden layer. (default: :obj:`4096`).
+    """
+
+    def __init__(self, input_size, output_size, hidden_size=512):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Linear(input_size, hidden_size, bias=True),
+            nn.PReLU(1),
+            nn.Linear(hidden_size, output_size, bias=True),
+        )
+        self.reset_parameters()
+
+    def forward(self, x):
+        return self.net(x)
+
+    def reset_parameters(self):
+        # kaiming_uniform
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                m.reset_parameters()
+
+
+class GCN(nn.Module):
+    def __init__(self, layer_sizes, batch_norm_mm=0.99):
+        super(GCN, self).__init__()
+
+        self.layers = nn.ModuleList()
+        for in_dim, out_dim in zip(layer_sizes[:-1], layer_sizes[1:]):
+            self.layers.append(GraphConv(in_dim, out_dim))
+            self.layers.append(BatchNorm1d(out_dim, momentum=batch_norm_mm))
+            self.layers.append(nn.PReLU())
+
+    def forward(self, g):
+        x = g.ndata["feat"]
+        for layer in self.layers:
+            if isinstance(layer, GraphConv):
+                x = layer(g, x)
+            else:
+                x = layer(x)
+        return x
+
+    def reset_parameters(self):
+        for layer in self.layers:
+            if hasattr(layer, "reset_parameters"):
+                layer.reset_parameters()
+
+
+class GraphSAGE_GCN(nn.Module):
+    def __init__(self, layer_sizes):
+        super().__init__()
+
+        input_size, hidden_size, embedding_size = layer_sizes
+
+        self.convs = nn.ModuleList(
+            [
+                SAGEConv(input_size, hidden_size, "mean"),
+                SAGEConv(hidden_size, hidden_size, "mean"),
+                SAGEConv(hidden_size, embedding_size, "mean"),
+            ]
+        )
+
+        self.skip_lins = nn.ModuleList(
+            [
+                nn.Linear(input_size, hidden_size, bias=False),
+                nn.Linear(input_size, hidden_size, bias=False),
+            ]
+        )
+
+        self.layer_norms = nn.ModuleList(
+            [
+                LayerNorm(hidden_size),
+                LayerNorm(hidden_size),
+                LayerNorm(embedding_size),
+            ]
+        )
+
+        self.activations = nn.ModuleList(
+            [
+                nn.PReLU(),
+                nn.PReLU(),
+                nn.PReLU(),
+            ]
+        )
+
+    def forward(self, g):
+        x = g.ndata["feat"]
+        if "batch" in g.ndata.keys():
+            batch = g.ndata["batch"]
+        else:
+            batch = None
+
+        h1 = self.convs[0](g, x)
+        h1 = self.layer_norms[0](h1, batch)
+        h1 = self.activations[0](h1)
+
+        x_skip_1 = self.skip_lins[0](x)
+        h2 = self.convs[1](g, h1 + x_skip_1)
+        h2 = self.layer_norms[1](h2, batch)
+        h2 = self.activations[1](h2)
+
+        x_skip_2 = self.skip_lins[1](x)
+        ret = self.convs[2](g, h1 + h2 + x_skip_2)
+        ret = self.layer_norms[2](ret, batch)
+        ret = self.activations[2](ret)
+        return ret
+
+    def reset_parameters(self):
+        for m in self.convs:
+            m.reset_parameters()
+        for m in self.skip_lins:
+            m.reset_parameters()
+        for m in self.activations:
+            m.weight.data.fill_(0.25)
+        for m in self.layer_norms:
+            m.reset_parameters()
+
+
+class BGRL(nn.Module):
+    r"""BGRL architecture for Graph representation learning.
+    Args:
+        encoder (torch.nn.Module): Encoder network to be duplicated and used in both online and target networks.
+        predictor (torch.nn.Module): Predictor network used to predict the target projection from the online projection.
+    .. note::
+        `encoder` must have a `reset_parameters` method, as the weights of the target network will be initialized
+        differently from the online network.
+    """
+
+    def __init__(self, encoder, predictor):
+        super(BGRL, self).__init__()
+        # online network
+        self.online_encoder = encoder
+        self.predictor = predictor
+
+        # target network
+        self.target_encoder = copy.deepcopy(encoder)
+
+        # reinitialize weights
+        self.target_encoder.reset_parameters()
+
+        # stop gradient
+        for param in self.target_encoder.parameters():
+            param.requires_grad = False
+
+    def trainable_parameters(self):
+        r"""Returns the parameters that will be updated via an optimizer."""
+        return list(self.online_encoder.parameters()) + list(
+            self.predictor.parameters()
+        )
+
+    @torch.no_grad()
+    def update_target_network(self, mm):
+        r"""Performs a momentum update of the target network's weights.
+        Args:
+            mm (float): Momentum used in moving average update.
+        """
+        for param_q, param_k in zip(
+            self.online_encoder.parameters(), self.target_encoder.parameters()
+        ):
+            param_k.data.mul_(mm).add_(param_q.data, alpha=1.0 - mm)
+
+    def forward(self, online_x, target_x):
+        # forward online network
+        online_y = self.online_encoder(online_x)
+
+        # prediction
+        online_q = self.predictor(online_y)
+
+        # forward target network
+        with torch.no_grad():
+            target_y = self.target_encoder(target_x).detach()
+        return online_q, target_y
+
+
+def compute_representations(net, dataset, device):
+    r"""Pre-computes the representations for the entire data.
+    Returns:
+        [torch.Tensor, torch.Tensor]: Representations and labels.
+    """
+    net.eval()
+    reps = []
+    labels = []
+
+    if len(dataset) == 1:
+        g = dataset[0]
+        g = dgl.add_self_loop(g)
+        g = g.to(device)
+        with torch.no_grad():
+            reps.append(net(g))
+            labels.append(g.ndata["label"])
+    else:
+        for g in dataset:
+            # forward
+            g = g.to(device)
+            with torch.no_grad():
+                reps.append(net(g))
+                labels.append(g.ndata["label"])
+
+    reps = torch.cat(reps, dim=0)
+    labels = torch.cat(labels, dim=0)
+    return [reps, labels]

examples/pytorch/bgrl/utils.py

@@ -0,0 +1,105 @@
+import copy
+
+import numpy as np
+import torch
+
+from dgl.data import (
+    AmazonCoBuyComputerDataset,
+    AmazonCoBuyPhotoDataset,
+    CoauthorCSDataset,
+    CoauthorPhysicsDataset,
+    PPIDataset,
+    WikiCSDataset,
+)
+from dgl.dataloading import GraphDataLoader
+from dgl.transforms import Compose, DropEdge, FeatMask, RowFeatNormalizer
+
+
+class CosineDecayScheduler:
+    def __init__(self, max_val, warmup_steps, total_steps):
+        self.max_val = max_val
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+
+    def get(self, step):
+        if step < self.warmup_steps:
+            return self.max_val * step / self.warmup_steps
+        elif self.warmup_steps <= step <= self.total_steps:
+            return (
+                self.max_val
+                * (
+                    1
+                    + np.cos(
+                        (step - self.warmup_steps)
+                        * np.pi
+                        / (self.total_steps - self.warmup_steps)
+                    )
+                )
+                / 2
+            )
+        else:
+            raise ValueError(
+                "Step ({}) > total number of steps ({}).".format(
+                    step, self.total_steps
+                )
+            )
+
+
+def get_graph_drop_transform(drop_edge_p, feat_mask_p):
+    transforms = list()
+
+    # make copy of graph
+    transforms.append(copy.deepcopy)
+
+    # drop edges
+    if drop_edge_p > 0.0:
+        transforms.append(DropEdge(drop_edge_p))
+
+    # drop features
+    if feat_mask_p > 0.0:
+        transforms.append(FeatMask(feat_mask_p, node_feat_names=["feat"]))
+
+    return Compose(transforms)
+
+
+def get_wiki_cs(transform=RowFeatNormalizer(subtract_min=True)):
+    dataset = WikiCSDataset(transform=transform)
+    g = dataset[0]
+    std, mean = torch.std_mean(g.ndata["feat"], dim=0, unbiased=False)
+    g.ndata["feat"] = (g.ndata["feat"] - mean) / std
+
+    return [g]
+
+
+def get_ppi():
+    train_dataset = PPIDataset(mode="train")
+    val_dataset = PPIDataset(mode="valid")
+    test_dataset = PPIDataset(mode="test")
+    train_val_dataset = [i for i in train_dataset] + [i for i in val_dataset]
+    for idx, data in enumerate(train_val_dataset):
+        data.ndata["batch"] = torch.zeros(data.num_nodes()) + idx
+        data.ndata["batch"] = data.ndata["batch"].long()
+
+    g = list(GraphDataLoader(train_val_dataset, batch_size=22, shuffle=True))
+
+    return g, PPIDataset(mode="train"), PPIDataset(mode="valid"), test_dataset
+
+
+def get_dataset(name, transform=RowFeatNormalizer(subtract_min=True)):
+    dgl_dataset_dict = {
+        "coauthor_cs": CoauthorCSDataset,
+        "coauthor_physics": CoauthorPhysicsDataset,
+        "amazon_computers": AmazonCoBuyComputerDataset,
+        "amazon_photos": AmazonCoBuyPhotoDataset,
+        "wiki_cs": get_wiki_cs,
+        "ppi": get_ppi,
+    }
+
+    dataset_class = dgl_dataset_dict[name]
+    train_data, val_data, test_data = None, None, None
+    if name != "ppi":
+        dataset = dataset_class(transform=transform)
+    else:
+        dataset, train_data, val_data, test_data = dataset_class()
+
+    return dataset, train_data, val_data, test_data

examples/pytorch/capsule/DGLDigitCapsule.py

@@ -0,0 +1,57 @@
+import dgl
+import dgl.function as fn
+import torch
+from DGLRoutingLayer import DGLRoutingLayer
+from torch import nn
+from torch.nn import functional as F
+
+
+class DGLDigitCapsuleLayer(nn.Module):
+    def __init__(
+        self,
+        in_nodes_dim=8,
+        in_nodes=1152,
+        out_nodes=10,
+        out_nodes_dim=16,
+        device="cpu",
+    ):
+        super(DGLDigitCapsuleLayer, self).__init__()
+        self.device = device
+        self.in_nodes_dim, self.out_nodes_dim = in_nodes_dim, out_nodes_dim
+        self.in_nodes, self.out_nodes = in_nodes, out_nodes
+        self.weight = nn.Parameter(
+            torch.randn(in_nodes, out_nodes, out_nodes_dim, in_nodes_dim)
+        )
+
+    def forward(self, x):
+        self.batch_size = x.size(0)
+        u_hat = self.compute_uhat(x)
+        routing = DGLRoutingLayer(
+            self.in_nodes,
+            self.out_nodes,
+            self.out_nodes_dim,
+            batch_size=self.batch_size,
+            device=self.device,
+        )
+        routing(u_hat, routing_num=3)
+        out_nodes_feature = routing.g.nodes[routing.out_indx].data["v"]
+        # shape transformation is for further classification
+        return (
+            out_nodes_feature.transpose(0, 1)
+            .unsqueeze(1)
+            .unsqueeze(4)
+            .squeeze(1)
+        )
+
+    def compute_uhat(self, x):
+        # x is the input vextor with shape [batch_size, in_nodes_dim, in_nodes]
+        # Transpose x to [batch_size, in_nodes, in_nodes_dim]
+        x = x.transpose(1, 2)
+        # Expand x to [batch_size, in_nodes, out_nodes, in_nodes_dim, 1]
+        x = torch.stack([x] * self.out_nodes, dim=2).unsqueeze(4)
+        # Expand W from [in_nodes, out_nodes, in_nodes_dim, out_nodes_dim]
+        # to [batch_size, in_nodes, out_nodes, out_nodes_dim, in_nodes_dim]
+        W = self.weight.expand(self.batch_size, *self.weight.size())
+        # u_hat's shape is [in_nodes, out_nodes, batch_size, out_nodes_dim]
+        u_hat = torch.matmul(W, x).permute(1, 2, 0, 3, 4).squeeze().contiguous()
+        return u_hat.view(-1, self.batch_size, self.out_nodes_dim)

examples/pytorch/capsule/DGLRoutingLayer.py

@@ -0,0 +1,83 @@
+import dgl
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DGLRoutingLayer(nn.Module):
+    def __init__(self, in_nodes, out_nodes, f_size, batch_size=0, device="cpu"):
+        super(DGLRoutingLayer, self).__init__()
+        self.batch_size = batch_size
+        self.g = init_graph(in_nodes, out_nodes, f_size, device=device)
+        self.in_nodes = in_nodes
+        self.out_nodes = out_nodes
+        self.in_indx = list(range(in_nodes))
+        self.out_indx = list(range(in_nodes, in_nodes + out_nodes))
+        self.device = device
+
+    def forward(self, u_hat, routing_num=1):
+        self.g.edata["u_hat"] = u_hat
+        batch_size = self.batch_size
+
+        # step 2 (line 5)
+        def cap_message(edges):
+            if batch_size:
+                return {"m": edges.data["c"].unsqueeze(1) * edges.data["u_hat"]}
+            else:
+                return {"m": edges.data["c"] * edges.data["u_hat"]}
+
+        def cap_reduce(nodes):
+            return {"s": th.sum(nodes.mailbox["m"], dim=1)}
+
+        for r in range(routing_num):
+            # step 1 (line 4): normalize over out edges
+            edges_b = self.g.edata["b"].view(self.in_nodes, self.out_nodes)
+            self.g.edata["c"] = F.softmax(edges_b, dim=1).view(-1, 1)
+
+            # Execute step 1 & 2
+            self.g.update_all(message_func=cap_message, reduce_func=cap_reduce)
+
+            # step 3 (line 6)
+            if self.batch_size:
+                self.g.nodes[self.out_indx].data["v"] = squash(
+                    self.g.nodes[self.out_indx].data["s"], dim=2
+                )
+            else:
+                self.g.nodes[self.out_indx].data["v"] = squash(
+                    self.g.nodes[self.out_indx].data["s"], dim=1
+                )
+            # step 4 (line 7)
+            v = th.cat(
+                [self.g.nodes[self.out_indx].data["v"]] * self.in_nodes, dim=0
+            )
+            if self.batch_size:
+                self.g.edata["b"] = self.g.edata["b"] + (
+                    self.g.edata["u_hat"] * v
+                ).mean(dim=1).sum(dim=1, keepdim=True)
+            else:
+                self.g.edata["b"] = self.g.edata["b"] + (
+                    self.g.edata["u_hat"] * v
+                ).sum(dim=1, keepdim=True)
+
+
+def squash(s, dim=1):
+    sq = th.sum(s**2, dim=dim, keepdim=True)
+    s_norm = th.sqrt(sq)
+    s = (sq / (1.0 + sq)) * (s / s_norm)
+    return s
+
+
+def init_graph(in_nodes, out_nodes, f_size, device="cpu"):
+    src, dst = [], []
+    in_indx = list(range(in_nodes))
+    out_indx = list(range(in_nodes, in_nodes + out_nodes))
+    # add edges use edge broadcasting
+    for u in in_indx:
+        src += [u] * len(out_indx)
+        dst += out_indx
+
+    g = dgl.graph((src, dst))  # dgl.graph once;
+    g.set_n_initializer(dgl.frame.zero_initializer)
+    g = g.to(device)
+    g.edata["b"] = th.zeros(in_nodes * out_nodes, 1).to(device)
+    return g

examples/pytorch/capsule/README.md

@@ -0,0 +1,23 @@
+DGL implementation of Capsule Network
+=====================================
+
+This repo implements Hinton and his team's [Capsule Network](https://arxiv.org/abs/1710.09829).
+Only margin loss is implemented, for simplicity to understand the DGL.
+
+Dependencies
+--------------
+* PyTorch 0.4.1+
+* torchvision
+
+```bash
+pip install torch torchvision
+```
+
+Training & Evaluation
+----------------------
+```bash
+# Run with default config
+python3 main.py
+# Run with train and test batch size 128, and for 50 epochs
+python3 main.py --batch-size 128 --test-batch-size 128 --epochs 50
+```

examples/pytorch/capsule/main.py

@@ -0,0 +1,157 @@
+import argparse
+
+import torch
+import torch.optim as optim
+from model import Net
+from torchvision import datasets, transforms
+
+
+def train(args, model, device, train_loader, optimizer, epoch):
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = model.margin_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            print(
+                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                    epoch,
+                    batch_idx * len(data),
+                    len(train_loader.dataset),
+                    100.0 * batch_idx / len(train_loader),
+                    loss.item(),
+                )
+            )
+
+
+def test(args, model, device, test_loader):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            test_loss += model.margin_loss(
+                output, target
+            ).item()  # sum up batch loss
+            pred = (
+                output.norm(dim=2).squeeze().max(1, keepdim=True)[1]
+            )  # get the index of the max log-probability
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+    print(
+        "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
+            test_loss,
+            correct,
+            len(test_loader.dataset),
+            100.0 * correct / len(test_loader.dataset),
+        )
+    )
+
+
+def main():
+    # Training settings
+    parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=512,
+        metavar="N",
+        help="input batch size for training (default: 64)",
+    )
+    parser.add_argument(
+        "--test-batch-size",
+        type=int,
+        default=512,
+        metavar="N",
+        help="input batch size for testing (default: 1000)",
+    )
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=10,
+        metavar="N",
+        help="number of epochs to train (default: 10)",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.01,
+        metavar="LR",
+        help="learning rate (default: 0.01)",
+    )
+    parser.add_argument(
+        "--no-cuda",
+        action="store_true",
+        default=False,
+        help="disables CUDA training",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=1,
+        metavar="S",
+        help="random seed (default: 1)",
+    )
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+    args = parser.parse_args()
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+
+    torch.manual_seed(args.seed)
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    kwargs = {"num_workers": 1, "pin_memory": True} if use_cuda else {}
+    train_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(
+            "../data",
+            train=True,
+            download=True,
+            transform=transforms.Compose(
+                [
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.1307,), (0.3081,)),
+                ]
+            ),
+        ),
+        batch_size=args.batch_size,
+        shuffle=True,
+        **kwargs
+    )
+    test_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(
+            "../data",
+            train=False,
+            transform=transforms.Compose(
+                [
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.1307,), (0.3081,)),
+                ]
+            ),
+        ),
+        batch_size=args.test_batch_size,
+        shuffle=True,
+        **kwargs
+    )
+
+    model = Net(device=device).to(device)
+    optimizer = optim.Adam(model.parameters(), lr=args.lr)
+
+    for epoch in range(1, args.epochs + 1):
+        train(args, model, device, train_loader, optimizer, epoch)
+        test(args, model, device, test_loader)
+
+
+if __name__ == "__main__":
+    main()

examples/pytorch/capsule/model.py

@@ -0,0 +1,59 @@
+import torch
+from DGLDigitCapsule import DGLDigitCapsuleLayer
+from DGLRoutingLayer import squash
+from torch import nn
+
+
+class Net(nn.Module):
+    def __init__(self, device="cpu"):
+        super(Net, self).__init__()
+        self.device = device
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_channels=1, out_channels=256, kernel_size=9, stride=1),
+            nn.ReLU(inplace=True),
+        )
+
+        self.primary = PrimaryCapsuleLayer(device=device)
+        self.digits = DGLDigitCapsuleLayer(device=device)
+
+    def forward(self, x):
+        out_conv1 = self.conv1(x)
+        out_primary_caps = self.primary(out_conv1)
+        out_digit_caps = self.digits(out_primary_caps)
+        return out_digit_caps
+
+    def margin_loss(self, input, target):
+        batch_s = target.size(0)
+        one_hot_vec = torch.zeros(batch_s, 10).to(self.device)
+        for i in range(batch_s):
+            one_hot_vec[i, target[i]] = 1.0
+        batch_size = input.size(0)
+        v_c = torch.sqrt((input**2).sum(dim=2, keepdim=True))
+        zero = torch.zeros(1).to(self.device)
+        m_plus = 0.9
+        m_minus = 0.1
+        loss_lambda = 0.5
+        max_left = torch.max(m_plus - v_c, zero).view(batch_size, -1) ** 2
+        max_right = torch.max(v_c - m_minus, zero).view(batch_size, -1) ** 2
+        t_c = one_hot_vec
+        l_c = t_c * max_left + loss_lambda * (1.0 - t_c) * max_right
+        l_c = l_c.sum(dim=1)
+        return l_c.mean()
+
+
+class PrimaryCapsuleLayer(nn.Module):
+    def __init__(self, in_channel=256, num_unit=8, device="cpu"):
+        super(PrimaryCapsuleLayer, self).__init__()
+        self.in_channel = in_channel
+        self.num_unit = num_unit
+        self.deivce = device
+        self.conv_units = nn.ModuleList(
+            [nn.Conv2d(self.in_channel, 32, 9, 2) for _ in range(self.num_unit)]
+        )
+
+    def forward(self, x):
+        unit = [self.conv_units[i](x) for i, l in enumerate(self.conv_units)]
+        unit = torch.stack(unit, dim=1)
+        batch_size = x.size(0)
+        unit = unit.view(batch_size, 8, -1)
+        return squash(unit, dim=2)

examples/pytorch/capsule/simple_routing.py

@@ -0,0 +1,40 @@
+import dgl
+import torch as th
+import torch.nn as nn
+from DGLRoutingLayer import DGLRoutingLayer
+from torch.nn import functional as F
+
+g = dgl.DGLGraph()
+g.graph_data = {}
+
+in_nodes = 20
+out_nodes = 10
+g.graph_data["in_nodes"] = in_nodes
+g.graph_data["out_nodes"] = out_nodes
+all_nodes = in_nodes + out_nodes
+g.add_nodes(all_nodes)
+
+
+in_indx = list(range(in_nodes))
+out_indx = list(range(in_nodes, in_nodes + out_nodes))
+g.graph_data["in_indx"] = in_indx
+g.graph_data["out_indx"] = out_indx
+
+# add edges use edge broadcasting
+for u in out_indx:
+    g.add_edges(in_indx, u)
+# init states
+f_size = 4
+g.ndata["v"] = th.zeros(all_nodes, f_size)
+g.edata["u_hat"] = th.randn(in_nodes * out_nodes, f_size)
+g.edata["b"] = th.randn(in_nodes * out_nodes, 1)
+
+routing_layer = DGLRoutingLayer(g)
+
+entropy_list = []
+for i in range(15):
+    routing_layer()
+    dist_matrix = g.edata["c"].view(in_nodes, out_nodes)
+    entropy = (-dist_matrix * th.log(dist_matrix)).sum(dim=0)
+    entropy_list.append(entropy.data.numpy())
+    std = dist_matrix.std(dim=0)

examples/pytorch/caregnn/README.md

@@ -0,0 +1,126 @@
+# DGL Implementation of the CARE-GNN Paper
+
+This DGL example implements the CAmouflage-REsistant GNN (CARE-GNN) model proposed in the paper [Enhancing Graph Neural Network-based Fraud Detectors against Camouflaged Fraudsters](https://arxiv.org/abs/2008.08692). The author's codes of implementation is [here](https://github.com/YingtongDou/CARE-GNN).
+
+**NOTE**: The sampling version of this model has been modified according to the feature of the DGL's NodeDataLoader. For the formula 2 in the paper, rather than using the embedding of the last layer, this version uses the embedding of the current layer in the previous epoch to measure the similarity between center nodes and their neighbors.
+
+Example implementor
+----------------------
+This example was implemented by [Kay Liu](https://github.com/kayzliu) during his SDE intern work at the AWS Shanghai AI Lab.
+
+Dependencies
+----------------------
+- Python 3.7.10
+- PyTorch 1.8.1
+- dgl 0.7.1
+- scikit-learn 0.23.2
+
+Dataset
+---------------------------------------
+The datasets used for node classification are DGL's built-in FraudDataset. The statistics are summarized as followings:
+
+**Amazon**
+
+- Nodes: 11,944
+- Edges:
+    - U-P-U: 351,216
+    - U-S-U: 7,132,958
+    - U-V-U: 2,073,474
+- Classes:
+    - Positive (fraudulent): 821
+    - Negative (benign): 7,818
+    - Unlabeled: 3,305
+- Positive-Negative ratio: 1 : 10.5
+- Node feature size: 25
+
+**YelpChi**
+
+- Nodes: 45,954
+- Edges:
+    - R-U-R: 98,630
+    - R-T-R: 1,147,232
+    - R-S-R: 6,805,486
+- Classes:
+    - Positive (spam): 6,677
+    - Negative (legitimate): 39,277
+- Positive-Negative ratio: 1 : 5.9
+- Node feature size: 32
+
+How to run
+--------------------------------
+To run the full graph version and use early stopping, in the care-gnn folder, run
+```
+python main.py --early-stop
+```
+
+If want to use a GPU, run
+```
+python main.py --gpu 0
+```
+
+To train on Yelp dataset instead of Amazon, run
+```
+python main.py --dataset yelp
+```
+
+To run the sampling version, run
+```
+python main_sampling.py
+```
+
+Performance
+-------------------------
+The result reported by the paper is the best validation results within 30 epochs, and the table below reports the val and test results (same setting in the paper except for the random seed, here `seed=717`). 
+
+<table>
+<thead>
+  <tr>
+    <th colspan="2">Dataset</th>
+    <th>Amazon</th>
+    <th>Yelp</th>
+  </tr>
+</thead>
+<tbody>
+  <tr>
+    <td>Metric (val / test)</td>
+    <td>Max Epoch</td>
+    <td>30</td>
+    <td>30 </td>
+  </tr>
+  <tr>
+    <td rowspan="3">AUC (val/test)</td>
+    <td>paper reported</td>
+    <td>0.8973 / -</td>
+    <td>0.7570 / -</td>
+  </tr>
+  <tr>
+    <td>DGL full graph</td>
+    <td>0.8849 / 0.8922</td>
+    <td>0.6856 / 0.6867</td>
+  </tr>
+  <tr>
+    <td>DGL sampling</td>
+    <td>0.9350 / 0.9331</td>
+    <td>0.7857 / 0.7890</td>
+  </tr>
+  <tr>
+    <td rowspan="3">Recall (val/test)</td>
+    <td>paper reported</td>
+    <td>0.8848 / -</td>
+    <td>0.7192 / -</td>
+  </tr>
+  <tr>
+    <td>DGL full graph</td>
+    <td>0.8615 / 0.8544</td>
+    <td>0.6667/ 0.6619</td>
+  </tr>
+  <tr>
+    <td>DGL sampling</td>
+    <td>0.9130 / 0.9045</td>
+    <td>0.7537 / 0.7540</td>
+  </tr>
+</tbody>
+</table>
+
+
+

examples/pytorch/caregnn/main.py

@@ -0,0 +1,208 @@
+import argparse
+
+import dgl
+
+import torch as th
+import torch.optim as optim
+from model import CAREGNN
+from sklearn.metrics import recall_score, roc_auc_score
+from torch.nn.functional import softmax
+from utils import EarlyStopping
+
+
+def main(args):
+    # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
+    # Load dataset
+    dataset = dgl.data.FraudDataset(args.dataset, train_size=0.4)
+    graph = dataset[0]
+    num_classes = dataset.num_classes
+
+    # check cuda
+    if args.gpu >= 0 and th.cuda.is_available():
+        device = "cuda:{}".format(args.gpu)
+    else:
+        device = "cpu"
+
+    # retrieve labels of ground truth
+    labels = graph.ndata["label"].to(device)
+
+    # Extract node features
+    feat = graph.ndata["feature"].to(device)
+
+    # retrieve masks for train/validation/test
+    train_mask = graph.ndata["train_mask"]
+    val_mask = graph.ndata["val_mask"]
+    test_mask = graph.ndata["test_mask"]
+
+    train_idx = th.nonzero(train_mask, as_tuple=False).squeeze(1).to(device)
+    val_idx = th.nonzero(val_mask, as_tuple=False).squeeze(1).to(device)
+    test_idx = th.nonzero(test_mask, as_tuple=False).squeeze(1).to(device)
+
+    # Reinforcement learning module only for positive training nodes
+    rl_idx = th.nonzero(
+        train_mask.to(device) & labels.bool(), as_tuple=False
+    ).squeeze(1)
+
+    graph = graph.to(device)
+
+    # Step 2: Create model =================================================================== #
+    model = CAREGNN(
+        in_dim=feat.shape[-1],
+        num_classes=num_classes,
+        hid_dim=args.hid_dim,
+        num_layers=args.num_layers,
+        activation=th.tanh,
+        step_size=args.step_size,
+        edges=graph.canonical_etypes,
+    )
+
+    model = model.to(device)
+
+    # Step 3: Create training components ===================================================== #
+    _, cnt = th.unique(labels, return_counts=True)
+    loss_fn = th.nn.CrossEntropyLoss(weight=1 / cnt)
+    optimizer = optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+    )
+    if args.early_stop:
+        stopper = EarlyStopping(patience=100)
+
+    # Step 4: training epochs =============================================================== #
+    for epoch in range(args.max_epoch):
+        # Training and validation using a full graph
+        model.train()
+        logits_gnn, logits_sim = model(graph, feat)
+
+        # compute loss
+        tr_loss = loss_fn(
+            logits_gnn[train_idx], labels[train_idx]
+        ) + args.sim_weight * loss_fn(logits_sim[train_idx], labels[train_idx])
+
+        tr_recall = recall_score(
+            labels[train_idx].cpu(),
+            logits_gnn.data[train_idx].argmax(dim=1).cpu(),
+        )
+        tr_auc = roc_auc_score(
+            labels[train_idx].cpu(),
+            softmax(logits_gnn, dim=1).data[train_idx][:, 1].cpu(),
+        )
+
+        # validation
+        val_loss = loss_fn(
+            logits_gnn[val_idx], labels[val_idx]
+        ) + args.sim_weight * loss_fn(logits_sim[val_idx], labels[val_idx])
+        val_recall = recall_score(
+            labels[val_idx].cpu(), logits_gnn.data[val_idx].argmax(dim=1).cpu()
+        )
+        val_auc = roc_auc_score(
+            labels[val_idx].cpu(),
+            softmax(logits_gnn, dim=1).data[val_idx][:, 1].cpu(),
+        )
+
+        # backward
+        optimizer.zero_grad()
+        tr_loss.backward()
+        optimizer.step()
+
+        # Print out performance
+        print(
+            "Epoch {}, Train: Recall: {:.4f} AUC: {:.4f} Loss: {:.4f} | Val: Recall: {:.4f} AUC: {:.4f} Loss: {:.4f}".format(
+                epoch,
+                tr_recall,
+                tr_auc,
+                tr_loss.item(),
+                val_recall,
+                val_auc,
+                val_loss.item(),
+            )
+        )
+
+        # Adjust p value with reinforcement learning module
+        model.RLModule(graph, epoch, rl_idx)
+
+        if args.early_stop:
+            if stopper.step(val_auc, model):
+                break
+
+    # Test after all epoch
+    model.eval()
+    if args.early_stop:
+        model.load_state_dict(th.load("es_checkpoint.pt"))
+
+    # forward
+    logits_gnn, logits_sim = model.forward(graph, feat)
+
+    # compute loss
+    test_loss = loss_fn(
+        logits_gnn[test_idx], labels[test_idx]
+    ) + args.sim_weight * loss_fn(logits_sim[test_idx], labels[test_idx])
+    test_recall = recall_score(
+        labels[test_idx].cpu(), logits_gnn[test_idx].argmax(dim=1).cpu()
+    )
+    test_auc = roc_auc_score(
+        labels[test_idx].cpu(),
+        softmax(logits_gnn, dim=1).data[test_idx][:, 1].cpu(),
+    )
+
+    print(
+        "Test Recall: {:.4f} AUC: {:.4f} Loss: {:.4f}".format(
+            test_recall, test_auc, test_loss.item()
+        )
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="GCN-based Anti-Spam Model")
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="amazon",
+        help="DGL dataset for this model (yelp, or amazon)",
+    )
+    parser.add_argument(
+        "--gpu", type=int, default=-1, help="GPU index. Default: -1, using CPU."
+    )
+    parser.add_argument(
+        "--hid_dim", type=int, default=64, help="Hidden layer dimension"
+    )
+    parser.add_argument(
+        "--num_layers", type=int, default=1, help="Number of layers"
+    )
+    parser.add_argument(
+        "--max_epoch",
+        type=int,
+        default=30,
+        help="The max number of epochs. Default: 30",
+    )
+    parser.add_argument(
+        "--lr", type=float, default=0.01, help="Learning rate. Default: 0.01"
+    )
+    parser.add_argument(
+        "--weight_decay",
+        type=float,
+        default=0.001,
+        help="Weight decay. Default: 0.001",
+    )
+    parser.add_argument(
+        "--step_size",
+        type=float,
+        default=0.02,
+        help="RL action step size (lambda 2). Default: 0.02",
+    )
+    parser.add_argument(
+        "--sim_weight",
+        type=float,
+        default=2,
+        help="Similarity loss weight (lambda 1). Default: 2",
+    )
+    parser.add_argument(
+        "--early-stop",
+        action="store_true",
+        default=False,
+        help="indicates whether to use early stop",
+    )
+
+    args = parser.parse_args()
+    print(args)
+    th.manual_seed(717)
+    main(args)

examples/pytorch/caregnn/main_sampling.py

@@ -0,0 +1,278 @@
+import argparse
+
+import dgl
+
+import torch as th
+import torch.optim as optim
+from model_sampling import _l1_dist, CAREGNN, CARESampler
+from sklearn.metrics import recall_score, roc_auc_score
+from torch.nn.functional import softmax
+from utils import EarlyStopping
+
+
+def evaluate(model, loss_fn, dataloader, device="cpu"):
+    loss = 0
+    auc = 0
+    recall = 0
+    num_blocks = 0
+    for input_nodes, output_nodes, blocks in dataloader:
+        blocks = [b.to(device) for b in blocks]
+        feature = blocks[0].srcdata["feature"]
+        label = blocks[-1].dstdata["label"]
+        logits_gnn, logits_sim = model(blocks, feature)
+
+        # compute loss
+        loss += (
+            loss_fn(logits_gnn, label).item()
+            + args.sim_weight * loss_fn(logits_sim, label).item()
+        )
+        recall += recall_score(
+            label.cpu(), logits_gnn.argmax(dim=1).detach().cpu()
+        )
+        auc += roc_auc_score(
+            label.cpu(), softmax(logits_gnn, dim=1)[:, 1].detach().cpu()
+        )
+        num_blocks += 1
+
+    return recall / num_blocks, auc / num_blocks, loss / num_blocks
+
+
+def main(args):
+    # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
+    # Load dataset
+    dataset = dgl.data.FraudDataset(args.dataset, train_size=0.4)
+    graph = dataset[0]
+    num_classes = dataset.num_classes
+
+    # check cuda
+    if args.gpu >= 0 and th.cuda.is_available():
+        device = "cuda:{}".format(args.gpu)
+        args.num_workers = 0
+    else:
+        device = "cpu"
+
+    # retrieve labels of ground truth
+    labels = graph.ndata["label"].to(device)
+
+    # Extract node features
+    feat = graph.ndata["feature"].to(device)
+    layers_feat = feat.expand(args.num_layers, -1, -1)
+
+    # retrieve masks for train/validation/test
+    train_mask = graph.ndata["train_mask"]
+    val_mask = graph.ndata["val_mask"]
+    test_mask = graph.ndata["test_mask"]
+
+    train_idx = th.nonzero(train_mask, as_tuple=False).squeeze(1).to(device)
+    val_idx = th.nonzero(val_mask, as_tuple=False).squeeze(1).to(device)
+    test_idx = th.nonzero(test_mask, as_tuple=False).squeeze(1).to(device)
+
+    # Reinforcement learning module only for positive training nodes
+    rl_idx = th.nonzero(
+        train_mask.to(device) & labels.bool(), as_tuple=False
+    ).squeeze(1)
+
+    graph = graph.to(device)
+
+    # Step 2: Create model =================================================================== #
+    model = CAREGNN(
+        in_dim=feat.shape[-1],
+        num_classes=num_classes,
+        hid_dim=args.hid_dim,
+        num_layers=args.num_layers,
+        activation=th.tanh,
+        step_size=args.step_size,
+        edges=graph.canonical_etypes,
+    )
+
+    model = model.to(device)
+
+    # Step 3: Create training components ===================================================== #
+    _, cnt = th.unique(labels, return_counts=True)
+    loss_fn = th.nn.CrossEntropyLoss(weight=1 / cnt)
+    optimizer = optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+    )
+    if args.early_stop:
+        stopper = EarlyStopping(patience=100)
+
+    # Step 4: training epochs =============================================================== #
+    for epoch in range(args.max_epoch):
+        # calculate the distance of each edges and sample based on the distance
+        dists = []
+        p = []
+        for i in range(args.num_layers):
+            dist = {}
+            graph.ndata["nd"] = th.tanh(model.layers[i].MLP(layers_feat[i]))
+            for etype in graph.canonical_etypes:
+                graph.apply_edges(_l1_dist, etype=etype)
+                dist[etype] = graph.edges[etype].data.pop("ed").detach().cpu()
+            dists.append(dist)
+            p.append(model.layers[i].p)
+        graph.ndata.pop("nd")
+        sampler = CARESampler(p, dists, args.num_layers)
+
+        # train
+        model.train()
+        tr_loss = 0
+        tr_recall = 0
+        tr_auc = 0
+        tr_blk = 0
+        train_dataloader = dgl.dataloading.DataLoader(
+            graph,
+            train_idx,
+            sampler,
+            batch_size=args.batch_size,
+            shuffle=True,
+            drop_last=False,
+            num_workers=args.num_workers,
+        )
+
+        for input_nodes, output_nodes, blocks in train_dataloader:
+            blocks = [b.to(device) for b in blocks]
+            train_feature = blocks[0].srcdata["feature"]
+            train_label = blocks[-1].dstdata["label"]
+            logits_gnn, logits_sim = model(blocks, train_feature)
+
+            # compute loss
+            blk_loss = loss_fn(
+                logits_gnn, train_label
+            ) + args.sim_weight * loss_fn(logits_sim, train_label)
+            tr_loss += blk_loss.item()
+            tr_recall += recall_score(
+                train_label.cpu(), logits_gnn.argmax(dim=1).detach().cpu()
+            )
+            tr_auc += roc_auc_score(
+                train_label.cpu(),
+                softmax(logits_gnn, dim=1)[:, 1].detach().cpu(),
+            )
+            tr_blk += 1
+
+            # backward
+            optimizer.zero_grad()
+            blk_loss.backward()
+            optimizer.step()
+
+        # Reinforcement learning module
+        model.RLModule(graph, epoch, rl_idx, dists)
+
+        # validation
+        model.eval()
+        val_dataloader = dgl.dataloading.DataLoader(
+            graph,
+            val_idx,
+            sampler,
+            batch_size=args.batch_size,
+            shuffle=True,
+            drop_last=False,
+            num_workers=args.num_workers,
+        )
+
+        val_recall, val_auc, val_loss = evaluate(
+            model, loss_fn, val_dataloader, device
+        )
+
+        # Print out performance
+        print(
+            "In epoch {}, Train Recall: {:.4f} | Train AUC: {:.4f} | Train Loss: {:.4f}; "
+            "Valid Recall: {:.4f} | Valid AUC: {:.4f} | Valid loss: {:.4f}".format(
+                epoch,
+                tr_recall / tr_blk,
+                tr_auc / tr_blk,
+                tr_loss / tr_blk,
+                val_recall,
+                val_auc,
+                val_loss,
+            )
+        )
+
+        if args.early_stop:
+            if stopper.step(val_auc, model):
+                break
+
+    # Test with mini batch after all epoch
+    model.eval()
+    if args.early_stop:
+        model.load_state_dict(th.load("es_checkpoint.pt"))
+    test_dataloader = dgl.dataloading.DataLoader(
+        graph,
+        test_idx,
+        sampler,
+        batch_size=args.batch_size,
+        shuffle=True,
+        drop_last=False,
+        num_workers=args.num_workers,
+    )
+
+    test_recall, test_auc, test_loss = evaluate(
+        model, loss_fn, test_dataloader, device
+    )
+
+    print(
+        "Test Recall: {:.4f} | Test AUC: {:.4f} | Test loss: {:.4f}".format(
+            test_recall, test_auc, test_loss
+        )
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="GCN-based Anti-Spam Model")
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="amazon",
+        help="DGL dataset for this model (yelp, or amazon)",
+    )
+    parser.add_argument(
+        "--gpu", type=int, default=-1, help="GPU index. Default: -1, using CPU."
+    )
+    parser.add_argument(
+        "--hid_dim", type=int, default=64, help="Hidden layer dimension"
+    )
+    parser.add_argument(
+        "--num_layers", type=int, default=1, help="Number of layers"
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=256, help="Size of mini-batch"
+    )
+    parser.add_argument(
+        "--max_epoch",
+        type=int,
+        default=30,
+        help="The max number of epochs. Default: 30",
+    )
+    parser.add_argument(
+        "--lr", type=float, default=0.01, help="Learning rate. Default: 0.01"
+    )
+    parser.add_argument(
+        "--weight_decay",
+        type=float,
+        default=0.001,
+        help="Weight decay. Default: 0.001",
+    )
+    parser.add_argument(
+        "--step_size",
+        type=float,
+        default=0.02,
+        help="RL action step size (lambda 2). Default: 0.02",
+    )
+    parser.add_argument(
+        "--sim_weight",
+        type=float,
+        default=2,
+        help="Similarity loss weight (lambda 1). Default: 0.001",
+    )
+    parser.add_argument(
+        "--num_workers", type=int, default=4, help="Number of node dataloader"
+    )
+    parser.add_argument(
+        "--early-stop",
+        action="store_true",
+        default=False,
+        help="indicates whether to use early stop",
+    )
+
+    args = parser.parse_args()
+    th.manual_seed(717)
+    print(args)
+    main(args)

examples/pytorch/caregnn/model.py

@@ -0,0 +1,208 @@
+import dgl.function as fn
+import numpy as np
+import torch as th
+import torch.nn as nn
+
+
+class CAREConv(nn.Module):
+    """One layer of CARE-GNN."""
+
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        num_classes,
+        edges,
+        activation=None,
+        step_size=0.02,
+    ):
+        super(CAREConv, self).__init__()
+
+        self.activation = activation
+        self.step_size = step_size
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.num_classes = num_classes
+        self.edges = edges
+        self.dist = {}
+
+        self.linear = nn.Linear(self.in_dim, self.out_dim)
+        self.MLP = nn.Linear(self.in_dim, self.num_classes)
+
+        self.p = {}
+        self.last_avg_dist = {}
+        self.f = {}
+        self.cvg = {}
+        for etype in edges:
+            self.p[etype] = 0.5
+            self.last_avg_dist[etype] = 0
+            self.f[etype] = []
+            self.cvg[etype] = False
+
+    def _calc_distance(self, edges):
+        # formula 2
+        d = th.norm(
+            th.tanh(self.MLP(edges.src["h"]))
+            - th.tanh(self.MLP(edges.dst["h"])),
+            1,
+            1,
+        )
+        return {"d": d}
+
+    def _top_p_sampling(self, g, p):
+        # this implementation is low efficient
+        # optimization requires dgl.sampling.select_top_p requested in issue #3100
+        dist = g.edata["d"]
+        neigh_list = []
+        for node in g.nodes():
+            edges = g.in_edges(node, form="eid")
+            num_neigh = th.ceil(g.in_degrees(node) * p).int().item()
+            neigh_dist = dist[edges]
+            if neigh_dist.shape[0] > num_neigh:
+                neigh_index = np.argpartition(
+                    neigh_dist.cpu().detach(), num_neigh
+                )[:num_neigh]
+            else:
+                neigh_index = np.arange(num_neigh)
+            neigh_list.append(edges[neigh_index])
+        return th.cat(neigh_list)
+
+    def forward(self, g, feat):
+        with g.local_scope():
+            g.ndata["h"] = feat
+
+            hr = {}
+            for i, etype in enumerate(g.canonical_etypes):
+                g.apply_edges(self._calc_distance, etype=etype)
+                self.dist[etype] = g.edges[etype].data["d"]
+                sampled_edges = self._top_p_sampling(g[etype], self.p[etype])
+
+                # formula 8
+                g.send_and_recv(
+                    sampled_edges,
+                    fn.copy_u("h", "m"),
+                    fn.mean("m", "h_%s" % etype[1]),
+                    etype=etype,
+                )
+                hr[etype] = g.ndata["h_%s" % etype[1]]
+                if self.activation is not None:
+                    hr[etype] = self.activation(hr[etype])
+
+            # formula 9 using mean as inter-relation aggregator
+            p_tensor = (
+                th.Tensor(list(self.p.values())).view(-1, 1, 1).to(g.device)
+            )
+            h_homo = th.sum(th.stack(list(hr.values())) * p_tensor, dim=0)
+            h_homo += feat
+            if self.activation is not None:
+                h_homo = self.activation(h_homo)
+
+            return self.linear(h_homo)
+
+
+class CAREGNN(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        num_classes,
+        hid_dim=64,
+        edges=None,
+        num_layers=2,
+        activation=None,
+        step_size=0.02,
+    ):
+        super(CAREGNN, self).__init__()
+        self.in_dim = in_dim
+        self.hid_dim = hid_dim
+        self.num_classes = num_classes
+        self.edges = edges
+        self.activation = activation
+        self.step_size = step_size
+        self.num_layers = num_layers
+
+        self.layers = nn.ModuleList()
+
+        if self.num_layers == 1:
+            # Single layer
+            self.layers.append(
+                CAREConv(
+                    self.in_dim,
+                    self.num_classes,
+                    self.num_classes,
+                    self.edges,
+                    activation=self.activation,
+                    step_size=self.step_size,
+                )
+            )
+
+        else:
+            # Input layer
+            self.layers.append(
+                CAREConv(
+                    self.in_dim,
+                    self.hid_dim,
+                    self.num_classes,
+                    self.edges,
+                    activation=self.activation,
+                    step_size=self.step_size,
+                )
+            )
+
+            # Hidden layers with n - 2 layers
+            for i in range(self.num_layers - 2):
+                self.layers.append(
+                    CAREConv(
+                        self.hid_dim,
+                        self.hid_dim,
+                        self.num_classes,
+                        self.edges,
+                        activation=self.activation,
+                        step_size=self.step_size,
+                    )
+                )
+
+            # Output layer
+            self.layers.append(
+                CAREConv(
+                    self.hid_dim,
+                    self.num_classes,
+                    self.num_classes,
+                    self.edges,
+                    activation=self.activation,
+                    step_size=self.step_size,
+                )
+            )
+
+    def forward(self, graph, feat):
+        # For full graph training, directly use the graph
+        # formula 4
+        sim = th.tanh(self.layers[0].MLP(feat))
+
+        # Forward of n layers of CARE-GNN
+        for layer in self.layers:
+            feat = layer(graph, feat)
+
+        return feat, sim
+
+    def RLModule(self, graph, epoch, idx):
+        for layer in self.layers:
+            for etype in self.edges:
+                if not layer.cvg[etype]:
+                    # formula 5
+                    eid = graph.in_edges(idx, form="eid", etype=etype)
+                    avg_dist = th.mean(layer.dist[etype][eid])
+
+                    # formula 6
+                    if layer.last_avg_dist[etype] < avg_dist:
+                        if layer.p[etype] - self.step_size > 0:
+                            layer.p[etype] -= self.step_size
+                        layer.f[etype].append(-1)
+                    else:
+                        if layer.p[etype] + self.step_size <= 1:
+                            layer.p[etype] += self.step_size
+                        layer.f[etype].append(+1)
+                    layer.last_avg_dist[etype] = avg_dist
+
+                    # formula 7
+                    if epoch >= 9 and abs(sum(layer.f[etype][-10:])) <= 2:
+                        layer.cvg[etype] = True

examples/pytorch/caregnn/model_sampling.py

@@ -0,0 +1,230 @@
+import dgl
+import dgl.function as fn
+import numpy as np
+import torch as th
+import torch.nn as nn
+
+
+def _l1_dist(edges):
+    # formula 2
+    ed = th.norm(edges.src["nd"] - edges.dst["nd"], 1, 1)
+    return {"ed": ed}
+
+
+class CARESampler(dgl.dataloading.BlockSampler):
+    def __init__(self, p, dists, num_layers):
+        super().__init__()
+        self.p = p
+        self.dists = dists
+        self.num_layers = num_layers
+
+    def sample_frontier(self, block_id, g, seed_nodes, *args, **kwargs):
+        with g.local_scope():
+            new_edges_masks = {}
+            for etype in g.canonical_etypes:
+                edge_mask = th.zeros(g.num_edges(etype))
+                # extract each node from dict because of single node type
+                for node in seed_nodes:
+                    edges = g.in_edges(node, form="eid", etype=etype)
+                    num_neigh = (
+                        th.ceil(
+                            g.in_degrees(node, etype=etype)
+                            * self.p[block_id][etype]
+                        )
+                        .int()
+                        .item()
+                    )
+                    neigh_dist = self.dists[block_id][etype][edges]
+                    if neigh_dist.shape[0] > num_neigh:
+                        neigh_index = np.argpartition(neigh_dist, num_neigh)[
+                            :num_neigh
+                        ]
+                    else:
+                        neigh_index = np.arange(num_neigh)
+                    edge_mask[edges[neigh_index]] = 1
+                new_edges_masks[etype] = edge_mask.bool()
+
+            return dgl.edge_subgraph(g, new_edges_masks, relabel_nodes=False)
+
+    def sample_blocks(self, g, seed_nodes, exclude_eids=None):
+        output_nodes = seed_nodes
+        blocks = []
+        for block_id in reversed(range(self.num_layers)):
+            frontier = self.sample_frontier(block_id, g, seed_nodes)
+            eid = frontier.edata[dgl.EID]
+            block = dgl.to_block(frontier, seed_nodes)
+            block.edata[dgl.EID] = eid
+            seed_nodes = block.srcdata[dgl.NID]
+            blocks.insert(0, block)
+
+        return seed_nodes, output_nodes, blocks
+
+    def __len__(self):
+        return self.num_layers
+
+
+class CAREConv(nn.Module):
+    """One layer of CARE-GNN."""
+
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        num_classes,
+        edges,
+        activation=None,
+        step_size=0.02,
+    ):
+        super(CAREConv, self).__init__()
+
+        self.activation = activation
+        self.step_size = step_size
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.num_classes = num_classes
+        self.edges = edges
+
+        self.linear = nn.Linear(self.in_dim, self.out_dim)
+        self.MLP = nn.Linear(self.in_dim, self.num_classes)
+
+        self.p = {}
+        self.last_avg_dist = {}
+        self.f = {}
+        # indicate whether the RL converges
+        self.cvg = {}
+        for etype in edges:
+            self.p[etype] = 0.5
+            self.last_avg_dist[etype] = 0
+            self.f[etype] = []
+            self.cvg[etype] = False
+
+    def forward(self, g, feat):
+        g.srcdata["h"] = feat
+
+        # formula 8
+        hr = {}
+        for etype in g.canonical_etypes:
+            g.update_all(fn.copy_u("h", "m"), fn.mean("m", "hr"), etype=etype)
+            hr[etype] = g.dstdata["hr"]
+            if self.activation is not None:
+                hr[etype] = self.activation(hr[etype])
+
+        # formula 9 using mean as inter-relation aggregator
+        p_tensor = (
+            th.Tensor(list(self.p.values())).view(-1, 1, 1).to(feat.device)
+        )
+        h_homo = th.sum(th.stack(list(hr.values())) * p_tensor, dim=0)
+        h_homo += feat[: g.number_of_dst_nodes()]
+        if self.activation is not None:
+            h_homo = self.activation(h_homo)
+
+        return self.linear(h_homo)
+
+
+class CAREGNN(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        num_classes,
+        hid_dim=64,
+        edges=None,
+        num_layers=2,
+        activation=None,
+        step_size=0.02,
+    ):
+        super(CAREGNN, self).__init__()
+        self.in_dim = in_dim
+        self.hid_dim = hid_dim
+        self.num_classes = num_classes
+        self.edges = edges
+        self.num_layers = num_layers
+        self.activation = activation
+        self.step_size = step_size
+
+        self.layers = nn.ModuleList()
+
+        if self.num_layers == 1:
+            # Single layer
+            self.layers.append(
+                CAREConv(
+                    self.in_dim,
+                    self.num_classes,
+                    self.num_classes,
+                    self.edges,
+                    activation=self.activation,
+                    step_size=self.step_size,
+                )
+            )
+
+        else:
+            # Input layer
+            self.layers.append(
+                CAREConv(
+                    self.in_dim,
+                    self.hid_dim,
+                    self.num_classes,
+                    self.edges,
+                    activation=self.activation,
+                    step_size=self.step_size,
+                )
+            )
+
+            # Hidden layers with n - 2 layers
+            for i in range(self.num_layers - 2):
+                self.layers.append(
+                    CAREConv(
+                        self.hid_dim,
+                        self.hid_dim,
+                        self.num_classes,
+                        self.edges,
+                        activation=self.activation,
+                        step_size=self.step_size,
+                    )
+                )
+
+            # Output layer
+            self.layers.append(
+                CAREConv(
+                    self.hid_dim,
+                    self.num_classes,
+                    self.num_classes,
+                    self.edges,
+                    activation=self.activation,
+                    step_size=self.step_size,
+                )
+            )
+
+    def forward(self, blocks, feat):
+        # formula 4
+        sim = th.tanh(self.layers[0].MLP(blocks[-1].dstdata["feature"].float()))
+
+        # Forward of n layers of CARE-GNN
+        for block, layer in zip(blocks, self.layers):
+            feat = layer(block, feat)
+        return feat, sim
+
+    def RLModule(self, graph, epoch, idx, dists):
+        for i, layer in enumerate(self.layers):
+            for etype in self.edges:
+                if not layer.cvg[etype]:
+                    # formula 5
+                    eid = graph.in_edges(idx, form="eid", etype=etype)
+                    avg_dist = th.mean(dists[i][etype][eid])
+
+                    # formula 6
+                    if layer.last_avg_dist[etype] < avg_dist:
+                        layer.p[etype] -= self.step_size
+                        layer.f[etype].append(-1)
+                        # avoid overflow, follow the author's implement
+                        if layer.p[etype] < 0:
+                            layer.p[etype] = 0.001
+                    else:
+                        layer.p[etype] += self.step_size
+                        layer.f[etype].append(+1)
+                        if layer.p[etype] > 1:
+                            layer.p[etype] = 0.999
+                    layer.last_avg_dist[etype] = avg_dist
+
+                    # formula 7
+                    if epoch >= 9 and abs(sum(layer.f[etype][-10:])) <= 2:
+                        layer.cvg[etype] = True

examples/pytorch/caregnn/utils.py

@@ -0,0 +1,34 @@
+"""
+From GAT utils
+"""
+import torch
+
+
+class EarlyStopping:
+    def __init__(self, patience=10):
+        self.patience = patience
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+
+    def step(self, acc, model):
+        score = acc
+        if self.best_score is None:
+            self.best_score = score
+            self.save_checkpoint(model)
+        elif score < self.best_score:
+            self.counter += 1
+            print(
+                f"EarlyStopping counter: {self.counter} out of {self.patience}"
+            )
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.save_checkpoint(model)
+            self.counter = 0
+        return self.early_stop
+
+    def save_checkpoint(self, model):
+        """Saves model when validation loss decrease."""
+        torch.save(model.state_dict(), "es_checkpoint.pt")

examples/pytorch/cluster_gcn/README.md

@@ -0,0 +1,19 @@
+Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks
+============
+- Paper link: [Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks](https://arxiv.org/abs/1905.07953)
+- Author's code repo: [https://github.com/google-research/google-research/blob/master/cluster_gcn/](https://github.com/google-research/google-research/blob/master/cluster_gcn/).
+
+This repo reproduce the reported speed and performance maximally on Reddit and PPI. However, the diag enhancement is not covered, as the GraphSage aggregator already achieves satisfying F1 score.
+
+Dependencies
+------------
+- Python 3.7+(for string formatting features)
+- PyTorch 1.9.0+
+- scikit-learn
+- TorchMetrics 0.11.4
+
+## Run Experiments
+
+```bash
+python cluster_gcn.py
+```

examples/pytorch/cluster_gcn/cluster_gcn.py

@@ -0,0 +1,123 @@
+import time
+
+import dgl
+import dgl.nn as dglnn
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchmetrics.functional as MF
+from ogb.nodeproppred import DglNodePropPredDataset
+
+
+class SAGE(nn.Module):
+    def __init__(self, in_feats, n_hidden, n_classes):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, "mean"))
+        self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, "mean"))
+        self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, sg, x):
+        h = x
+        for l, layer in enumerate(self.layers):
+            h = layer(sg, h)
+            if l != len(self.layers) - 1:
+                h = F.relu(h)
+                h = self.dropout(h)
+        return h
+
+
+dataset = dgl.data.AsNodePredDataset(DglNodePropPredDataset("ogbn-products"))
+graph = dataset[
+    0
+]  # already prepares ndata['label'/'train_mask'/'val_mask'/'test_mask']
+
+model = SAGE(graph.ndata["feat"].shape[1], 256, dataset.num_classes).cuda()
+opt = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-4)
+
+num_partitions = 1000
+sampler = dgl.dataloading.ClusterGCNSampler(
+    graph,
+    num_partitions,
+    prefetch_ndata=["feat", "label", "train_mask", "val_mask", "test_mask"],
+)
+# DataLoader for generic dataloading with a graph, a set of indices (any indices, like
+# partition IDs here), and a graph sampler.
+dataloader = dgl.dataloading.DataLoader(
+    graph,
+    torch.arange(num_partitions).to("cuda"),
+    sampler,
+    device="cuda",
+    batch_size=100,
+    shuffle=True,
+    drop_last=False,
+    num_workers=0,
+    use_uva=True,
+)
+
+durations = []
+for epoch in range(10):
+    t0 = time.time()
+    model.train()
+    for it, sg in enumerate(dataloader):
+        x = sg.ndata["feat"]
+        y = sg.ndata["label"]
+        m = sg.ndata["train_mask"].bool()
+        y_hat = model(sg, x)
+        loss = F.cross_entropy(y_hat[m], y[m])
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+        if it % 20 == 0:
+            acc = MF.accuracy(
+                y_hat[m],
+                y[m],
+                task="multiclass",
+                num_classes=dataset.num_classes,
+            )
+            mem = torch.cuda.max_memory_allocated() / 1000000
+            print("Loss", loss.item(), "Acc", acc.item(), "GPU Mem", mem, "MB")
+
+    tt = time.time() - t0
+    print("Run time for epoch# %d: %.2fs" % (epoch, tt))
+    durations.append(tt)
+
+    model.eval()
+    with torch.no_grad():
+        val_preds, test_preds = [], []
+        val_labels, test_labels = [], []
+        for it, sg in enumerate(dataloader):
+            x = sg.ndata["feat"]
+            y = sg.ndata["label"]
+            m_val = sg.ndata["val_mask"].bool()
+            m_test = sg.ndata["test_mask"].bool()
+            y_hat = model(sg, x)
+            val_preds.append(y_hat[m_val])
+            val_labels.append(y[m_val])
+            test_preds.append(y_hat[m_test])
+            test_labels.append(y[m_test])
+        val_preds = torch.cat(val_preds, 0)
+        val_labels = torch.cat(val_labels, 0)
+        test_preds = torch.cat(test_preds, 0)
+        test_labels = torch.cat(test_labels, 0)
+        val_acc = MF.accuracy(
+            val_preds,
+            val_labels,
+            task="multiclass",
+            num_classes=dataset.num_classes,
+        )
+        test_acc = MF.accuracy(
+            test_preds,
+            test_labels,
+            task="multiclass",
+            num_classes=dataset.num_classes,
+        )
+        print("Validation acc:", val_acc.item(), "Test acc:", test_acc.item())
+
+print(
+    "Average run time for last %d epochs: %.2fs standard deviation: %.3f"
+    % ((epoch - 3), np.mean(durations[4:]), np.std(durations[4:]))
+)

examples/pytorch/compGCN/README.md

@@ -0,0 +1,78 @@
+# DGL Implementation of the CompGCN Paper
+
+This DGL example implements the GNN model proposed in the paper [CompositionGCN](https://arxiv.org/abs/1911.03082). 
+The author's codes of implementation is in [here](https://github.com/malllabiisc/CompGCN)
+
+Example implementor
+----------------------
+This example was implemented by [zhjwy9343](https://github.com/zhjwy9343) and [KounianhuaDu](https://github.com/KounianhuaDu) at the AWS Shanghai AI Lab.
+
+Dependencies
+----------------------
+- pytorch 1.9.0
+- dgl 0.7.1
+- numpy 1.20.3
+- ordered_set 4.0.2
+
+Dataset
+---------------------------------------
+The datasets used for link predictions are FB15k-237 constructed from Freebase and WN18RR constructed from WordNet. The statistics are summarized as followings:
+
+**FB15k-237** 
+
+- Nodes: 14541
+- Relation types: 237
+- Reversed relation types: 237
+- Train: 272115
+- Valid: 17535
+- Test: 20466
+
+**WN18RR** 
+
+- Nodes: 40943
+- Relation types: 11
+- Reversed relation types: 11
+- Train: 86835
+- Valid: 3034
+- Test: 3134
+
+How to run
+--------------------------------
+First to get the data, one can run 
+
+```python
+sh get_fb15k-237.sh
+```
+```python
+sh get_wn18rr.sh
+```
+
+Then for FB15k-237, run
+
+```python
+python main.py --score_func conve --opn ccorr --gpu 0 --data FB15k-237
+```
+
+For WN18RR, run
+
+```python
+python main.py --score_func conve --opn ccorr --gpu 0 --data wn18rr
+```
+
+
+Performance
+-------------------------
+**Link Prediction Results**
+
+| Dataset |        FB15k-237         |          WN18RR          |
+|---------| ------------------------ | ------------------------ |
+|  Metric |    Paper   /  ours (dgl) |    Paper   /  ours (dgl) |
+|   MRR   |    0.355   /    0.348    |    0.479   /    0.466    |
+|   MR    |     197    /     208     |    3533    /     3542    |
+| Hit@10  |    0.535   /   0.527     |    0.546   /    0.525    |
+|  Hit@3  |    0.390   /    0.380    |    0.494   /    0.476    |
+|  Hit@1  |    0.264   /    0.259    |    0.443   /    0.435    |
+
+
+
+

examples/pytorch/compGCN/data_loader.py

@@ -0,0 +1,283 @@
+from collections import defaultdict as ddict
+
+import dgl
+
+import numpy as np
+import torch
+from ordered_set import OrderedSet
+from torch.utils.data import DataLoader, Dataset
+
+
+class TrainDataset(Dataset):
+    """
+    Training Dataset class.
+    Parameters
+    ----------
+    triples: The triples used for training the model
+    num_ent: Number of entities in the knowledge graph
+    lbl_smooth: Label smoothing
+
+    Returns
+    -------
+    A training Dataset class instance used by DataLoader
+    """
+
+    def __init__(self, triples, num_ent, lbl_smooth):
+        self.triples = triples
+        self.num_ent = num_ent
+        self.lbl_smooth = lbl_smooth
+        self.entities = np.arange(self.num_ent, dtype=np.int32)
+
+    def __len__(self):
+        return len(self.triples)
+
+    def __getitem__(self, idx):
+        ele = self.triples[idx]
+        triple, label = torch.LongTensor(ele["triple"]), np.int32(ele["label"])
+        trp_label = self.get_label(label)
+        # label smoothing
+        if self.lbl_smooth != 0.0:
+            trp_label = (1.0 - self.lbl_smooth) * trp_label + (
+                1.0 / self.num_ent
+            )
+
+        return triple, trp_label
+
+    @staticmethod
+    def collate_fn(data):
+        triples = []
+        labels = []
+        for triple, label in data:
+            triples.append(triple)
+            labels.append(label)
+        triple = torch.stack(triples, dim=0)
+        trp_label = torch.stack(labels, dim=0)
+        return triple, trp_label
+
+    # for edges that exist in the graph, the entry is 1.0, otherwise the entry is 0.0
+    def get_label(self, label):
+        y = np.zeros([self.num_ent], dtype=np.float32)
+        for e2 in label:
+            y[e2] = 1.0
+        return torch.FloatTensor(y)
+
+
+class TestDataset(Dataset):
+    """
+    Evaluation Dataset class.
+    Parameters
+    ----------
+    triples: The triples used for evaluating the model
+    num_ent: Number of entities in the knowledge graph
+
+    Returns
+    -------
+    An evaluation Dataset class instance used by DataLoader for model evaluation
+    """
+
+    def __init__(self, triples, num_ent):
+        self.triples = triples
+        self.num_ent = num_ent
+
+    def __len__(self):
+        return len(self.triples)
+
+    def __getitem__(self, idx):
+        ele = self.triples[idx]
+        triple, label = torch.LongTensor(ele["triple"]), np.int32(ele["label"])
+        label = self.get_label(label)
+
+        return triple, label
+
+    @staticmethod
+    def collate_fn(data):
+        triples = []
+        labels = []
+        for triple, label in data:
+            triples.append(triple)
+            labels.append(label)
+        triple = torch.stack(triples, dim=0)
+        label = torch.stack(labels, dim=0)
+        return triple, label
+
+    # for edges that exist in the graph, the entry is 1.0, otherwise the entry is 0.0
+    def get_label(self, label):
+        y = np.zeros([self.num_ent], dtype=np.float32)
+        for e2 in label:
+            y[e2] = 1.0
+        return torch.FloatTensor(y)
+
+
+class Data(object):
+    def __init__(self, dataset, lbl_smooth, num_workers, batch_size):
+        """
+        Reading in raw triples and converts it into a standard format.
+        Parameters
+        ----------
+        dataset:           The name of the dataset
+        lbl_smooth:        Label smoothing
+        num_workers:       Number of workers of dataloaders
+        batch_size:        Batch size of dataloaders
+
+        Returns
+        -------
+        self.ent2id:            Entity to unique identifier mapping
+        self.rel2id:            Relation to unique identifier mapping
+        self.id2ent:            Inverse mapping of self.ent2id
+        self.id2rel:            Inverse mapping of self.rel2id
+        self.num_ent:           Number of entities in the knowledge graph
+        self.num_rel:           Number of relations in the knowledge graph
+
+        self.g:                 The dgl graph constucted from the edges in the traing set and all the entities in the knowledge graph
+        self.data['train']:     Stores the triples corresponding to training dataset
+        self.data['valid']:     Stores the triples corresponding to validation dataset
+        self.data['test']:      Stores the triples corresponding to test dataset
+        self.data_iter:		The dataloader for different data splits
+        """
+        self.dataset = dataset
+        self.lbl_smooth = lbl_smooth
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+
+        # read in raw data and get mappings
+        ent_set, rel_set = OrderedSet(), OrderedSet()
+        for split in ["train", "test", "valid"]:
+            for line in open("./{}/{}.txt".format(self.dataset, split)):
+                sub, rel, obj = map(str.lower, line.strip().split("\t"))
+                ent_set.add(sub)
+                rel_set.add(rel)
+                ent_set.add(obj)
+
+        self.ent2id = {ent: idx for idx, ent in enumerate(ent_set)}
+        self.rel2id = {rel: idx for idx, rel in enumerate(rel_set)}
+        self.rel2id.update(
+            {
+                rel + "_reverse": idx + len(self.rel2id)
+                for idx, rel in enumerate(rel_set)
+            }
+        )
+
+        self.id2ent = {idx: ent for ent, idx in self.ent2id.items()}
+        self.id2rel = {idx: rel for rel, idx in self.rel2id.items()}
+
+        self.num_ent = len(self.ent2id)
+        self.num_rel = len(self.rel2id) // 2
+
+        # read in ids of subjects, relations, and objects for train/test/valid
+        self.data = ddict(list)  # stores the triples
+        sr2o = ddict(
+            set
+        )  # The key of sr20 is (subject, relation), and the items are all the successors following (subject, relation)
+        src = []
+        dst = []
+        rels = []
+        inver_src = []
+        inver_dst = []
+        inver_rels = []
+
+        for split in ["train", "test", "valid"]:
+            for line in open("./{}/{}.txt".format(self.dataset, split)):
+                sub, rel, obj = map(str.lower, line.strip().split("\t"))
+                sub_id, rel_id, obj_id = (
+                    self.ent2id[sub],
+                    self.rel2id[rel],
+                    self.ent2id[obj],
+                )
+                self.data[split].append((sub_id, rel_id, obj_id))
+
+                if split == "train":
+                    sr2o[(sub_id, rel_id)].add(obj_id)
+                    sr2o[(obj_id, rel_id + self.num_rel)].add(
+                        sub_id
+                    )  # append the reversed edges
+                    src.append(sub_id)
+                    dst.append(obj_id)
+                    rels.append(rel_id)
+                    inver_src.append(obj_id)
+                    inver_dst.append(sub_id)
+                    inver_rels.append(rel_id + self.num_rel)
+
+        # construct dgl graph
+        src = src + inver_src
+        dst = dst + inver_dst
+        rels = rels + inver_rels
+        self.g = dgl.graph((src, dst), num_nodes=self.num_ent)
+        self.g.edata["etype"] = torch.Tensor(rels).long()
+
+        # identify in and out edges
+        in_edges_mask = [True] * (self.g.num_edges() // 2) + [False] * (
+            self.g.num_edges() // 2
+        )
+        out_edges_mask = [False] * (self.g.num_edges() // 2) + [True] * (
+            self.g.num_edges() // 2
+        )
+        self.g.edata["in_edges_mask"] = torch.Tensor(in_edges_mask)
+        self.g.edata["out_edges_mask"] = torch.Tensor(out_edges_mask)
+
+        # Prepare train/valid/test data
+        self.data = dict(self.data)
+        self.sr2o = {
+            k: list(v) for k, v in sr2o.items()
+        }  # store only the train data
+
+        for split in ["test", "valid"]:
+            for sub, rel, obj in self.data[split]:
+                sr2o[(sub, rel)].add(obj)
+                sr2o[(obj, rel + self.num_rel)].add(sub)
+
+        self.sr2o_all = {
+            k: list(v) for k, v in sr2o.items()
+        }  # store all the data
+        self.triples = ddict(list)
+
+        for (sub, rel), obj in self.sr2o.items():
+            self.triples["train"].append(
+                {"triple": (sub, rel, -1), "label": self.sr2o[(sub, rel)]}
+            )
+
+        for split in ["test", "valid"]:
+            for sub, rel, obj in self.data[split]:
+                rel_inv = rel + self.num_rel
+                self.triples["{}_{}".format(split, "tail")].append(
+                    {
+                        "triple": (sub, rel, obj),
+                        "label": self.sr2o_all[(sub, rel)],
+                    }
+                )
+                self.triples["{}_{}".format(split, "head")].append(
+                    {
+                        "triple": (obj, rel_inv, sub),
+                        "label": self.sr2o_all[(obj, rel_inv)],
+                    }
+                )
+
+        self.triples = dict(self.triples)
+
+        def get_train_data_loader(split, batch_size, shuffle=True):
+            return DataLoader(
+                TrainDataset(
+                    self.triples[split], self.num_ent, self.lbl_smooth
+                ),
+                batch_size=batch_size,
+                shuffle=shuffle,
+                num_workers=max(0, self.num_workers),
+                collate_fn=TrainDataset.collate_fn,
+            )
+
+        def get_test_data_loader(split, batch_size, shuffle=True):
+            return DataLoader(
+                TestDataset(self.triples[split], self.num_ent),
+                batch_size=batch_size,
+                shuffle=shuffle,
+                num_workers=max(0, self.num_workers),
+                collate_fn=TestDataset.collate_fn,
+            )
+
+        # train/valid/test dataloaders
+        self.data_iter = {
+            "train": get_train_data_loader("train", self.batch_size),
+            "valid_head": get_test_data_loader("valid_head", self.batch_size),
+            "valid_tail": get_test_data_loader("valid_tail", self.batch_size),
+            "test_head": get_test_data_loader("test_head", self.batch_size),
+            "test_tail": get_test_data_loader("test_tail", self.batch_size),
+        }

examples/pytorch/compGCN/get_fb15k-237.sh

@@ -0,0 +1,2 @@
+wget https://dgl-data.s3.cn-north-1.amazonaws.com.cn/dataset/FB15k-237.zip
+unzip FB15k-237.zip

examples/pytorch/compGCN/get_wn18rr.sh

@@ -0,0 +1,2 @@
+wget https://dgl-data.s3.cn-north-1.amazonaws.com.cn/dataset/wn18rr.zip
+unzip wn18rr.zip

examples/pytorch/compGCN/main.py

@@ -0,0 +1,361 @@
+import argparse
+from time import time
+
+import numpy as np
+import torch as th
+import torch.optim as optim
+from data_loader import Data
+from models import CompGCN_ConvE
+from utils import in_out_norm
+
+
+# predict the tail for (head, rel, -1) or head for (-1, rel, tail)
+def predict(model, graph, device, data_iter, split="valid", mode="tail"):
+    model.eval()
+    with th.no_grad():
+        results = {}
+        train_iter = iter(data_iter["{}_{}".format(split, mode)])
+
+        for step, batch in enumerate(train_iter):
+            triple, label = batch[0].to(device), batch[1].to(device)
+            sub, rel, obj, label = (
+                triple[:, 0],
+                triple[:, 1],
+                triple[:, 2],
+                label,
+            )
+            pred = model(graph, sub, rel)
+            b_range = th.arange(pred.size()[0], device=device)
+            target_pred = pred[b_range, obj]
+            pred = th.where(label.bool(), -th.ones_like(pred) * 10000000, pred)
+            pred[b_range, obj] = target_pred
+
+            # compute metrics
+            ranks = (
+                1
+                + th.argsort(
+                    th.argsort(pred, dim=1, descending=True),
+                    dim=1,
+                    descending=False,
+                )[b_range, obj]
+            )
+            ranks = ranks.float()
+            results["count"] = th.numel(ranks) + results.get("count", 0.0)
+            results["mr"] = th.sum(ranks).item() + results.get("mr", 0.0)
+            results["mrr"] = th.sum(1.0 / ranks).item() + results.get(
+                "mrr", 0.0
+            )
+            for k in [1, 3, 10]:
+                results["hits@{}".format(k)] = th.numel(
+                    ranks[ranks <= (k)]
+                ) + results.get("hits@{}".format(k), 0.0)
+
+    return results
+
+
+# evaluation function, evaluate the head and tail prediction and then combine the results
+def evaluate(model, graph, device, data_iter, split="valid"):
+    # predict for head and tail
+    left_results = predict(model, graph, device, data_iter, split, mode="tail")
+    right_results = predict(model, graph, device, data_iter, split, mode="head")
+    results = {}
+    count = float(left_results["count"])
+
+    # combine the head and tail prediction results
+    # Metrics: MRR, MR, and Hit@k
+    results["left_mr"] = round(left_results["mr"] / count, 5)
+    results["left_mrr"] = round(left_results["mrr"] / count, 5)
+    results["right_mr"] = round(right_results["mr"] / count, 5)
+    results["right_mrr"] = round(right_results["mrr"] / count, 5)
+    results["mr"] = round(
+        (left_results["mr"] + right_results["mr"]) / (2 * count), 5
+    )
+    results["mrr"] = round(
+        (left_results["mrr"] + right_results["mrr"]) / (2 * count), 5
+    )
+    for k in [1, 3, 10]:
+        results["left_hits@{}".format(k)] = round(
+            left_results["hits@{}".format(k)] / count, 5
+        )
+        results["right_hits@{}".format(k)] = round(
+            right_results["hits@{}".format(k)] / count, 5
+        )
+        results["hits@{}".format(k)] = round(
+            (
+                left_results["hits@{}".format(k)]
+                + right_results["hits@{}".format(k)]
+            )
+            / (2 * count),
+            5,
+        )
+    return results
+
+
+def main(args):
+    # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
+    # check cuda
+    if args.gpu >= 0 and th.cuda.is_available():
+        device = "cuda:{}".format(args.gpu)
+    else:
+        device = "cpu"
+
+    # construct graph, split in/out edges and prepare train/validation/test data_loader
+    data = Data(
+        args.dataset, args.lbl_smooth, args.num_workers, args.batch_size
+    )
+    data_iter = data.data_iter  # train/validation/test data_loader
+    graph = data.g.to(device)
+    num_rel = th.max(graph.edata["etype"]).item() + 1
+
+    # Compute in/out edge norms and store in edata
+    graph = in_out_norm(graph)
+
+    # Step 2: Create model =================================================================== #
+    compgcn_model = CompGCN_ConvE(
+        num_bases=args.num_bases,
+        num_rel=num_rel,
+        num_ent=graph.num_nodes(),
+        in_dim=args.init_dim,
+        layer_size=args.layer_size,
+        comp_fn=args.opn,
+        batchnorm=True,
+        dropout=args.dropout,
+        layer_dropout=args.layer_dropout,
+        num_filt=args.num_filt,
+        hid_drop=args.hid_drop,
+        feat_drop=args.feat_drop,
+        ker_sz=args.ker_sz,
+        k_w=args.k_w,
+        k_h=args.k_h,
+    )
+    compgcn_model = compgcn_model.to(device)
+
+    # Step 3: Create training components ===================================================== #
+    loss_fn = th.nn.BCELoss()
+    optimizer = optim.Adam(
+        compgcn_model.parameters(), lr=args.lr, weight_decay=args.l2
+    )
+
+    # Step 4: training epoches =============================================================== #
+    best_mrr = 0.0
+    kill_cnt = 0
+    for epoch in range(args.max_epochs):
+        # Training and validation using a full graph
+        compgcn_model.train()
+        train_loss = []
+        t0 = time()
+        for step, batch in enumerate(data_iter["train"]):
+            triple, label = batch[0].to(device), batch[1].to(device)
+            sub, rel, obj, label = (
+                triple[:, 0],
+                triple[:, 1],
+                triple[:, 2],
+                label,
+            )
+            logits = compgcn_model(graph, sub, rel)
+
+            # compute loss
+            tr_loss = loss_fn(logits, label)
+            train_loss.append(tr_loss.item())
+
+            # backward
+            optimizer.zero_grad()
+            tr_loss.backward()
+            optimizer.step()
+
+        train_loss = np.sum(train_loss)
+
+        t1 = time()
+        val_results = evaluate(
+            compgcn_model, graph, device, data_iter, split="valid"
+        )
+        t2 = time()
+
+        # validate
+        if val_results["mrr"] > best_mrr:
+            best_mrr = val_results["mrr"]
+            th.save(
+                compgcn_model.state_dict(), "comp_link" + "_" + args.dataset
+            )
+            kill_cnt = 0
+            print("saving model...")
+        else:
+            kill_cnt += 1
+            if kill_cnt > 100:
+                print("early stop.")
+                break
+        print(
+            "In epoch {}, Train Loss: {:.4f}, Valid MRR: {:.5}, Train time: {}, Valid time: {}".format(
+                epoch, train_loss, val_results["mrr"], t1 - t0, t2 - t1
+            )
+        )
+
+    # test use the best model
+    compgcn_model.eval()
+    compgcn_model.load_state_dict(th.load("comp_link" + "_" + args.dataset))
+    test_results = evaluate(
+        compgcn_model, graph, device, data_iter, split="test"
+    )
+    print(
+        "Test MRR: {:.5}\n, MR: {:.10}\n, H@10: {:.5}\n, H@3: {:.5}\n, H@1: {:.5}\n".format(
+            test_results["mrr"],
+            test_results["mr"],
+            test_results["hits@10"],
+            test_results["hits@3"],
+            test_results["hits@1"],
+        )
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Parser For Arguments",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--data",
+        dest="dataset",
+        default="FB15k-237",
+        help="Dataset to use, default: FB15k-237",
+    )
+    parser.add_argument(
+        "--model", dest="model", default="compgcn", help="Model Name"
+    )
+    parser.add_argument(
+        "--score_func",
+        dest="score_func",
+        default="conve",
+        help="Score Function for Link prediction",
+    )
+    parser.add_argument(
+        "--opn",
+        dest="opn",
+        default="ccorr",
+        help="Composition Operation to be used in CompGCN",
+    )
+
+    parser.add_argument(
+        "--batch", dest="batch_size", default=1024, type=int, help="Batch size"
+    )
+    parser.add_argument(
+        "--gpu",
+        type=int,
+        default="0",
+        help="Set GPU Ids : Eg: For CPU = -1, For Single GPU = 0",
+    )
+    parser.add_argument(
+        "--epoch",
+        dest="max_epochs",
+        type=int,
+        default=500,
+        help="Number of epochs",
+    )
+    parser.add_argument(
+        "--l2", type=float, default=0.0, help="L2 Regularization for Optimizer"
+    )
+    parser.add_argument(
+        "--lr", type=float, default=0.001, help="Starting Learning Rate"
+    )
+    parser.add_argument(
+        "--lbl_smooth",
+        dest="lbl_smooth",
+        type=float,
+        default=0.1,
+        help="Label Smoothing",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=10,
+        help="Number of processes to construct batches",
+    )
+    parser.add_argument(
+        "--seed",
+        dest="seed",
+        default=41504,
+        type=int,
+        help="Seed for randomization",
+    )
+
+    parser.add_argument(
+        "--num_bases",
+        dest="num_bases",
+        default=-1,
+        type=int,
+        help="Number of basis relation vectors to use",
+    )
+    parser.add_argument(
+        "--init_dim",
+        dest="init_dim",
+        default=100,
+        type=int,
+        help="Initial dimension size for entities and relations",
+    )
+    parser.add_argument(
+        "--layer_size",
+        nargs="?",
+        default="[200]",
+        help="List of output size for each compGCN layer",
+    )
+    parser.add_argument(
+        "--gcn_drop",
+        dest="dropout",
+        default=0.1,
+        type=float,
+        help="Dropout to use in GCN Layer",
+    )
+    parser.add_argument(
+        "--layer_dropout",
+        nargs="?",
+        default="[0.3]",
+        help="List of dropout value after each compGCN layer",
+    )
+
+    # ConvE specific hyperparameters
+    parser.add_argument(
+        "--hid_drop",
+        dest="hid_drop",
+        default=0.3,
+        type=float,
+        help="ConvE: Hidden dropout",
+    )
+    parser.add_argument(
+        "--feat_drop",
+        dest="feat_drop",
+        default=0.3,
+        type=float,
+        help="ConvE: Feature Dropout",
+    )
+    parser.add_argument(
+        "--k_w", dest="k_w", default=10, type=int, help="ConvE: k_w"
+    )
+    parser.add_argument(
+        "--k_h", dest="k_h", default=20, type=int, help="ConvE: k_h"
+    )
+    parser.add_argument(
+        "--num_filt",
+        dest="num_filt",
+        default=200,
+        type=int,
+        help="ConvE: Number of filters in convolution",
+    )
+    parser.add_argument(
+        "--ker_sz",
+        dest="ker_sz",
+        default=7,
+        type=int,
+        help="ConvE: Kernel size to use",
+    )
+
+    args = parser.parse_args()
+
+    np.random.seed(args.seed)
+    th.manual_seed(args.seed)
+
+    print(args)
+
+    args.layer_size = eval(args.layer_size)
+    args.layer_dropout = eval(args.layer_dropout)
+
+    main(args)

examples/pytorch/compGCN/models.py

@@ -0,0 +1,305 @@
+import dgl
+import dgl.function as fn
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from utils import ccorr
+
+
+class CompGraphConv(nn.Module):
+    """One layer of CompGCN."""
+
+    def __init__(
+        self, in_dim, out_dim, comp_fn="sub", batchnorm=True, dropout=0.1
+    ):
+        super(CompGraphConv, self).__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.comp_fn = comp_fn
+        self.actvation = th.tanh
+        self.batchnorm = batchnorm
+
+        # define dropout layer
+        self.dropout = nn.Dropout(dropout)
+
+        # define batch norm layer
+        if self.batchnorm:
+            self.bn = nn.BatchNorm1d(out_dim)
+
+        # define in/out/loop transform layer
+        self.W_O = nn.Linear(self.in_dim, self.out_dim)
+        self.W_I = nn.Linear(self.in_dim, self.out_dim)
+        self.W_S = nn.Linear(self.in_dim, self.out_dim)
+
+        # define relation transform layer
+        self.W_R = nn.Linear(self.in_dim, self.out_dim)
+
+        # self loop embedding
+        self.loop_rel = nn.Parameter(th.Tensor(1, self.in_dim))
+        nn.init.xavier_normal_(self.loop_rel)
+
+    def forward(self, g, n_in_feats, r_feats):
+        with g.local_scope():
+            # Assign values to source nodes. In a homogeneous graph, this is equal to
+            # assigning them to all nodes.
+            g.srcdata["h"] = n_in_feats
+            # append loop_rel embedding to r_feats
+            r_feats = th.cat((r_feats, self.loop_rel), 0)
+            # Assign features to all edges with the corresponding relation embeddings
+            g.edata["h"] = r_feats[g.edata["etype"]] * g.edata["norm"]
+
+            # Compute composition function in 4 steps
+            # Step 1: compute composition by edge in the edge direction, and store results in edges.
+            if self.comp_fn == "sub":
+                g.apply_edges(fn.u_sub_e("h", "h", out="comp_h"))
+            elif self.comp_fn == "mul":
+                g.apply_edges(fn.u_mul_e("h", "h", out="comp_h"))
+            elif self.comp_fn == "ccorr":
+                g.apply_edges(
+                    lambda edges: {
+                        "comp_h": ccorr(edges.src["h"], edges.data["h"])
+                    }
+                )
+            else:
+                raise Exception("Only supports sub, mul, and ccorr")
+
+            # Step 2: use extracted edge direction to compute in and out edges
+            comp_h = g.edata["comp_h"]
+
+            in_edges_idx = th.nonzero(
+                g.edata["in_edges_mask"], as_tuple=False
+            ).squeeze()
+            out_edges_idx = th.nonzero(
+                g.edata["out_edges_mask"], as_tuple=False
+            ).squeeze()
+
+            comp_h_O = self.W_O(comp_h[out_edges_idx])
+            comp_h_I = self.W_I(comp_h[in_edges_idx])
+
+            new_comp_h = th.zeros(comp_h.shape[0], self.out_dim).to(
+                comp_h.device
+            )
+            new_comp_h[out_edges_idx] = comp_h_O
+            new_comp_h[in_edges_idx] = comp_h_I
+
+            g.edata["new_comp_h"] = new_comp_h
+
+            # Step 3: sum comp results to both src and dst nodes
+            g.update_all(fn.copy_e("new_comp_h", "m"), fn.sum("m", "comp_edge"))
+
+            # Step 4: add results of self-loop
+            if self.comp_fn == "sub":
+                comp_h_s = n_in_feats - r_feats[-1]
+            elif self.comp_fn == "mul":
+                comp_h_s = n_in_feats * r_feats[-1]
+            elif self.comp_fn == "ccorr":
+                comp_h_s = ccorr(n_in_feats, r_feats[-1])
+            else:
+                raise Exception("Only supports sub, mul, and ccorr")
+
+            # Sum all of the comp results as output of nodes and dropout
+            n_out_feats = (
+                self.W_S(comp_h_s) + self.dropout(g.ndata["comp_edge"])
+            ) * (1 / 3)
+
+            # Compute relation output
+            r_out_feats = self.W_R(r_feats)
+
+            # Batch norm
+            if self.batchnorm:
+                n_out_feats = self.bn(n_out_feats)
+
+            # Activation function
+            if self.actvation is not None:
+                n_out_feats = self.actvation(n_out_feats)
+
+        return n_out_feats, r_out_feats[:-1]
+
+
+class CompGCN(nn.Module):
+    def __init__(
+        self,
+        num_bases,
+        num_rel,
+        num_ent,
+        in_dim=100,
+        layer_size=[200],
+        comp_fn="sub",
+        batchnorm=True,
+        dropout=0.1,
+        layer_dropout=[0.3],
+    ):
+        super(CompGCN, self).__init__()
+
+        self.num_bases = num_bases
+        self.num_rel = num_rel
+        self.num_ent = num_ent
+        self.in_dim = in_dim
+        self.layer_size = layer_size
+        self.comp_fn = comp_fn
+        self.batchnorm = batchnorm
+        self.dropout = dropout
+        self.layer_dropout = layer_dropout
+        self.num_layer = len(layer_size)
+
+        # CompGCN layers
+        self.layers = nn.ModuleList()
+        self.layers.append(
+            CompGraphConv(
+                self.in_dim,
+                self.layer_size[0],
+                comp_fn=self.comp_fn,
+                batchnorm=self.batchnorm,
+                dropout=self.dropout,
+            )
+        )
+        for i in range(self.num_layer - 1):
+            self.layers.append(
+                CompGraphConv(
+                    self.layer_size[i],
+                    self.layer_size[i + 1],
+                    comp_fn=self.comp_fn,
+                    batchnorm=self.batchnorm,
+                    dropout=self.dropout,
+                )
+            )
+
+        # Initial relation embeddings
+        if self.num_bases > 0:
+            self.basis = nn.Parameter(th.Tensor(self.num_bases, self.in_dim))
+            self.weights = nn.Parameter(th.Tensor(self.num_rel, self.num_bases))
+            nn.init.xavier_normal_(self.basis)
+            nn.init.xavier_normal_(self.weights)
+        else:
+            self.rel_embds = nn.Parameter(th.Tensor(self.num_rel, self.in_dim))
+            nn.init.xavier_normal_(self.rel_embds)
+
+        # Node embeddings
+        self.n_embds = nn.Parameter(th.Tensor(self.num_ent, self.in_dim))
+        nn.init.xavier_normal_(self.n_embds)
+
+        # Dropout after compGCN layers
+        self.dropouts = nn.ModuleList()
+        for i in range(self.num_layer):
+            self.dropouts.append(nn.Dropout(self.layer_dropout[i]))
+
+    def forward(self, graph):
+        # node and relation features
+        n_feats = self.n_embds
+        if self.num_bases > 0:
+            r_embds = th.mm(self.weights, self.basis)
+            r_feats = r_embds
+        else:
+            r_feats = self.rel_embds
+
+        for layer, dropout in zip(self.layers, self.dropouts):
+            n_feats, r_feats = layer(graph, n_feats, r_feats)
+            n_feats = dropout(n_feats)
+
+        return n_feats, r_feats
+
+
+# Use convE as the score function
+class CompGCN_ConvE(nn.Module):
+    def __init__(
+        self,
+        num_bases,
+        num_rel,
+        num_ent,
+        in_dim,
+        layer_size,
+        comp_fn="sub",
+        batchnorm=True,
+        dropout=0.1,
+        layer_dropout=[0.3],
+        num_filt=200,
+        hid_drop=0.3,
+        feat_drop=0.3,
+        ker_sz=5,
+        k_w=5,
+        k_h=5,
+    ):
+        super(CompGCN_ConvE, self).__init__()
+
+        self.embed_dim = layer_size[-1]
+        self.hid_drop = hid_drop
+        self.feat_drop = feat_drop
+        self.ker_sz = ker_sz
+        self.k_w = k_w
+        self.k_h = k_h
+        self.num_filt = num_filt
+
+        # compGCN model to get sub/rel embs
+        self.compGCN_Model = CompGCN(
+            num_bases,
+            num_rel,
+            num_ent,
+            in_dim,
+            layer_size,
+            comp_fn,
+            batchnorm,
+            dropout,
+            layer_dropout,
+        )
+
+        # batchnorms to the combined (sub+rel) emb
+        self.bn0 = th.nn.BatchNorm2d(1)
+        self.bn1 = th.nn.BatchNorm2d(self.num_filt)
+        self.bn2 = th.nn.BatchNorm1d(self.embed_dim)
+
+        # dropouts and conv module to the combined (sub+rel) emb
+        self.hidden_drop = th.nn.Dropout(self.hid_drop)
+        self.feature_drop = th.nn.Dropout(self.feat_drop)
+        self.m_conv1 = th.nn.Conv2d(
+            1,
+            out_channels=self.num_filt,
+            kernel_size=(self.ker_sz, self.ker_sz),
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+
+        flat_sz_h = int(2 * self.k_w) - self.ker_sz + 1
+        flat_sz_w = self.k_h - self.ker_sz + 1
+        self.flat_sz = flat_sz_h * flat_sz_w * self.num_filt
+        self.fc = th.nn.Linear(self.flat_sz, self.embed_dim)
+
+        # bias to the score
+        self.bias = nn.Parameter(th.zeros(num_ent))
+
+    # combine entity embeddings and relation embeddings
+    def concat(self, e1_embed, rel_embed):
+        e1_embed = e1_embed.view(-1, 1, self.embed_dim)
+        rel_embed = rel_embed.view(-1, 1, self.embed_dim)
+        stack_inp = th.cat([e1_embed, rel_embed], 1)
+        stack_inp = th.transpose(stack_inp, 2, 1).reshape(
+            (-1, 1, 2 * self.k_w, self.k_h)
+        )
+        return stack_inp
+
+    def forward(self, graph, sub, rel):
+        # get sub_emb and rel_emb via compGCN
+        n_feats, r_feats = self.compGCN_Model(graph)
+        sub_emb = n_feats[sub, :]
+        rel_emb = r_feats[rel, :]
+
+        # combine the sub_emb and rel_emb
+        stk_inp = self.concat(sub_emb, rel_emb)
+        # use convE to score the combined emb
+        x = self.bn0(stk_inp)
+        x = self.m_conv1(x)
+        x = self.bn1(x)
+        x = F.relu(x)
+        x = self.feature_drop(x)
+        x = x.view(-1, self.flat_sz)
+        x = self.fc(x)
+        x = self.hidden_drop(x)
+        x = self.bn2(x)
+        x = F.relu(x)
+        # compute score
+        x = th.mm(x, n_feats.transpose(1, 0))
+        # add in bias
+        x += self.bias.expand_as(x)
+        score = th.sigmoid(x)
+        return score

examples/pytorch/compGCN/utils.py

@@ -0,0 +1,66 @@
+# This file is based on the CompGCN author's implementation
+# <https://github.com/malllabiisc/CompGCN/blob/master/helper.py>.
+# It implements the operation of circular convolution in the ccorr function and an additional in_out_norm function for norm computation.
+
+import dgl
+import torch as th
+
+
+def com_mult(a, b):
+    r1, i1 = a[..., 0], a[..., 1]
+    r2, i2 = b[..., 0], b[..., 1]
+    return th.stack([r1 * r2 - i1 * i2, r1 * i2 + i1 * r2], dim=-1)
+
+
+def conj(a):
+    a[..., 1] = -a[..., 1]
+    return a
+
+
+def ccorr(a, b):
+    """
+    Compute circular correlation of two tensors.
+    Parameters
+    ----------
+    a: Tensor, 1D or 2D
+    b: Tensor, 1D or 2D
+
+    Notes
+    -----
+    Input a and b should have the same dimensions. And this operation supports broadcasting.
+
+    Returns
+    -------
+    Tensor, having the same dimension as the input a.
+    """
+    return th.fft.irfftn(
+        th.conj(th.fft.rfftn(a, (-1))) * th.fft.rfftn(b, (-1)), (-1)
+    )
+
+
+# identify in/out edges, compute edge norm for each and store in edata
+def in_out_norm(graph):
+    src, dst, EID = graph.edges(form="all")
+    graph.edata["norm"] = th.ones(EID.shape[0]).to(graph.device)
+
+    in_edges_idx = th.nonzero(
+        graph.edata["in_edges_mask"], as_tuple=False
+    ).squeeze()
+    out_edges_idx = th.nonzero(
+        graph.edata["out_edges_mask"], as_tuple=False
+    ).squeeze()
+
+    for idx in [in_edges_idx, out_edges_idx]:
+        u, v = src[idx], dst[idx]
+        deg = th.zeros(graph.num_nodes()).to(graph.device)
+        n_idx, inverse_index, count = th.unique(
+            v, return_inverse=True, return_counts=True
+        )
+        deg[n_idx] = count.float()
+        deg_inv = deg.pow(-0.5)  # D^{-0.5}
+        deg_inv[deg_inv == float("inf")] = 0
+        norm = deg_inv[u] * deg_inv[v]
+        graph.edata["norm"][idx] = norm
+    graph.edata["norm"] = graph.edata["norm"].unsqueeze(1)
+
+    return graph

examples/pytorch/correct_and_smooth/README.md

@@ -0,0 +1,79 @@
+# DGL Implementation of CorrectAndSmooth
+
+This DGL example implements the GNN model proposed in the paper [Combining Label Propagation and Simple Models Out-performs Graph Neural Networks](https://arxiv.org/abs/2010.13993). For the original implementation, see [here](https://github.com/CUAI/CorrectAndSmooth).
+
+Contributor: [xnuohz](https://github.com/xnuohz)
+
+### Requirements
+The codebase is implemented in Python 3.7. For version requirement of packages, see below.
+
+```
+dgl 0.6.0.post1
+torch 1.7.0
+ogb 1.3.0
+```
+
+### Limitations
+
+Spectral and Diffusion Embeddings used by the authors for feature augmentation are not currently implemented. Without these feature augmentations only the "Plain" (without feature augmentations) results from the authors can be replicated.
+
+### The graph datasets used in this example
+
+Open Graph Benchmark(OGB). Dataset summary:
+
+|    Dataset    |  #Nodes   |   #Edges   | #Node Feats |  Metric  |
+| :-----------: | :-------: | :--------: | :---------: | :------: |
+|  ogbn-arxiv   |  169,343  | 1,166,243  |     128     | Accuracy |
+| ogbn-products | 2,449,029 | 61,859,140 |     100     | Accuracy |
+
+### Usage
+
+Training a **Base predictor** and using **Correct&Smooth** which follows the original hyperparameters on different datasets.
+
+##### ogbn-arxiv
+
+* **Plain MLP + C&S**
+
+```bash
+python main.py --dropout 0.5
+python main.py --pretrain --correction-adj DA --smoothing-adj AD --autoscale
+```
+
+* **Plain Linear + C&S**
+
+```bash
+python main.py --model linear --dropout 0.5 --epochs 1000
+python main.py --model linear --pretrain --correction-alpha 0.87 --smoothing-alpha 0.81 --correction-adj AD --autoscale
+```
+
+##### ogbn-products
+
+* **Plain Linear + C&S**
+
+```bash
+python main.py --dataset ogbn-products --model linear --dropout 0.5 --epochs 1000 --lr 0.1
+python main.py --dataset ogbn-products --model linear --pretrain --correction-alpha 1. --smoothing-alpha 0.9
+```
+
+### Performance
+
+#### ogbn-arxiv
+
+|                 | Linear | Plain Linear + C&S |
+| :-------------: | :----: |    :----------:    |
+| Results(Author) | 52.5   |       71.26        |
+|  Results(DGL)   | 52.48  |       71.26        |
+
+#### ogbn-products
+
+|                 | Plain Linear | Plain Linear + C&S |
+| :-------------: | :----: | :----------: |
+| Results(Author) | 47.67  |    82.34     |
+|  Results(DGL)   | 47.65  |    82.86     |
+
+### Speed
+
+|      ogb-arxiv       |      Time     | GPU Memory | Params  |
+| :------------------: | :-----------: | :--------: | :-----: |
+| Author, Plain Linear + C&S | 6.3 * 10 ^ -3 |   1,248M   |  5,160  |
+|   DGL, Plain Linear + C&S  | 5.6 * 10 ^ -3 |   1,252M   |  5,160  |

examples/pytorch/correct_and_smooth/main.py

@@ -0,0 +1,198 @@
+import argparse
+import copy
+import os
+
+import dgl
+
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+from model import CorrectAndSmooth, MLP, MLPLinear
+from ogb.nodeproppred import DglNodePropPredDataset, Evaluator
+
+
+def evaluate(y_pred, y_true, idx, evaluator):
+    return evaluator.eval({"y_true": y_true[idx], "y_pred": y_pred[idx]})["acc"]
+
+
+def main():
+    # check cuda
+    device = (
+        f"cuda:{args.gpu}"
+        if torch.cuda.is_available() and args.gpu >= 0
+        else "cpu"
+    )
+    # load data
+    dataset = DglNodePropPredDataset(name=args.dataset)
+    evaluator = Evaluator(name=args.dataset)
+
+    split_idx = dataset.get_idx_split()
+    g, labels = dataset[
+        0
+    ]  # graph: DGLGraph object, label: torch tensor of shape (num_nodes, num_tasks)
+
+    if args.dataset == "ogbn-arxiv":
+        g = dgl.to_bidirected(g, copy_ndata=True)
+
+        feat = g.ndata["feat"]
+        feat = (feat - feat.mean(0)) / feat.std(0)
+        g.ndata["feat"] = feat
+
+    g = g.to(device)
+    feats = g.ndata["feat"]
+    labels = labels.to(device)
+
+    # load masks for train / validation / test
+    train_idx = split_idx["train"].to(device)
+    valid_idx = split_idx["valid"].to(device)
+    test_idx = split_idx["test"].to(device)
+
+    n_features = feats.size()[-1]
+    n_classes = dataset.num_classes
+
+    # load model
+    if args.model == "mlp":
+        model = MLP(
+            n_features, args.hid_dim, n_classes, args.num_layers, args.dropout
+        )
+    elif args.model == "linear":
+        model = MLPLinear(n_features, n_classes)
+    else:
+        raise NotImplementedError(f"Model {args.model} is not supported.")
+
+    model = model.to(device)
+    print(f"Model parameters: {sum(p.numel() for p in model.parameters())}")
+
+    if args.pretrain:
+        print("---------- Before ----------")
+        model.load_state_dict(
+            torch.load(
+                f"base/{args.dataset}-{args.model}.pt", weights_only=False
+            )
+        )
+        model.eval()
+
+        y_soft = model(feats).exp()
+
+        y_pred = y_soft.argmax(dim=-1, keepdim=True)
+        valid_acc = evaluate(y_pred, labels, valid_idx, evaluator)
+        test_acc = evaluate(y_pred, labels, test_idx, evaluator)
+        print(f"Valid acc: {valid_acc:.4f} | Test acc: {test_acc:.4f}")
+
+        print("---------- Correct & Smoothing ----------")
+        cs = CorrectAndSmooth(
+            num_correction_layers=args.num_correction_layers,
+            correction_alpha=args.correction_alpha,
+            correction_adj=args.correction_adj,
+            num_smoothing_layers=args.num_smoothing_layers,
+            smoothing_alpha=args.smoothing_alpha,
+            smoothing_adj=args.smoothing_adj,
+            autoscale=args.autoscale,
+            scale=args.scale,
+        )
+
+        y_soft = cs.correct(g, y_soft, labels[train_idx], train_idx)
+        y_soft = cs.smooth(g, y_soft, labels[train_idx], train_idx)
+        y_pred = y_soft.argmax(dim=-1, keepdim=True)
+        valid_acc = evaluate(y_pred, labels, valid_idx, evaluator)
+        test_acc = evaluate(y_pred, labels, test_idx, evaluator)
+        print(f"Valid acc: {valid_acc:.4f} | Test acc: {test_acc:.4f}")
+    else:
+        opt = optim.Adam(model.parameters(), lr=args.lr)
+
+        best_acc = 0
+        best_model = copy.deepcopy(model)
+
+        # training
+        print("---------- Training ----------")
+        for i in range(args.epochs):
+            model.train()
+            opt.zero_grad()
+
+            logits = model(feats)
+
+            train_loss = F.nll_loss(
+                logits[train_idx], labels.squeeze(1)[train_idx]
+            )
+            train_loss.backward()
+
+            opt.step()
+
+            model.eval()
+            with torch.no_grad():
+                logits = model(feats)
+
+                y_pred = logits.argmax(dim=-1, keepdim=True)
+
+                train_acc = evaluate(y_pred, labels, train_idx, evaluator)
+                valid_acc = evaluate(y_pred, labels, valid_idx, evaluator)
+
+                print(
+                    f"Epoch {i} | Train loss: {train_loss.item():.4f} | Train acc: {train_acc:.4f} | Valid acc {valid_acc:.4f}"
+                )
+
+                if valid_acc > best_acc:
+                    best_acc = valid_acc
+                    best_model = copy.deepcopy(model)
+
+        # testing & saving model
+        print("---------- Testing ----------")
+        best_model.eval()
+
+        logits = best_model(feats)
+
+        y_pred = logits.argmax(dim=-1, keepdim=True)
+        test_acc = evaluate(y_pred, labels, test_idx, evaluator)
+        print(f"Test acc: {test_acc:.4f}")
+
+        if not os.path.exists("base"):
+            os.makedirs("base")
+
+        torch.save(
+            best_model.state_dict(), f"base/{args.dataset}-{args.model}.pt"
+        )
+
+
+if __name__ == "__main__":
+    """
+    Correct & Smoothing Hyperparameters
+    """
+    parser = argparse.ArgumentParser(description="Base predictor(C&S)")
+
+    # Dataset
+    parser.add_argument("--gpu", type=int, default=0, help="-1 for cpu")
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="ogbn-arxiv",
+        choices=["ogbn-arxiv", "ogbn-products"],
+    )
+    # Base predictor
+    parser.add_argument(
+        "--model", type=str, default="mlp", choices=["mlp", "linear"]
+    )
+    parser.add_argument("--num-layers", type=int, default=3)
+    parser.add_argument("--hid-dim", type=int, default=256)
+    parser.add_argument("--dropout", type=float, default=0.4)
+    parser.add_argument("--lr", type=float, default=0.01)
+    parser.add_argument("--epochs", type=int, default=300)
+    # extra options for gat
+    parser.add_argument("--n-heads", type=int, default=3)
+    parser.add_argument("--attn_drop", type=float, default=0.05)
+    # C & S
+    parser.add_argument(
+        "--pretrain", action="store_true", help="Whether to perform C & S"
+    )
+    parser.add_argument("--num-correction-layers", type=int, default=50)
+    parser.add_argument("--correction-alpha", type=float, default=0.979)
+    parser.add_argument("--correction-adj", type=str, default="DAD")
+    parser.add_argument("--num-smoothing-layers", type=int, default=50)
+    parser.add_argument("--smoothing-alpha", type=float, default=0.756)
+    parser.add_argument("--smoothing-adj", type=str, default="DAD")
+    parser.add_argument("--autoscale", action="store_true")
+    parser.add_argument("--scale", type=float, default=20.0)
+
+    args = parser.parse_args()
+    print(args)
+
+    main()

examples/pytorch/correct_and_smooth/model.py

@@ -0,0 +1,230 @@
+import dgl.function as fn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MLPLinear(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super(MLPLinear, self).__init__()
+        self.linear = nn.Linear(in_dim, out_dim)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.linear.reset_parameters()
+
+    def forward(self, x):
+        return F.log_softmax(self.linear(x), dim=-1)
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim, hid_dim, out_dim, num_layers, dropout=0.0):
+        super(MLP, self).__init__()
+        assert num_layers >= 2
+
+        self.linears = nn.ModuleList()
+        self.bns = nn.ModuleList()
+        self.linears.append(nn.Linear(in_dim, hid_dim))
+        self.bns.append(nn.BatchNorm1d(hid_dim))
+
+        for _ in range(num_layers - 2):
+            self.linears.append(nn.Linear(hid_dim, hid_dim))
+            self.bns.append(nn.BatchNorm1d(hid_dim))
+
+        self.linears.append(nn.Linear(hid_dim, out_dim))
+        self.dropout = dropout
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for layer in self.linears:
+            layer.reset_parameters()
+        for layer in self.bns:
+            layer.reset_parameters()
+
+    def forward(self, x):
+        for linear, bn in zip(self.linears[:-1], self.bns):
+            x = linear(x)
+            x = F.relu(x, inplace=True)
+            x = bn(x)
+            x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.linears[-1](x)
+        return F.log_softmax(x, dim=-1)
+
+
+class LabelPropagation(nn.Module):
+    r"""
+
+    Description
+    -----------
+    Introduced in `Learning from Labeled and Unlabeled Data with Label Propagation <https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.14.3864&rep=rep1&type=pdf>`_
+
+    .. math::
+        \mathbf{Y}^{\prime} = \alpha \cdot \mathbf{D}^{-1/2} \mathbf{A}
+        \mathbf{D}^{-1/2} \mathbf{Y} + (1 - \alpha) \mathbf{Y},
+
+    where unlabeled data is inferred by labeled data via propagation.
+
+    Parameters
+    ----------
+        num_layers: int
+            The number of propagations.
+        alpha: float
+            The :math:`\alpha` coefficient.
+        adj: str
+            'DAD': D^-0.5 * A * D^-0.5
+            'DA': D^-1 * A
+            'AD': A * D^-1
+    """
+
+    def __init__(self, num_layers, alpha, adj="DAD"):
+        super(LabelPropagation, self).__init__()
+
+        self.num_layers = num_layers
+        self.alpha = alpha
+        self.adj = adj
+
+    @torch.no_grad()
+    def forward(
+        self, g, labels, mask=None, post_step=lambda y: y.clamp_(0.0, 1.0)
+    ):
+        with g.local_scope():
+            if labels.dtype == torch.long:
+                labels = F.one_hot(labels.view(-1)).to(torch.float32)
+
+            y = labels
+            if mask is not None:
+                y = torch.zeros_like(labels)
+                y[mask] = labels[mask]
+
+            last = (1 - self.alpha) * y
+            degs = g.in_degrees().float().clamp(min=1)
+            norm = (
+                torch.pow(degs, -0.5 if self.adj == "DAD" else -1)
+                .to(labels.device)
+                .unsqueeze(1)
+            )
+
+            for _ in range(self.num_layers):
+                # Assume the graphs to be undirected
+                if self.adj in ["DAD", "AD"]:
+                    y = norm * y
+
+                g.ndata["h"] = y
+                g.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
+                y = self.alpha * g.ndata.pop("h")
+
+                if self.adj in ["DAD", "DA"]:
+                    y = y * norm
+
+                y = post_step(last + y)
+
+            return y
+
+
+class CorrectAndSmooth(nn.Module):
+    r"""
+
+    Description
+    -----------
+    Introduced in `Combining Label Propagation and Simple Models Out-performs Graph Neural Networks <https://arxiv.org/abs/2010.13993>`_
+
+    Parameters
+    ----------
+        num_correction_layers: int
+            The number of correct propagations.
+        correction_alpha: float
+            The coefficient of correction.
+        correction_adj: str
+            'DAD': D^-0.5 * A * D^-0.5
+            'DA': D^-1 * A
+            'AD': A * D^-1
+        num_smoothing_layers: int
+            The number of smooth propagations.
+        smoothing_alpha: float
+            The coefficient of smoothing.
+        smoothing_adj: str
+            'DAD': D^-0.5 * A * D^-0.5
+            'DA': D^-1 * A
+            'AD': A * D^-1
+        autoscale: bool, optional
+            If set to True, will automatically determine the scaling factor :math:`\sigma`. Default is True.
+        scale: float, optional
+            The scaling factor :math:`\sigma`, in case :obj:`autoscale = False`. Default is 1.
+    """
+
+    def __init__(
+        self,
+        num_correction_layers,
+        correction_alpha,
+        correction_adj,
+        num_smoothing_layers,
+        smoothing_alpha,
+        smoothing_adj,
+        autoscale=True,
+        scale=1.0,
+    ):
+        super(CorrectAndSmooth, self).__init__()
+
+        self.autoscale = autoscale
+        self.scale = scale
+
+        self.prop1 = LabelPropagation(
+            num_correction_layers, correction_alpha, correction_adj
+        )
+        self.prop2 = LabelPropagation(
+            num_smoothing_layers, smoothing_alpha, smoothing_adj
+        )
+
+    def correct(self, g, y_soft, y_true, mask):
+        with g.local_scope():
+            assert abs(float(y_soft.sum()) / y_soft.size(0) - 1.0) < 1e-2
+            numel = (
+                int(mask.sum()) if mask.dtype == torch.bool else mask.size(0)
+            )
+            assert y_true.size(0) == numel
+
+            if y_true.dtype == torch.long:
+                y_true = F.one_hot(y_true.view(-1), y_soft.size(-1)).to(
+                    y_soft.dtype
+                )
+
+            error = torch.zeros_like(y_soft)
+            error[mask] = y_true - y_soft[mask]
+
+            if self.autoscale:
+                smoothed_error = self.prop1(
+                    g, error, post_step=lambda x: x.clamp_(-1.0, 1.0)
+                )
+                sigma = error[mask].abs().sum() / numel
+                scale = sigma / smoothed_error.abs().sum(dim=1, keepdim=True)
+                scale[scale.isinf() | (scale > 1000)] = 1.0
+
+                result = y_soft + scale * smoothed_error
+                result[result.isnan()] = y_soft[result.isnan()]
+                return result
+            else:
+
+                def fix_input(x):
+                    x[mask] = error[mask]
+                    return x
+
+                smoothed_error = self.prop1(g, error, post_step=fix_input)
+
+                result = y_soft + self.scale * smoothed_error
+                result[result.isnan()] = y_soft[result.isnan()]
+                return result
+
+    def smooth(self, g, y_soft, y_true, mask):
+        with g.local_scope():
+            numel = (
+                int(mask.sum()) if mask.dtype == torch.bool else mask.size(0)
+            )
+            assert y_true.size(0) == numel
+
+            if y_true.dtype == torch.long:
+                y_true = F.one_hot(y_true.view(-1), y_soft.size(-1)).to(
+                    y_soft.dtype
+                )
+
+            y_soft[mask] = y_true
+            return self.prop2(g, y_soft)

examples/pytorch/dagnn/README.md

@@ -0,0 +1,69 @@
+# DAGNN
+
+This DGL example implements the GNN model proposed in the paper [Towards Deeper Graph Neural Networks](https://arxiv.org/abs/2007.09296).
+
+Paper link: https://arxiv.org/abs/2007.09296
+
+Author's code: https://github.com/divelab/DeeperGNN
+
+Contributor: Liu Tang ([@lt610](https://github.com/lt610))
+
+## Dependecies
+- Python 3.6.10
+- PyTorch 1.4.0
+- numpy 1.18.1
+- dgl 0.5.3
+- tqdm 4.44.1
+
+## Dataset
+
+The DGL's built-in Cora, Pubmed and Citeseer datasets. Dataset summary:
+
+| Dataset | #Nodes | #Edges | #Feats | #Classes | #Train Nodes | #Val Nodes | #Test Nodes |
+| :-: | :-: | :-: | :-: | :-: | :-: | :-: | :-: |
+| Citeseer | 3,327 | 9,228 | 3,703 | 6 | 120 | 500 | 1000 |
+| Cora | 2,708 | 10,556 | 1,433 | 7 | 140 | 500 | 1000 |
+| Pubmed | 19,717 | 88,651 | 500 | 3 | 60 | 500 | 1000 |
+
+## Arguments
+
+###### Dataset options
+```
+--dataset          str     The graph dataset name.             Default is 'Cora'.
+```
+
+###### GPU options
+```
+--gpu              int     GPU index.                          Default is -1, using CPU.
+```
+
+###### Model options
+```
+--runs             int     Number of training runs.               Default is 1
+--epochs           int     Number of training epochs.             Default is 1500.
+--early-stopping   int     Early stopping patience rounds.        Default is 100.
+--lr               float   Adam optimizer learning rate.          Default is 0.01.
+--lamb             float   L2 regularization coefficient.         Default is 5e-3.
+--k                int     Number of propagation layers.          Default is 10.
+--hid-dim          int     Hidden layer dimensionalities.         Default is 64.
+--dropout          float   Dropout rate                           Default is 0.8
+```
+
+## Examples
+
+Train a model which follows the original hyperparameters on different datasets.
+```bash
+# Cora:
+python main.py --dataset Cora --gpu 0 --runs 100 --lamb 0.005 --k 12
+# Citeseer:
+python main.py --dataset Citeseer --gpu 0 --runs 100 --lamb 0.02 --k 16
+# Pubmed:
+python main.py --dataset Pubmed --gpu 0 --runs 100 --lamb 0.005 --k 20
+```
+### Performance
+
+#### On Cora, Citeseer and Pubmed
+| Dataset | Cora | Citeseer | Pubmed |
+| :-: | :-: | :-: | :-: |
+| Accuracy Reported(100 runs) | 84.4 ± 0.5 | 73.3 ± 0.6 | 80.5 ± 0.5 |
+| Accuracy DGL(100 runs) | 84.3 ± 0.5 | 73.1 ± 0.9 | 80.5 ± 0.4 |

examples/pytorch/dagnn/main.py

@@ -0,0 +1,281 @@
+import argparse
+
+import dgl.function as fn
+
+import numpy as np
+import torch
+from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
+from torch import nn
+from torch.nn import functional as F, Parameter
+from tqdm import trange
+from utils import evaluate, generate_random_seeds, set_random_state
+
+
+class DAGNNConv(nn.Module):
+    def __init__(self, in_dim, k):
+        super(DAGNNConv, self).__init__()
+
+        self.s = Parameter(torch.FloatTensor(in_dim, 1))
+        self.k = k
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        gain = nn.init.calculate_gain("sigmoid")
+        nn.init.xavier_uniform_(self.s, gain=gain)
+
+    def forward(self, graph, feats):
+        with graph.local_scope():
+            results = [feats]
+
+            degs = graph.in_degrees().float()
+            norm = torch.pow(degs, -0.5)
+            norm = norm.to(feats.device).unsqueeze(1)
+
+            for _ in range(self.k):
+                feats = feats * norm
+                graph.ndata["h"] = feats
+                graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
+                feats = graph.ndata["h"]
+                feats = feats * norm
+                results.append(feats)
+
+            H = torch.stack(results, dim=1)
+            S = F.sigmoid(torch.matmul(H, self.s))
+            S = S.permute(0, 2, 1)
+            H = torch.matmul(S, H).squeeze()
+
+            return H
+
+
+class MLPLayer(nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, activation=None, dropout=0):
+        super(MLPLayer, self).__init__()
+
+        self.linear = nn.Linear(in_dim, out_dim, bias=bias)
+        self.activation = activation
+        self.dropout = nn.Dropout(dropout)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        gain = 1.0
+        if self.activation is F.relu:
+            gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_uniform_(self.linear.weight, gain=gain)
+        if self.linear.bias is not None:
+            nn.init.zeros_(self.linear.bias)
+
+    def forward(self, feats):
+        feats = self.dropout(feats)
+        feats = self.linear(feats)
+        if self.activation:
+            feats = self.activation(feats)
+
+        return feats
+
+
+class DAGNN(nn.Module):
+    def __init__(
+        self,
+        k,
+        in_dim,
+        hid_dim,
+        out_dim,
+        bias=True,
+        activation=F.relu,
+        dropout=0,
+    ):
+        super(DAGNN, self).__init__()
+        self.mlp = nn.ModuleList()
+        self.mlp.append(
+            MLPLayer(
+                in_dim=in_dim,
+                out_dim=hid_dim,
+                bias=bias,
+                activation=activation,
+                dropout=dropout,
+            )
+        )
+        self.mlp.append(
+            MLPLayer(
+                in_dim=hid_dim,
+                out_dim=out_dim,
+                bias=bias,
+                activation=None,
+                dropout=dropout,
+            )
+        )
+        self.dagnn = DAGNNConv(in_dim=out_dim, k=k)
+
+    def forward(self, graph, feats):
+        for layer in self.mlp:
+            feats = layer(feats)
+        feats = self.dagnn(graph, feats)
+        return feats
+
+
+def main(args):
+    # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
+    # Load from DGL dataset
+    if args.dataset == "Cora":
+        dataset = CoraGraphDataset()
+    elif args.dataset == "Citeseer":
+        dataset = CiteseerGraphDataset()
+    elif args.dataset == "Pubmed":
+        dataset = PubmedGraphDataset()
+    else:
+        raise ValueError("Dataset {} is invalid.".format(args.dataset))
+
+    graph = dataset[0]
+    graph = graph.add_self_loop()
+
+    # check cuda
+    if args.gpu >= 0 and torch.cuda.is_available():
+        device = "cuda:{}".format(args.gpu)
+    else:
+        device = "cpu"
+
+    # retrieve the number of classes
+    n_classes = dataset.num_classes
+
+    # retrieve labels of ground truth
+    labels = graph.ndata.pop("label").to(device).long()
+
+    # Extract node features
+    feats = graph.ndata.pop("feat").to(device)
+    n_features = feats.shape[-1]
+
+    # retrieve masks for train/validation/test
+    train_mask = graph.ndata.pop("train_mask")
+    val_mask = graph.ndata.pop("val_mask")
+    test_mask = graph.ndata.pop("test_mask")
+
+    train_idx = torch.nonzero(train_mask, as_tuple=False).squeeze().to(device)
+    val_idx = torch.nonzero(val_mask, as_tuple=False).squeeze().to(device)
+    test_idx = torch.nonzero(test_mask, as_tuple=False).squeeze().to(device)
+
+    graph = graph.to(device)
+
+    # Step 2: Create model =================================================================== #
+    model = DAGNN(
+        k=args.k,
+        in_dim=n_features,
+        hid_dim=args.hid_dim,
+        out_dim=n_classes,
+        dropout=args.dropout,
+    )
+    model = model.to(device)
+
+    # Step 3: Create training components ===================================================== #
+    loss_fn = F.cross_entropy
+    opt = torch.optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.lamb
+    )
+
+    # Step 4: training epochs =============================================================== #
+    loss = float("inf")
+    best_acc = 0
+    no_improvement = 0
+    epochs = trange(args.epochs, desc="Accuracy & Loss")
+
+    for _ in epochs:
+        model.train()
+
+        logits = model(graph, feats)
+
+        # compute loss
+        train_loss = loss_fn(logits[train_idx], labels[train_idx])
+
+        # backward
+        opt.zero_grad()
+        train_loss.backward()
+        opt.step()
+
+        (
+            train_loss,
+            train_acc,
+            valid_loss,
+            valid_acc,
+            test_loss,
+            test_acc,
+        ) = evaluate(
+            model, graph, feats, labels, (train_idx, val_idx, test_idx)
+        )
+
+        # Print out performance
+        epochs.set_description(
+            "Train Acc {:.4f} | Train Loss {:.4f} | Val Acc {:.4f} | Val loss {:.4f}".format(
+                train_acc, train_loss.item(), valid_acc, valid_loss.item()
+            )
+        )
+
+        if valid_loss > loss:
+            no_improvement += 1
+            if no_improvement == args.early_stopping:
+                print("Early stop.")
+                break
+        else:
+            no_improvement = 0
+            loss = valid_loss
+            best_acc = test_acc
+
+    print("Test Acc {:.4f}".format(best_acc))
+    return best_acc
+
+
+if __name__ == "__main__":
+    """
+    DAGNN Model Hyperparameters
+    """
+    parser = argparse.ArgumentParser(description="DAGNN")
+    # data source params
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="Cora",
+        choices=["Cora", "Citeseer", "Pubmed"],
+        help="Name of dataset.",
+    )
+    # cuda params
+    parser.add_argument(
+        "--gpu", type=int, default=-1, help="GPU index. Default: -1, using CPU."
+    )
+    # training params
+    parser.add_argument("--runs", type=int, default=1, help="Training runs.")
+    parser.add_argument(
+        "--epochs", type=int, default=1500, help="Training epochs."
+    )
+    parser.add_argument(
+        "--early-stopping",
+        type=int,
+        default=100,
+        help="Patient epochs to wait before early stopping.",
+    )
+    parser.add_argument("--lr", type=float, default=0.01, help="Learning rate.")
+    parser.add_argument("--lamb", type=float, default=0.005, help="L2 reg.")
+    # model params
+    parser.add_argument(
+        "--k", type=int, default=12, help="Number of propagation layers."
+    )
+    parser.add_argument(
+        "--hid-dim", type=int, default=64, help="Hidden layer dimensionalities."
+    )
+    parser.add_argument("--dropout", type=float, default=0.8, help="dropout")
+    args = parser.parse_args()
+    print(args)
+
+    acc_lists = []
+    random_seeds = generate_random_seeds(seed=1222, nums=args.runs)
+
+    for run in range(args.runs):
+        set_random_state(random_seeds[run])
+        acc_lists.append(main(args))
+
+    acc_lists = np.array(acc_lists)
+
+    mean = np.around(np.mean(acc_lists, axis=0), decimals=4)
+    std = np.around(np.std(acc_lists, axis=0), decimals=4)
+
+    print("Total acc: ", acc_lists)
+    print("mean", mean)
+    print("std", std)

examples/pytorch/dagnn/utils.py

@@ -0,0 +1,33 @@
+import random
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+
+def evaluate(model, graph, feats, labels, idxs):
+    model.eval()
+    with torch.no_grad():
+        logits = model(graph, feats)
+        results = ()
+        for idx in idxs:
+            loss = F.cross_entropy(logits[idx], labels[idx])
+            acc = torch.sum(
+                logits[idx].argmax(dim=1) == labels[idx]
+            ).item() / len(idx)
+            results += (loss, acc)
+    return results
+
+
+def generate_random_seeds(seed, nums):
+    random.seed(seed)
+    return [random.randint(1, 999999999) for _ in range(nums)]
+
+
+def set_random_state(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True

examples/pytorch/deepergcn/README.md

@@ -0,0 +1,51 @@
+# DGL Implementation of DeeperGCN
+
+This DGL example implements the GNN model proposed in the paper [DeeperGCN: All You Need to Train Deeper GCNs](https://arxiv.org/abs/2006.07739). For the original implementation, see [here](https://github.com/lightaime/deep_gcns_torch).
+
+Contributor: [xnuohz](https://github.com/xnuohz)
+
+### Requirements
+The codebase is implemented in Python 3.7. For version requirement of packages, see below.
+
+```
+dgl 0.6.0.post1
+torch 1.7.0
+ogb 1.3.0
+```
+
+### The graph datasets used in this example
+
+Open Graph Benchmark(OGB). Dataset summary:
+
+###### Graph Property Prediction
+
+|   Dataset   | #Graphs | #Node Feats | #Edge Feats | Metric  |
+| :---------: | :-----: | :---------: | :---------: | :-----: |
+| ogbg-molhiv | 41,127  |      9      |      3      | ROC-AUC |
+
+### Usage
+
+Train a model which follows the original hyperparameters on different datasets.
+```bash
+# ogbg-molhiv
+python main.py --gpu 0 --learn-beta
+```
+
+### Performance
+
+* Table 6: Numbers associated with "Table 6" are the ones from table 6 in the paper.
+* Author: Numbers associated with "Author" are the ones we got by running the original code.
+* DGL: Numbers associated with "DGL" are the ones we got by running the DGL example.
+
+|     Dataset      | ogbg-molhiv |
+| :--------------: | :---------: |
+| Results(Table 6) |    0.786    |
+| Results(Author)  |    0.781    |
+|   Results(DGL)   |    0.778    |
+
+### Speed
+
+|     Dataset     | ogbg-molhiv |
+| :-------------: | :---------: |
+| Results(Author) |   11.833    |
+|  Results(DGL)   |    8.965    |

examples/pytorch/deepergcn/layers.py

@@ -0,0 +1,118 @@
+import dgl.function as fn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.nn.functional import edge_softmax
+from modules import MessageNorm, MLP
+from ogb.graphproppred.mol_encoder import BondEncoder
+
+
+class GENConv(nn.Module):
+    r"""
+
+    Description
+    -----------
+    Generalized Message Aggregator was introduced in "DeeperGCN: All You Need to Train Deeper GCNs <https://arxiv.org/abs/2006.07739>"
+
+    Parameters
+    ----------
+    in_dim: int
+        Input size.
+    out_dim: int
+        Output size.
+    aggregator: str
+        Type of aggregation. Default is 'softmax'.
+    beta: float
+        A continuous variable called an inverse temperature. Default is 1.0.
+    learn_beta: bool
+        Whether beta is a learnable variable or not. Default is False.
+    p: float
+        Initial power for power mean aggregation. Default is 1.0.
+    learn_p: bool
+        Whether p is a learnable variable or not. Default is False.
+    msg_norm: bool
+        Whether message normalization is used. Default is False.
+    learn_msg_scale: bool
+        Whether s is a learnable scaling factor or not in message normalization. Default is False.
+    mlp_layers: int
+        The number of MLP layers. Default is 1.
+    eps: float
+        A small positive constant in message construction function. Default is 1e-7.
+    """
+
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        aggregator="softmax",
+        beta=1.0,
+        learn_beta=False,
+        p=1.0,
+        learn_p=False,
+        msg_norm=False,
+        learn_msg_scale=False,
+        mlp_layers=1,
+        eps=1e-7,
+    ):
+        super(GENConv, self).__init__()
+
+        self.aggr = aggregator
+        self.eps = eps
+
+        channels = [in_dim]
+        for _ in range(mlp_layers - 1):
+            channels.append(in_dim * 2)
+        channels.append(out_dim)
+
+        self.mlp = MLP(channels)
+        self.msg_norm = MessageNorm(learn_msg_scale) if msg_norm else None
+
+        self.beta = (
+            nn.Parameter(torch.Tensor([beta]), requires_grad=True)
+            if learn_beta and self.aggr == "softmax"
+            else beta
+        )
+        self.p = (
+            nn.Parameter(torch.Tensor([p]), requires_grad=True)
+            if learn_p
+            else p
+        )
+
+        self.edge_encoder = BondEncoder(in_dim)
+
+    def forward(self, g, node_feats, edge_feats):
+        with g.local_scope():
+            # Node and edge feature size need to match.
+            g.ndata["h"] = node_feats
+            g.edata["h"] = self.edge_encoder(edge_feats)
+            g.apply_edges(fn.u_add_e("h", "h", "m"))
+
+            if self.aggr == "softmax":
+                g.edata["m"] = F.relu(g.edata["m"]) + self.eps
+                g.edata["a"] = edge_softmax(g, g.edata["m"] * self.beta)
+                g.update_all(
+                    lambda edge: {"x": edge.data["m"] * edge.data["a"]},
+                    fn.sum("x", "m"),
+                )
+
+            elif self.aggr == "power":
+                minv, maxv = 1e-7, 1e1
+                torch.clamp_(g.edata["m"], minv, maxv)
+                g.update_all(
+                    lambda edge: {"x": torch.pow(edge.data["m"], self.p)},
+                    fn.mean("x", "m"),
+                )
+                torch.clamp_(g.ndata["m"], minv, maxv)
+                g.ndata["m"] = torch.pow(g.ndata["m"], self.p)
+
+            else:
+                raise NotImplementedError(
+                    f"Aggregator {self.aggr} is not supported."
+                )
+
+            if self.msg_norm is not None:
+                g.ndata["m"] = self.msg_norm(node_feats, g.ndata["m"])
+
+            feats = node_feats + g.ndata["m"]
+
+            return self.mlp(feats)

examples/pytorch/deepergcn/main.py

@@ -0,0 +1,165 @@
+import argparse
+import copy
+import time
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from models import DeeperGCN
+from ogb.graphproppred import collate_dgl, DglGraphPropPredDataset, Evaluator
+from torch.utils.data import DataLoader
+
+
+def train(model, device, data_loader, opt, loss_fn):
+    model.train()
+
+    train_loss = []
+    for g, labels in data_loader:
+        g = g.to(device)
+        labels = labels.to(torch.float32).to(device)
+        logits = model(g, g.edata["feat"], g.ndata["feat"])
+        loss = loss_fn(logits, labels)
+        train_loss.append(loss.item())
+
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+
+    return sum(train_loss) / len(train_loss)
+
+
+@torch.no_grad()
+def test(model, device, data_loader, evaluator):
+    model.eval()
+    y_true, y_pred = [], []
+
+    for g, labels in data_loader:
+        g = g.to(device)
+        logits = model(g, g.edata["feat"], g.ndata["feat"])
+        y_true.append(labels.detach().cpu())
+        y_pred.append(logits.detach().cpu())
+
+    y_true = torch.cat(y_true, dim=0).numpy()
+    y_pred = torch.cat(y_pred, dim=0).numpy()
+
+    return evaluator.eval({"y_true": y_true, "y_pred": y_pred})["rocauc"]
+
+
+def main():
+    # check cuda
+    device = (
+        f"cuda:{args.gpu}"
+        if args.gpu >= 0 and torch.cuda.is_available()
+        else "cpu"
+    )
+
+    # load ogb dataset & evaluator
+    dataset = DglGraphPropPredDataset(name="ogbg-molhiv")
+    evaluator = Evaluator(name="ogbg-molhiv")
+
+    g, _ = dataset[0]
+    node_feat_dim = g.ndata["feat"].size()[-1]
+    edge_feat_dim = g.edata["feat"].size()[-1]
+    n_classes = dataset.num_tasks
+
+    split_idx = dataset.get_idx_split()
+    train_loader = DataLoader(
+        dataset[split_idx["train"]],
+        batch_size=args.batch_size,
+        shuffle=True,
+        collate_fn=collate_dgl,
+    )
+    valid_loader = DataLoader(
+        dataset[split_idx["valid"]],
+        batch_size=args.batch_size,
+        shuffle=False,
+        collate_fn=collate_dgl,
+    )
+    test_loader = DataLoader(
+        dataset[split_idx["test"]],
+        batch_size=args.batch_size,
+        shuffle=False,
+        collate_fn=collate_dgl,
+    )
+
+    # load model
+    model = DeeperGCN(
+        node_feat_dim=node_feat_dim,
+        edge_feat_dim=edge_feat_dim,
+        hid_dim=args.hid_dim,
+        out_dim=n_classes,
+        num_layers=args.num_layers,
+        dropout=args.dropout,
+        learn_beta=args.learn_beta,
+    ).to(device)
+
+    print(model)
+
+    opt = optim.Adam(model.parameters(), lr=args.lr)
+    loss_fn = nn.BCEWithLogitsLoss()
+
+    # training & validation & testing
+    best_auc = 0
+    best_model = copy.deepcopy(model)
+    times = []
+
+    print("---------- Training ----------")
+    for i in range(args.epochs):
+        t1 = time.time()
+        train_loss = train(model, device, train_loader, opt, loss_fn)
+        t2 = time.time()
+
+        if i >= 5:
+            times.append(t2 - t1)
+
+        train_auc = test(model, device, train_loader, evaluator)
+        valid_auc = test(model, device, valid_loader, evaluator)
+
+        print(
+            f"Epoch {i} | Train Loss: {train_loss:.4f} | Train Auc: {train_auc:.4f} | Valid Auc: {valid_auc:.4f}"
+        )
+
+        if valid_auc > best_auc:
+            best_auc = valid_auc
+            best_model = copy.deepcopy(model)
+
+    print("---------- Testing ----------")
+    test_auc = test(best_model, device, test_loader, evaluator)
+    print(f"Test Auc: {test_auc}")
+    if len(times) > 0:
+        print("Times/epoch: ", sum(times) / len(times))
+
+
+if __name__ == "__main__":
+    """
+    DeeperGCN Hyperparameters
+    """
+    parser = argparse.ArgumentParser(description="DeeperGCN")
+    # training
+    parser.add_argument(
+        "--gpu", type=int, default=-1, help="GPU index, -1 for CPU."
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=300, help="Number of epochs to train."
+    )
+    parser.add_argument("--lr", type=float, default=0.01, help="Learning rate.")
+    parser.add_argument(
+        "--dropout", type=float, default=0.2, help="Dropout rate."
+    )
+    parser.add_argument(
+        "--batch-size", type=int, default=2048, help="Batch size."
+    )
+    # model
+    parser.add_argument(
+        "--num-layers", type=int, default=7, help="Number of GNN layers."
+    )
+    parser.add_argument(
+        "--hid-dim", type=int, default=256, help="Hidden channel size."
+    )
+    # learnable parameters in aggr
+    parser.add_argument("--learn-beta", action="store_true")
+
+    args = parser.parse_args()
+    print(args)
+
+    main()

examples/pytorch/deepergcn/models.py

@@ -0,0 +1,90 @@
+import dgl.function as fn
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.nn.pytorch.glob import AvgPooling
+from layers import GENConv
+from ogb.graphproppred.mol_encoder import AtomEncoder
+
+
+class DeeperGCN(nn.Module):
+    r"""
+
+    Description
+    -----------
+    Introduced in "DeeperGCN: All You Need to Train Deeper GCNs <https://arxiv.org/abs/2006.07739>"
+
+    Parameters
+    ----------
+    node_feat_dim: int
+        Size of node feature.
+    edge_feat_dim: int
+        Size of edge feature.
+    hid_dim: int
+        Size of hidden representations.
+    out_dim: int
+        Size of output.
+    num_layers: int
+        Number of graph convolutional layers.
+    dropout: float
+        Dropout rate. Default is 0.
+    beta: float
+        A continuous variable called an inverse temperature. Default is 1.0.
+    learn_beta: bool
+        Whether beta is a learnable weight. Default is False.
+    aggr: str
+        Type of aggregation. Default is 'softmax'.
+    mlp_layers: int
+        Number of MLP layers in message normalization. Default is 1.
+    """
+
+    def __init__(
+        self,
+        node_feat_dim,
+        edge_feat_dim,
+        hid_dim,
+        out_dim,
+        num_layers,
+        dropout=0.0,
+        beta=1.0,
+        learn_beta=False,
+        aggr="softmax",
+        mlp_layers=1,
+    ):
+        super(DeeperGCN, self).__init__()
+
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.gcns = nn.ModuleList()
+        self.norms = nn.ModuleList()
+
+        for _ in range(self.num_layers):
+            conv = GENConv(
+                in_dim=hid_dim,
+                out_dim=hid_dim,
+                aggregator=aggr,
+                beta=beta,
+                learn_beta=learn_beta,
+                mlp_layers=mlp_layers,
+            )
+
+            self.gcns.append(conv)
+            self.norms.append(nn.BatchNorm1d(hid_dim, affine=True))
+
+        self.node_encoder = AtomEncoder(hid_dim)
+        self.pooling = AvgPooling()
+        self.output = nn.Linear(hid_dim, out_dim)
+
+    def forward(self, g, edge_feats, node_feats=None):
+        with g.local_scope():
+            hv = self.node_encoder(node_feats)
+            he = edge_feats
+
+            for layer in range(self.num_layers):
+                hv1 = self.norms[layer](hv)
+                hv1 = F.relu(hv1)
+                hv1 = F.dropout(hv1, p=self.dropout, training=self.training)
+                hv = self.gcns[layer](g, hv1, he) + hv
+
+            h_g = self.pooling(g, hv)
+
+            return self.output(h_g)

examples/pytorch/deepergcn/modules.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MLP(nn.Sequential):
+    r"""
+
+    Description
+    -----------
+    From equation (5) in "DeeperGCN: All You Need to Train Deeper GCNs <https://arxiv.org/abs/2006.07739>"
+    """
+
+    def __init__(self, channels, act="relu", dropout=0.0, bias=True):
+        layers = []
+
+        for i in range(1, len(channels)):
+            layers.append(nn.Linear(channels[i - 1], channels[i], bias))
+            if i < len(channels) - 1:
+                layers.append(nn.BatchNorm1d(channels[i], affine=True))
+                layers.append(nn.ReLU())
+                layers.append(nn.Dropout(dropout))
+
+        super(MLP, self).__init__(*layers)
+
+
+class MessageNorm(nn.Module):
+    r"""
+
+    Description
+    -----------
+    Message normalization was introduced in "DeeperGCN: All You Need to Train Deeper GCNs <https://arxiv.org/abs/2006.07739>"
+
+    Parameters
+    ----------
+    learn_scale: bool
+        Whether s is a learnable scaling factor or not. Default is False.
+    """
+
+    def __init__(self, learn_scale=False):
+        super(MessageNorm, self).__init__()
+        self.scale = nn.Parameter(
+            torch.FloatTensor([1.0]), requires_grad=learn_scale
+        )
+
+    def forward(self, feats, msg, p=2):
+        msg = F.normalize(msg, p=2, dim=-1)
+        feats_norm = feats.norm(p=p, dim=-1, keepdim=True)
+        return msg * feats_norm * self.scale

examples/pytorch/deepwalk/README.md

@@ -0,0 +1,5 @@
+# DeepWalk
+
+- Paper link: [here](https://arxiv.org/pdf/1403.6652.pdf)
+
+The example code was moved to examples/pytorch/ogb/deepwalk.

examples/pytorch/dgi/README.md

@@ -0,0 +1,38 @@
+Deep Graph Infomax (DGI)
+========================
+
+- Paper link: [https://arxiv.org/abs/1809.10341](https://arxiv.org/abs/1809.10341)
+- Author's code repo (in Pytorch):
+  [https://github.com/PetarV-/DGI](https://github.com/PetarV-/DGI)
+
+Dependencies
+------------
+- PyTorch 0.4.1+
+- requests
+
+```bash
+pip install torch requests
+```
+
+How to run
+----------
+
+Run with following:
+
+```bash
+python3 train.py --dataset=cora --gpu=0 --self-loop
+```
+
+```bash
+python3 train.py --dataset=citeseer --gpu=0
+```
+
+```bash
+python3 train.py --dataset=pubmed --gpu=0
+```
+
+Results
+-------
+* cora: ~81.6 (81.2-82.1) (paper: 82.3)
+* citeseer: ~69.4 (paper: 71.8)
+* pubmed: ~76.1 (paper: 76.8)

examples/pytorch/dgi/dgi.py

@@ -0,0 +1,87 @@
+"""
+Deep Graph Infomax in DGL
+
+References
+----------
+Papers: https://arxiv.org/abs/1809.10341
+Author's code: https://github.com/PetarV-/DGI
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+from gcn import GCN
+
+
+class Encoder(nn.Module):
+    def __init__(self, g, in_feats, n_hidden, n_layers, activation, dropout):
+        super(Encoder, self).__init__()
+        self.g = g
+        self.conv = GCN(
+            g, in_feats, n_hidden, n_hidden, n_layers, activation, dropout
+        )
+
+    def forward(self, features, corrupt=False):
+        if corrupt:
+            perm = torch.randperm(self.g.num_nodes())
+            features = features[perm]
+        features = self.conv(features)
+        return features
+
+
+class Discriminator(nn.Module):
+    def __init__(self, n_hidden):
+        super(Discriminator, self).__init__()
+        self.weight = nn.Parameter(torch.Tensor(n_hidden, n_hidden))
+        self.reset_parameters()
+
+    def uniform(self, size, tensor):
+        bound = 1.0 / math.sqrt(size)
+        if tensor is not None:
+            tensor.data.uniform_(-bound, bound)
+
+    def reset_parameters(self):
+        size = self.weight.size(0)
+        self.uniform(size, self.weight)
+
+    def forward(self, features, summary):
+        features = torch.matmul(features, torch.matmul(self.weight, summary))
+        return features
+
+
+class DGI(nn.Module):
+    def __init__(self, g, in_feats, n_hidden, n_layers, activation, dropout):
+        super(DGI, self).__init__()
+        self.encoder = Encoder(
+            g, in_feats, n_hidden, n_layers, activation, dropout
+        )
+        self.discriminator = Discriminator(n_hidden)
+        self.loss = nn.BCEWithLogitsLoss()
+
+    def forward(self, features):
+        positive = self.encoder(features, corrupt=False)
+        negative = self.encoder(features, corrupt=True)
+        summary = torch.sigmoid(positive.mean(dim=0))
+
+        positive = self.discriminator(positive, summary)
+        negative = self.discriminator(negative, summary)
+
+        l1 = self.loss(positive, torch.ones_like(positive))
+        l2 = self.loss(negative, torch.zeros_like(negative))
+
+        return l1 + l2
+
+
+class Classifier(nn.Module):
+    def __init__(self, n_hidden, n_classes):
+        super(Classifier, self).__init__()
+        self.fc = nn.Linear(n_hidden, n_classes)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.fc.reset_parameters()
+
+    def forward(self, features):
+        features = self.fc(features)
+        return torch.log_softmax(features, dim=-1)

examples/pytorch/dgi/gcn.py

@@ -0,0 +1,34 @@
+"""
+This code was copied from the GCN implementation in DGL examples.
+"""
+import torch
+import torch.nn as nn
+
+from dgl.nn.pytorch import GraphConv
+
+
+class GCN(nn.Module):
+    def __init__(
+        self, g, in_feats, n_hidden, n_classes, n_layers, activation, dropout
+    ):
+        super(GCN, self).__init__()
+        self.g = g
+        self.layers = nn.ModuleList()
+        # input layer
+        self.layers.append(GraphConv(in_feats, n_hidden, activation=activation))
+        # hidden layers
+        for i in range(n_layers - 1):
+            self.layers.append(
+                GraphConv(n_hidden, n_hidden, activation=activation)
+            )
+        # output layer
+        self.layers.append(GraphConv(n_hidden, n_classes))
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, features):
+        h = features
+        for i, layer in enumerate(self.layers):
+            if i != 0:
+                h = self.dropout(h)
+            h = layer(self.g, h)
+        return h

examples/pytorch/dgi/train.py

@@ -0,0 +1,213 @@
+import argparse, time
+
+import dgl
+import networkx as nx
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgi import Classifier, DGI
+from dgl import DGLGraph
+from dgl.data import load_data, register_data_args
+
+
+def evaluate(model, features, labels, mask):
+    model.eval()
+    with torch.no_grad():
+        logits = model(features)
+        logits = logits[mask]
+        labels = labels[mask]
+        _, indices = torch.max(logits, dim=1)
+        correct = torch.sum(indices == labels)
+        return correct.item() * 1.0 / len(labels)
+
+
+def main(args):
+    # load and preprocess dataset
+    data = load_data(args)
+    g = data[0]
+    features = torch.FloatTensor(g.ndata["feat"])
+    labels = torch.LongTensor(g.ndata["label"])
+    if hasattr(torch, "BoolTensor"):
+        train_mask = torch.BoolTensor(g.ndata["train_mask"])
+        val_mask = torch.BoolTensor(g.ndata["val_mask"])
+        test_mask = torch.BoolTensor(g.ndata["test_mask"])
+    else:
+        train_mask = torch.ByteTensor(g.ndata["train_mask"])
+        val_mask = torch.ByteTensor(g.ndata["val_mask"])
+        test_mask = torch.ByteTensor(g.ndata["test_mask"])
+    in_feats = features.shape[1]
+    n_classes = data.num_classes
+    n_edges = g.num_edges()
+
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        torch.cuda.set_device(args.gpu)
+        features = features.cuda()
+        labels = labels.cuda()
+        train_mask = train_mask.cuda()
+        val_mask = val_mask.cuda()
+        test_mask = test_mask.cuda()
+
+    # add self loop
+    if args.self_loop:
+        g = dgl.remove_self_loop(g)
+        g = dgl.add_self_loop(g)
+    n_edges = g.num_edges()
+
+    if args.gpu >= 0:
+        g = g.to(args.gpu)
+    # create DGI model
+    dgi = DGI(
+        g,
+        in_feats,
+        args.n_hidden,
+        args.n_layers,
+        nn.PReLU(args.n_hidden),
+        args.dropout,
+    )
+
+    if cuda:
+        dgi.cuda()
+
+    dgi_optimizer = torch.optim.Adam(
+        dgi.parameters(), lr=args.dgi_lr, weight_decay=args.weight_decay
+    )
+
+    # train deep graph infomax
+    cnt_wait = 0
+    best = 1e9
+    best_t = 0
+    mean = 0
+    for epoch in range(args.n_dgi_epochs):
+        dgi.train()
+        if epoch >= 3:
+            t0 = time.time()
+
+        dgi_optimizer.zero_grad()
+        loss = dgi(features)
+        loss.backward()
+        dgi_optimizer.step()
+
+        if loss < best:
+            best = loss
+            best_t = epoch
+            cnt_wait = 0
+            torch.save(dgi.state_dict(), "best_dgi.pkl")
+        else:
+            cnt_wait += 1
+
+        if cnt_wait == args.patience:
+            print("Early stopping!")
+            break
+
+        if epoch >= 3:
+            mean = (mean * (epoch - 3) + (time.time() - t0)) / (epoch - 2)
+
+            print(
+                "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | "
+                "ETputs(KTEPS) {:.2f}".format(
+                    epoch, mean, loss.item(), n_edges / mean / 1000
+                )
+            )
+
+    # create classifier model
+    classifier = Classifier(args.n_hidden, n_classes)
+    if cuda:
+        classifier.cuda()
+
+    classifier_optimizer = torch.optim.Adam(
+        classifier.parameters(),
+        lr=args.classifier_lr,
+        weight_decay=args.weight_decay,
+    )
+
+    # train classifier
+    print("Loading {}th epoch".format(best_t))
+    dgi.load_state_dict(torch.load("best_dgi.pkl", weights_only=False))
+    embeds = dgi.encoder(features, corrupt=False)
+    embeds = embeds.detach()
+    mean = 0
+    for epoch in range(args.n_classifier_epochs):
+        classifier.train()
+        if epoch >= 3:
+            t0 = time.time()
+
+        classifier_optimizer.zero_grad()
+        preds = classifier(embeds)
+        loss = F.nll_loss(preds[train_mask], labels[train_mask])
+        loss.backward()
+        classifier_optimizer.step()
+
+        if epoch >= 3:
+            mean = (mean * (epoch - 3) + (time.time() - t0)) / (epoch - 2)
+
+            acc = evaluate(classifier, embeds, labels, val_mask)
+            print(
+                "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
+                "ETputs(KTEPS) {:.2f}".format(
+                    epoch,
+                    mean,
+                    loss.item(),
+                    acc,
+                    n_edges / mean / 1000,
+                )
+            )
+
+    print()
+    acc = evaluate(classifier, embeds, labels, test_mask)
+    print("Test Accuracy {:.4f}".format(acc))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="DGI")
+    register_data_args(parser)
+    parser.add_argument(
+        "--dropout", type=float, default=0.0, help="dropout probability"
+    )
+    parser.add_argument("--gpu", type=int, default=-1, help="gpu")
+    parser.add_argument(
+        "--dgi-lr", type=float, default=1e-3, help="dgi learning rate"
+    )
+    parser.add_argument(
+        "--classifier-lr",
+        type=float,
+        default=1e-2,
+        help="classifier learning rate",
+    )
+    parser.add_argument(
+        "--n-dgi-epochs",
+        type=int,
+        default=300,
+        help="number of training epochs",
+    )
+    parser.add_argument(
+        "--n-classifier-epochs",
+        type=int,
+        default=300,
+        help="number of training epochs",
+    )
+    parser.add_argument(
+        "--n-hidden", type=int, default=512, help="number of hidden gcn units"
+    )
+    parser.add_argument(
+        "--n-layers", type=int, default=1, help="number of hidden gcn layers"
+    )
+    parser.add_argument(
+        "--weight-decay", type=float, default=0.0, help="Weight for L2 loss"
+    )
+    parser.add_argument(
+        "--patience", type=int, default=20, help="early stop patience condition"
+    )
+    parser.add_argument(
+        "--self-loop",
+        action="store_true",
+        help="graph self-loop (default=False)",
+    )
+    parser.set_defaults(self_loop=False)
+    args = parser.parse_args()
+    print(args)
+
+    main(args)

examples/pytorch/dgmg/README.md

@@ -0,0 +1,28 @@
+# Learning Deep Generative Models of Graphs
+
+This is an implementation of [Learning Deep Generative Models of Graphs](https://arxiv.org/pdf/1803.03324.pdf) by 
+Yujia Li, Oriol Vinyals, Chris Dyer, Razvan Pascanu, Peter Battaglia.
+
+For molecule generation, see 
+[DGL-LifeSci](https://github.com/awslabs/dgl-lifesci/tree/master/examples/generative_models/dgmg).
+
+## Dependencies
+- Python 3.5.2
+- [Pytorch 0.4.1](https://pytorch.org/)
+- [Matplotlib 2.2.2](https://matplotlib.org/)
+
+## Usage
+
+`python3 main.py`
+
+## Performance
+
+90% accuracy for cycles compared with 84% accuracy reported in the original paper.
+
+## Speed
+
+On AWS p3.2x instance (w/ V100), one epoch takes ~526s.
+
+## Acknowledgement
+
+We would like to thank Yujia Li for providing details on the implementation.

examples/pytorch/dgmg/configure.py

@@ -0,0 +1,33 @@
+"""We intend to make our reproduction as close as possible to the original paper.
+The configuration in the file is mostly from the description in the original paper
+and will be loaded when setting up."""
+
+
+def dataset_based_configure(opts):
+    if opts["dataset"] == "cycles":
+        ds_configure = cycles_configure
+    else:
+        raise ValueError("Unsupported dataset: {}".format(opts["dataset"]))
+
+    opts = {**opts, **ds_configure}
+
+    return opts
+
+
+synthetic_dataset_configure = {
+    "node_hidden_size": 16,
+    "num_propagation_rounds": 2,
+    "optimizer": "Adam",
+    "nepochs": 25,
+    "ds_size": 4000,
+    "num_generated_samples": 10000,
+}
+
+cycles_configure = {
+    **synthetic_dataset_configure,
+    **{
+        "min_size": 10,
+        "max_size": 20,
+        "lr": 5e-4,
+    },
+}

examples/pytorch/dgmg/cycles.py

@@ -0,0 +1,226 @@
+import os
+import pickle
+import random
+
+import matplotlib.pyplot as plt
+import networkx as nx
+from torch.utils.data import Dataset
+
+
+def get_previous(i, v_max):
+    if i == 0:
+        return v_max
+    else:
+        return i - 1
+
+
+def get_next(i, v_max):
+    if i == v_max:
+        return 0
+    else:
+        return i + 1
+
+
+def is_cycle(g):
+    size = g.num_nodes()
+
+    if size < 3:
+        return False
+
+    for node in range(size):
+        neighbors = g.successors(node)
+
+        if len(neighbors) != 2:
+            return False
+
+        if get_previous(node, size - 1) not in neighbors:
+            return False
+
+        if get_next(node, size - 1) not in neighbors:
+            return False
+
+    return True
+
+
+def get_decision_sequence(size):
+    """
+    Get the decision sequence for generating valid cycles with DGMG for teacher
+    forcing optimization.
+    """
+    decision_sequence = []
+
+    for i in range(size):
+        decision_sequence.append(0)  # Add node
+
+        if i != 0:
+            decision_sequence.append(0)  # Add edge
+            decision_sequence.append(
+                i - 1
+            )  # Set destination to be previous node.
+
+        if i == size - 1:
+            decision_sequence.append(0)  # Add edge
+            decision_sequence.append(0)  # Set destination to be the root.
+
+        decision_sequence.append(1)  # Stop adding edge
+
+    decision_sequence.append(1)  # Stop adding node
+
+    return decision_sequence
+
+
+def generate_dataset(v_min, v_max, n_samples, fname):
+    samples = []
+    for _ in range(n_samples):
+        size = random.randint(v_min, v_max)
+        samples.append(get_decision_sequence(size))
+
+    with open(fname, "wb") as f:
+        pickle.dump(samples, f)
+
+
+class CycleDataset(Dataset):
+    def __init__(self, fname):
+        super(CycleDataset, self).__init__()
+
+        with open(fname, "rb") as f:
+            self.dataset = pickle.load(f)
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        return self.dataset[index]
+
+    def collate_single(self, batch):
+        assert len(batch) == 1, "Currently we do not support batched training"
+        return batch[0]
+
+    def collate_batch(self, batch):
+        return batch
+
+
+def dglGraph_to_adj_list(g):
+    adj_list = {}
+    for node in range(g.num_nodes()):
+        # For undirected graph. successors and
+        # predecessors are equivalent.
+        adj_list[node] = g.successors(node).tolist()
+    return adj_list
+
+
+class CycleModelEvaluation(object):
+    def __init__(self, v_min, v_max, dir):
+        super(CycleModelEvaluation, self).__init__()
+
+        self.v_min = v_min
+        self.v_max = v_max
+
+        self.dir = dir
+
+    def rollout_and_examine(self, model, num_samples):
+        assert not model.training, "You need to call model.eval()."
+
+        num_total_size = 0
+        num_valid_size = 0
+        num_cycle = 0
+        num_valid = 0
+        plot_times = 0
+        adj_lists_to_plot = []
+
+        for i in range(num_samples):
+            sampled_graph = model()
+            if isinstance(sampled_graph, list):
+                # When the model is a batched implementation, a list of
+                # DGLGraph objects is returned. Note that with model(),
+                # we generate a single graph as with the non-batched
+                # implementation. We actually support batched generation
+                # during the inference so feel free to modify the code.
+                sampled_graph = sampled_graph[0]
+
+            sampled_adj_list = dglGraph_to_adj_list(sampled_graph)
+            adj_lists_to_plot.append(sampled_adj_list)
+
+            graph_size = sampled_graph.num_nodes()
+            valid_size = self.v_min <= graph_size <= self.v_max
+            cycle = is_cycle(sampled_graph)
+
+            num_total_size += graph_size
+
+            if valid_size:
+                num_valid_size += 1
+
+            if cycle:
+                num_cycle += 1
+
+            if valid_size and cycle:
+                num_valid += 1
+
+            if len(adj_lists_to_plot) >= 4:
+                plot_times += 1
+                fig, ((ax0, ax1), (ax2, ax3)) = plt.subplots(2, 2)
+                axes = {0: ax0, 1: ax1, 2: ax2, 3: ax3}
+                for i in range(4):
+                    nx.draw_circular(
+                        nx.from_dict_of_lists(adj_lists_to_plot[i]),
+                        with_labels=True,
+                        ax=axes[i],
+                    )
+
+                plt.savefig(self.dir + "/samples/{:d}".format(plot_times))
+                plt.close()
+
+                adj_lists_to_plot = []
+
+        self.num_samples_examined = num_samples
+        self.average_size = num_total_size / num_samples
+        self.valid_size_ratio = num_valid_size / num_samples
+        self.cycle_ratio = num_cycle / num_samples
+        self.valid_ratio = num_valid / num_samples
+
+    def write_summary(self):
+        def _format_value(v):
+            if isinstance(v, float):
+                return "{:.4f}".format(v)
+            elif isinstance(v, int):
+                return "{:d}".format(v)
+            else:
+                return "{}".format(v)
+
+        statistics = {
+            "num_samples": self.num_samples_examined,
+            "v_min": self.v_min,
+            "v_max": self.v_max,
+            "average_size": self.average_size,
+            "valid_size_ratio": self.valid_size_ratio,
+            "cycle_ratio": self.cycle_ratio,
+            "valid_ratio": self.valid_ratio,
+        }
+
+        model_eval_path = os.path.join(self.dir, "model_eval.txt")
+
+        with open(model_eval_path, "w") as f:
+            for key, value in statistics.items():
+                msg = "{}\t{}\n".format(key, _format_value(value))
+                f.write(msg)
+
+        print("Saved model evaluation statistics to {}".format(model_eval_path))
+
+
+class CyclePrinting(object):
+    def __init__(self, num_epochs, num_batches):
+        super(CyclePrinting, self).__init__()
+
+        self.num_epochs = num_epochs
+        self.num_batches = num_batches
+        self.batch_count = 0
+
+    def update(self, epoch, metrics):
+        self.batch_count = (self.batch_count) % self.num_batches + 1
+
+        msg = "epoch {:d}/{:d}, batch {:d}/{:d}".format(
+            epoch, self.num_epochs, self.batch_count, self.num_batches
+        )
+        for key, value in metrics.items():
+            msg += ", {}: {:4f}".format(key, value)
+        print(msg)

examples/pytorch/dgmg/main.py

@@ -0,0 +1,178 @@
+"""
+Learning Deep Generative Models of Graphs
+Paper: https://arxiv.org/pdf/1803.03324.pdf
+
+This implementation works with a minibatch of size 1 only for both training and inference.
+"""
+import argparse
+import datetime
+import time
+
+import torch
+from model import DGMG
+from torch.nn.utils import clip_grad_norm_
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+
+
+def main(opts):
+    t1 = time.time()
+
+    # Setup dataset and data loader
+    if opts["dataset"] == "cycles":
+        from cycles import CycleDataset, CycleModelEvaluation, CyclePrinting
+
+        dataset = CycleDataset(fname=opts["path_to_dataset"])
+        evaluator = CycleModelEvaluation(
+            v_min=opts["min_size"], v_max=opts["max_size"], dir=opts["log_dir"]
+        )
+        printer = CyclePrinting(
+            num_epochs=opts["nepochs"],
+            num_batches=opts["ds_size"] // opts["batch_size"],
+        )
+    else:
+        raise ValueError("Unsupported dataset: {}".format(opts["dataset"]))
+
+    data_loader = DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=True,
+        num_workers=0,
+        collate_fn=dataset.collate_single,
+    )
+
+    # Initialize_model
+    model = DGMG(
+        v_max=opts["max_size"],
+        node_hidden_size=opts["node_hidden_size"],
+        num_prop_rounds=opts["num_propagation_rounds"],
+    )
+
+    # Initialize optimizer
+    if opts["optimizer"] == "Adam":
+        optimizer = Adam(model.parameters(), lr=opts["lr"])
+    else:
+        raise ValueError("Unsupported argument for the optimizer")
+
+    t2 = time.time()
+
+    # Training
+    model.train()
+    for epoch in range(opts["nepochs"]):
+        batch_count = 0
+        batch_loss = 0
+        batch_prob = 0
+        optimizer.zero_grad()
+
+        for i, data in enumerate(data_loader):
+            log_prob = model(actions=data)
+            prob = log_prob.detach().exp()
+
+            loss = -log_prob / opts["batch_size"]
+            prob_averaged = prob / opts["batch_size"]
+
+            loss.backward()
+
+            batch_loss += loss.item()
+            batch_prob += prob_averaged.item()
+            batch_count += 1
+
+            if batch_count % opts["batch_size"] == 0:
+                printer.update(
+                    epoch + 1,
+                    {"averaged_loss": batch_loss, "averaged_prob": batch_prob},
+                )
+
+                if opts["clip_grad"]:
+                    clip_grad_norm_(model.parameters(), opts["clip_bound"])
+
+                optimizer.step()
+
+                batch_loss = 0
+                batch_prob = 0
+                optimizer.zero_grad()
+
+    t3 = time.time()
+
+    model.eval()
+    evaluator.rollout_and_examine(model, opts["num_generated_samples"])
+    evaluator.write_summary()
+
+    t4 = time.time()
+
+    print("It took {} to setup.".format(datetime.timedelta(seconds=t2 - t1)))
+    print(
+        "It took {} to finish training.".format(
+            datetime.timedelta(seconds=t3 - t2)
+        )
+    )
+    print(
+        "It took {} to finish evaluation.".format(
+            datetime.timedelta(seconds=t4 - t3)
+        )
+    )
+    print(
+        "--------------------------------------------------------------------------"
+    )
+    print(
+        "On average, an epoch takes {}.".format(
+            datetime.timedelta(seconds=(t3 - t2) / opts["nepochs"])
+        )
+    )
+
+    del model.g
+    torch.save(model, "./model.pth")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="DGMG")
+
+    # configure
+    parser.add_argument("--seed", type=int, default=9284, help="random seed")
+
+    # dataset
+    parser.add_argument(
+        "--dataset", choices=["cycles"], default="cycles", help="dataset to use"
+    )
+    parser.add_argument(
+        "--path-to-dataset",
+        type=str,
+        default="cycles.p",
+        help="load the dataset if it exists, "
+        "generate it and save to the path otherwise",
+    )
+
+    # log
+    parser.add_argument(
+        "--log-dir",
+        default="./results",
+        help="folder to save info like experiment configuration "
+        "or model evaluation results",
+    )
+
+    # optimization
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=10,
+        help="batch size to use for training",
+    )
+    parser.add_argument(
+        "--clip-grad",
+        action="store_true",
+        default=True,
+        help="gradient clipping is required to prevent gradient explosion",
+    )
+    parser.add_argument(
+        "--clip-bound",
+        type=float,
+        default=0.25,
+        help="constraint of gradient norm for gradient clipping",
+    )
+
+    args = parser.parse_args()
+    from utils import setup
+
+    opts = setup(args)
+
+    main(opts)

examples/pytorch/dgmg/model.py

@@ -0,0 +1,346 @@
+from functools import partial
+
+import dgl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Bernoulli, Categorical
+
+
+class GraphEmbed(nn.Module):
+    def __init__(self, node_hidden_size):
+        super(GraphEmbed, self).__init__()
+
+        # Setting from the paper
+        self.graph_hidden_size = 2 * node_hidden_size
+
+        # Embed graphs
+        self.node_gating = nn.Sequential(
+            nn.Linear(node_hidden_size, 1), nn.Sigmoid()
+        )
+        self.node_to_graph = nn.Linear(node_hidden_size, self.graph_hidden_size)
+
+    def forward(self, g):
+        if g.num_nodes() == 0:
+            return torch.zeros(1, self.graph_hidden_size)
+        else:
+            # Node features are stored as hv in ndata.
+            hvs = g.ndata["hv"]
+            return (self.node_gating(hvs) * self.node_to_graph(hvs)).sum(
+                0, keepdim=True
+            )
+
+
+class GraphProp(nn.Module):
+    def __init__(self, num_prop_rounds, node_hidden_size):
+        super(GraphProp, self).__init__()
+
+        self.num_prop_rounds = num_prop_rounds
+
+        # Setting from the paper
+        self.node_activation_hidden_size = 2 * node_hidden_size
+
+        message_funcs = []
+        self.reduce_funcs = []
+        node_update_funcs = []
+
+        for t in range(num_prop_rounds):
+            # input being [hv, hu, xuv]
+            message_funcs.append(
+                nn.Linear(
+                    2 * node_hidden_size + 1, self.node_activation_hidden_size
+                )
+            )
+
+            self.reduce_funcs.append(partial(self.dgmg_reduce, round=t))
+            node_update_funcs.append(
+                nn.GRUCell(self.node_activation_hidden_size, node_hidden_size)
+            )
+
+        self.message_funcs = nn.ModuleList(message_funcs)
+        self.node_update_funcs = nn.ModuleList(node_update_funcs)
+
+    def dgmg_msg(self, edges):
+        """For an edge u->v, return concat([h_u, x_uv])"""
+        return {"m": torch.cat([edges.src["hv"], edges.data["he"]], dim=1)}
+
+    def dgmg_reduce(self, nodes, round):
+        hv_old = nodes.data["hv"]
+        m = nodes.mailbox["m"]
+        message = torch.cat(
+            [hv_old.unsqueeze(1).expand(-1, m.size(1), -1), m], dim=2
+        )
+        node_activation = (self.message_funcs[round](message)).sum(1)
+
+        return {"a": node_activation}
+
+    def forward(self, g):
+        if g.num_edges() == 0:
+            return
+        else:
+            for t in range(self.num_prop_rounds):
+                g.update_all(
+                    message_func=self.dgmg_msg, reduce_func=self.reduce_funcs[t]
+                )
+                g.ndata["hv"] = self.node_update_funcs[t](
+                    g.ndata["a"], g.ndata["hv"]
+                )
+
+
+def bernoulli_action_log_prob(logit, action):
+    """Calculate the log p of an action with respect to a Bernoulli
+    distribution. Use logit rather than prob for numerical stability."""
+    if action == 0:
+        return F.logsigmoid(-logit)
+    else:
+        return F.logsigmoid(logit)
+
+
+class AddNode(nn.Module):
+    def __init__(self, graph_embed_func, node_hidden_size):
+        super(AddNode, self).__init__()
+
+        self.graph_op = {"embed": graph_embed_func}
+
+        self.stop = 1
+        self.add_node = nn.Linear(graph_embed_func.graph_hidden_size, 1)
+
+        # If to add a node, initialize its hv
+        self.node_type_embed = nn.Embedding(1, node_hidden_size)
+        self.initialize_hv = nn.Linear(
+            node_hidden_size + graph_embed_func.graph_hidden_size,
+            node_hidden_size,
+        )
+
+        self.init_node_activation = torch.zeros(1, 2 * node_hidden_size)
+
+    def _initialize_node_repr(self, g, node_type, graph_embed):
+        num_nodes = g.num_nodes()
+        hv_init = self.initialize_hv(
+            torch.cat(
+                [
+                    self.node_type_embed(torch.LongTensor([node_type])),
+                    graph_embed,
+                ],
+                dim=1,
+            )
+        )
+        g.nodes[num_nodes - 1].data["hv"] = hv_init
+        g.nodes[num_nodes - 1].data["a"] = self.init_node_activation
+
+    def prepare_training(self):
+        self.log_prob = []
+
+    def forward(self, g, action=None):
+        graph_embed = self.graph_op["embed"](g)
+
+        logit = self.add_node(graph_embed)
+        prob = torch.sigmoid(logit)
+
+        if not self.training:
+            action = Bernoulli(prob).sample().item()
+        stop = bool(action == self.stop)
+
+        if not stop:
+            g.add_nodes(1)
+            self._initialize_node_repr(g, action, graph_embed)
+
+        if self.training:
+            sample_log_prob = bernoulli_action_log_prob(logit, action)
+            self.log_prob.append(sample_log_prob)
+
+        return stop
+
+
+class AddEdge(nn.Module):
+    def __init__(self, graph_embed_func, node_hidden_size):
+        super(AddEdge, self).__init__()
+
+        self.graph_op = {"embed": graph_embed_func}
+        self.add_edge = nn.Linear(
+            graph_embed_func.graph_hidden_size + node_hidden_size, 1
+        )
+
+    def prepare_training(self):
+        self.log_prob = []
+
+    def forward(self, g, action=None):
+        graph_embed = self.graph_op["embed"](g)
+        src_embed = g.nodes[g.num_nodes() - 1].data["hv"]
+
+        logit = self.add_edge(torch.cat([graph_embed, src_embed], dim=1))
+        prob = torch.sigmoid(logit)
+
+        if not self.training:
+            action = Bernoulli(prob).sample().item()
+        to_add_edge = bool(action == 0)
+
+        if self.training:
+            sample_log_prob = bernoulli_action_log_prob(logit, action)
+            self.log_prob.append(sample_log_prob)
+
+        return to_add_edge
+
+
+class ChooseDestAndUpdate(nn.Module):
+    def __init__(self, graph_prop_func, node_hidden_size):
+        super(ChooseDestAndUpdate, self).__init__()
+
+        self.graph_op = {"prop": graph_prop_func}
+        self.choose_dest = nn.Linear(2 * node_hidden_size, 1)
+
+    def _initialize_edge_repr(self, g, src_list, dest_list):
+        # For untyped edges, we only add 1 to indicate its existence.
+        # For multiple edge types, we can use a one hot representation
+        # or an embedding module.
+        edge_repr = torch.ones(len(src_list), 1)
+        g.edges[src_list, dest_list].data["he"] = edge_repr
+
+    def prepare_training(self):
+        self.log_prob = []
+
+    def forward(self, g, dest):
+        src = g.num_nodes() - 1
+        possible_dests = range(src)
+
+        src_embed_expand = g.nodes[src].data["hv"].expand(src, -1)
+        possible_dests_embed = g.nodes[possible_dests].data["hv"]
+
+        dests_scores = self.choose_dest(
+            torch.cat([possible_dests_embed, src_embed_expand], dim=1)
+        ).view(1, -1)
+        dests_probs = F.softmax(dests_scores, dim=1)
+
+        if not self.training:
+            dest = Categorical(dests_probs).sample().item()
+
+        if not g.has_edges_between(src, dest):
+            # For undirected graphs, we add edges for both directions
+            # so that we can perform graph propagation.
+            src_list = [src, dest]
+            dest_list = [dest, src]
+
+            g.add_edges(src_list, dest_list)
+            self._initialize_edge_repr(g, src_list, dest_list)
+
+            self.graph_op["prop"](g)
+
+        if self.training:
+            if dests_probs.nelement() > 1:
+                self.log_prob.append(
+                    F.log_softmax(dests_scores, dim=1)[:, dest : dest + 1]
+                )
+
+
+class DGMG(nn.Module):
+    def __init__(self, v_max, node_hidden_size, num_prop_rounds):
+        super(DGMG, self).__init__()
+
+        # Graph configuration
+        self.v_max = v_max
+
+        # Graph embedding module
+        self.graph_embed = GraphEmbed(node_hidden_size)
+
+        # Graph propagation module
+        self.graph_prop = GraphProp(num_prop_rounds, node_hidden_size)
+
+        # Actions
+        self.add_node_agent = AddNode(self.graph_embed, node_hidden_size)
+        self.add_edge_agent = AddEdge(self.graph_embed, node_hidden_size)
+        self.choose_dest_agent = ChooseDestAndUpdate(
+            self.graph_prop, node_hidden_size
+        )
+
+        # Weight initialization
+        self.init_weights()
+
+    def init_weights(self):
+        from utils import dgmg_message_weight_init, weights_init
+
+        self.graph_embed.apply(weights_init)
+        self.graph_prop.apply(weights_init)
+        self.add_node_agent.apply(weights_init)
+        self.add_edge_agent.apply(weights_init)
+        self.choose_dest_agent.apply(weights_init)
+
+        self.graph_prop.message_funcs.apply(dgmg_message_weight_init)
+
+    @property
+    def action_step(self):
+        old_step_count = self.step_count
+        self.step_count += 1
+
+        return old_step_count
+
+    def prepare_for_train(self):
+        self.step_count = 0
+
+        self.add_node_agent.prepare_training()
+        self.add_edge_agent.prepare_training()
+        self.choose_dest_agent.prepare_training()
+
+    def add_node_and_update(self, a=None):
+        """Decide if to add a new node.
+        If a new node should be added, update the graph."""
+
+        return self.add_node_agent(self.g, a)
+
+    def add_edge_or_not(self, a=None):
+        """Decide if a new edge should be added."""
+
+        return self.add_edge_agent(self.g, a)
+
+    def choose_dest_and_update(self, a=None):
+        """Choose destination and connect it to the latest node.
+        Add edges for both directions and update the graph."""
+
+        self.choose_dest_agent(self.g, a)
+
+    def get_log_prob(self):
+        return (
+            torch.cat(self.add_node_agent.log_prob).sum()
+            + torch.cat(self.add_edge_agent.log_prob).sum()
+            + torch.cat(self.choose_dest_agent.log_prob).sum()
+        )
+
+    def forward_train(self, actions):
+        self.prepare_for_train()
+
+        stop = self.add_node_and_update(a=actions[self.action_step])
+
+        while not stop:
+            to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
+            while to_add_edge:
+                self.choose_dest_and_update(a=actions[self.action_step])
+                to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
+            stop = self.add_node_and_update(a=actions[self.action_step])
+
+        return self.get_log_prob()
+
+    def forward_inference(self):
+        stop = self.add_node_and_update()
+        while (not stop) and (self.g.num_nodes() < self.v_max + 1):
+            num_trials = 0
+            to_add_edge = self.add_edge_or_not()
+            while to_add_edge and (num_trials < self.g.num_nodes() - 1):
+                self.choose_dest_and_update()
+                num_trials += 1
+                to_add_edge = self.add_edge_or_not()
+            stop = self.add_node_and_update()
+
+        return self.g
+
+    def forward(self, actions=None):
+        # The graph we will work on
+        self.g = dgl.DGLGraph()
+
+        # If there are some features for nodes and edges,
+        # zero tensors will be set for those of new nodes and edges.
+        self.g.set_n_initializer(dgl.frame.zero_initializer)
+        self.g.set_e_initializer(dgl.frame.zero_initializer)
+
+        if self.training:
+            return self.forward_train(actions)
+        else:
+            return self.forward_inference()

examples/pytorch/dgmg/utils.py

@@ -0,0 +1,156 @@
+import datetime
+import os
+import random
+from pprint import pprint
+
+import matplotlib.pyplot as plt
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.init as init
+
+########################################################################################################################
+#                                                    configuration                                                     #
+########################################################################################################################
+
+
+def mkdir_p(path):
+    import errno
+
+    try:
+        os.makedirs(path)
+        print("Created directory {}".format(path))
+    except OSError as exc:
+        if exc.errno == errno.EEXIST and os.path.isdir(path):
+            print("Directory {} already exists.".format(path))
+        else:
+            raise
+
+
+def date_filename(base_dir="./"):
+    dt = datetime.datetime.now()
+    return os.path.join(
+        base_dir,
+        "{}_{:02d}-{:02d}-{:02d}".format(
+            dt.date(), dt.hour, dt.minute, dt.second
+        ),
+    )
+
+
+def setup_log_dir(opts):
+    log_dir = "{}".format(date_filename(opts["log_dir"]))
+    mkdir_p(log_dir)
+    return log_dir
+
+
+def save_arg_dict(opts, filename="settings.txt"):
+    def _format_value(v):
+        if isinstance(v, float):
+            return "{:.4f}".format(v)
+        elif isinstance(v, int):
+            return "{:d}".format(v)
+        else:
+            return "{}".format(v)
+
+    save_path = os.path.join(opts["log_dir"], filename)
+    with open(save_path, "w") as f:
+        for key, value in opts.items():
+            f.write("{}\t{}\n".format(key, _format_value(value)))
+    print("Saved settings to {}".format(save_path))
+
+
+def setup(args):
+    opts = args.__dict__.copy()
+
+    cudnn.benchmark = False
+    cudnn.deterministic = True
+
+    # Seed
+    if opts["seed"] is None:
+        opts["seed"] = random.randint(1, 10000)
+    random.seed(opts["seed"])
+    torch.manual_seed(opts["seed"])
+
+    # Dataset
+    from configure import dataset_based_configure
+
+    opts = dataset_based_configure(opts)
+
+    assert (
+        opts["path_to_dataset"] is not None
+    ), "Expect path to dataset to be set."
+    if not os.path.exists(opts["path_to_dataset"]):
+        if opts["dataset"] == "cycles":
+            from cycles import generate_dataset
+
+            generate_dataset(
+                opts["min_size"],
+                opts["max_size"],
+                opts["ds_size"],
+                opts["path_to_dataset"],
+            )
+        else:
+            raise ValueError("Unsupported dataset: {}".format(opts["dataset"]))
+
+    # Optimization
+    if opts["clip_grad"]:
+        assert (
+            opts["clip_grad"] is not None
+        ), "Expect the gradient norm constraint to be set."
+
+    # Log
+    print("Prepare logging directory...")
+    log_dir = setup_log_dir(opts)
+    opts["log_dir"] = log_dir
+    mkdir_p(log_dir + "/samples")
+
+    plt.switch_backend("Agg")
+
+    save_arg_dict(opts)
+    pprint(opts)
+
+    return opts
+
+
+########################################################################################################################
+#                                                         model                                                        #
+########################################################################################################################
+
+
+def weights_init(m):
+    """
+    Code from https://gist.github.com/jeasinema/ed9236ce743c8efaf30fa2ff732749f5
+    Usage:
+        model = Model()
+        model.apply(weight_init)
+    """
+    if isinstance(m, nn.Linear):
+        init.xavier_normal_(m.weight.data)
+        init.normal_(m.bias.data)
+    elif isinstance(m, nn.GRUCell):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+
+
+def dgmg_message_weight_init(m):
+    """
+    This is similar as the function above where we initialize linear layers from a normal distribution with std
+    1./10 as suggested by the author. This should only be used for the message passing functions, i.e. fe's in the
+    paper.
+    """
+
+    def _weight_init(m):
+        if isinstance(m, nn.Linear):
+            init.normal_(m.weight.data, std=1.0 / 10)
+            init.normal_(m.bias.data, std=1.0 / 10)
+        else:
+            raise ValueError("Expected the input to be of type nn.Linear!")
+
+    if isinstance(m, nn.ModuleList):
+        for layer in m:
+            layer.apply(_weight_init)
+    else:
+        m.apply(_weight_init)

examples/pytorch/diffpool/README.md

@@ -0,0 +1,59 @@
+Hierarchical Graph Representation Learning with Differentiable Pooling
+============
+
+
+Paper link: [https://arxiv.org/abs/1806.08804](https://arxiv.org/abs/1806.08804)
+
+Author's code repo: [https://github.com/RexYing/diffpool](https://github.com/RexYing/diffpool)
+
+This folder contains a DGL implementation of the DiffPool model. The first pooling layer is computed with DGL, and following pooling layers are computed with tensorized operation since the pooled graphs are dense.
+
+Dependencies
+------------
+* PyTorch 1.0+
+
+How to run
+----------
+
+```bash
+python train.py --dataset ENZYMES --pool_ratio 0.10 --num_pool 1 --epochs 1000
+python train.py --dataset DD --pool_ratio 0.15 --num_pool 1  --batch-size 10
+```
+Performance
+-----------
+ENZYMES 63.33% (with early stopping)
+DD 79.31% (with early stopping)
+
+
+## Update (2021-03-09)
+
+**Changes:**
+
+* Fix bug in Diffpool: the wrong `assign_dim` parameter
+* Improve efficiency of DiffPool, make the model independent of batch size. Remove redundant computation.
+
+
+**Efficiency:**
+
+On V100-SXM2 16GB
+
+|                    | Train time/epoch (original) (s) | Train time/epoch (improved) (s) |
+| ------------------ | ------------------------------: | ------------------------------: |
+| DD (batch_size=10) |                          21.302 |                      **17.282** |
+| DD (batch_size=20) |                             OOM |                      **44.682** |
+| ENZYMES            |                           1.749 |                       **1.685** |
+
+|                    | Memory usage (original) (MB) | Memory usage (improved) (MB) |
+| ------------------ | ---------------------------: | ---------------------------: |
+| DD (batch_size=10) |                     5274.620 |                 **2928.568** |
+| DD (batch_size=20) |                          OOM |                **10088.889** |
+| ENZYMES            |                       25.685 |                   **21.909** |
+
+**Accuracy**
+
+Each experiment with improved model is only conducted once, thus the result may has noise.
+
+|         |   Original |   Improved |
+| ------- | ---------: | ---------: |
+| DD      | **79.31%** |     78.33% |
+| ENZYMES |     63.33% | **68.33%** |

examples/pytorch/diffpool/data_utils.py

@@ -0,0 +1,40 @@
+import numpy as np
+import torch
+
+
+def one_hotify(labels, pad=-1):
+    """
+    cast label to one hot vector
+    """
+    num_instances = len(labels)
+    if pad <= 0:
+        dim_embedding = np.max(labels) + 1  # zero-indexed assumed
+    else:
+        assert pad > 0, "result_dim for padding one hot embedding not set!"
+        dim_embedding = pad + 1
+    embeddings = np.zeros((num_instances, dim_embedding))
+    embeddings[np.arange(num_instances), labels] = 1
+
+    return embeddings
+
+
+def pre_process(dataset, prog_args):
+    """
+    diffpool specific data partition, pre-process and shuffling
+    """
+    if prog_args.data_mode != "default":
+        print("overwrite node attributes with DiffPool's preprocess setting")
+        if prog_args.data_mode == "id":
+            for g, _ in dataset:
+                id_list = np.arange(g.num_nodes())
+                g.ndata["feat"] = one_hotify(id_list, pad=dataset.max_num_node)
+
+        elif prog_args.data_mode == "deg-num":
+            for g, _ in dataset:
+                g.ndata["feat"] = np.expand_dims(g.in_degrees(), axis=1)
+
+        elif prog_args.data_mode == "deg":
+            for g in dataset:
+                degs = list(g.in_degrees())
+                degs_one_hot = one_hotify(degs, pad=dataset.max_degrees)
+                g.ndata["feat"] = degs_one_hot

examples/pytorch/diffpool/model/dgl_layers/__init__.py

@@ -0,0 +1 @@
+from .gnn import DiffPoolBatchedGraphLayer, GraphSage, GraphSageLayer

examples/pytorch/diffpool/model/dgl_layers/aggregator.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Aggregator(nn.Module):
+    """
+    Base Aggregator class. Adapting
+    from PR# 403
+
+    This class is not supposed to be called
+    """
+
+    def __init__(self):
+        super(Aggregator, self).__init__()
+
+    def forward(self, node):
+        neighbour = node.mailbox["m"]
+        c = self.aggre(neighbour)
+        return {"c": c}
+
+    def aggre(self, neighbour):
+        # N x F
+        raise NotImplementedError
+
+
+class MeanAggregator(Aggregator):
+    """
+    Mean Aggregator for graphsage
+    """
+
+    def __init__(self):
+        super(MeanAggregator, self).__init__()
+
+    def aggre(self, neighbour):
+        mean_neighbour = torch.mean(neighbour, dim=1)
+        return mean_neighbour
+
+
+class MaxPoolAggregator(Aggregator):
+    """
+    Maxpooling aggregator for graphsage
+    """
+
+    def __init__(self, in_feats, out_feats, activation, bias):
+        super(MaxPoolAggregator, self).__init__()
+        self.linear = nn.Linear(in_feats, out_feats, bias=bias)
+        self.activation = activation
+        # Xavier initialization of weight
+        nn.init.xavier_uniform_(
+            self.linear.weight, gain=nn.init.calculate_gain("relu")
+        )
+
+    def aggre(self, neighbour):
+        neighbour = self.linear(neighbour)
+        if self.activation:
+            neighbour = self.activation(neighbour)
+        maxpool_neighbour = torch.max(neighbour, dim=1)[0]
+        return maxpool_neighbour
+
+
+class LSTMAggregator(Aggregator):
+    """
+    LSTM aggregator for graphsage
+    """
+
+    def __init__(self, in_feats, hidden_feats):
+        super(LSTMAggregator, self).__init__()
+        self.lstm = nn.LSTM(in_feats, hidden_feats, batch_first=True)
+        self.hidden_dim = hidden_feats
+        self.hidden = self.init_hidden()
+
+        nn.init.xavier_uniform_(
+            self.lstm.weight, gain=nn.init.calculate_gain("relu")
+        )
+
+    def init_hidden(self):
+        """
+        Defaulted to initialite all zero
+        """
+        return (
+            torch.zeros(1, 1, self.hidden_dim),
+            torch.zeros(1, 1, self.hidden_dim),
+        )
+
+    def aggre(self, neighbours):
+        """
+        aggregation function
+        """
+        # N X F
+        rand_order = torch.randperm(neighbours.size()[1])
+        neighbours = neighbours[:, rand_order, :]
+
+        (lstm_out, self.hidden) = self.lstm(
+            neighbours.view(neighbours.size()[0], neighbours.size()[1], -1)
+        )
+        return lstm_out[:, -1, :]
+
+    def forward(self, node):
+        neighbour = node.mailbox["m"]
+        c = self.aggre(neighbour)
+        return {"c": c}

examples/pytorch/diffpool/model/dgl_layers/bundler.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Bundler(nn.Module):
+    """
+    Bundler, which will be the node_apply function in DGL paradigm
+    """
+
+    def __init__(self, in_feats, out_feats, activation, dropout, bias=True):
+        super(Bundler, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        self.linear = nn.Linear(in_feats * 2, out_feats, bias)
+        self.activation = activation
+
+        nn.init.xavier_uniform_(
+            self.linear.weight, gain=nn.init.calculate_gain("relu")
+        )
+
+    def concat(self, h, aggre_result):
+        bundle = torch.cat((h, aggre_result), 1)
+        bundle = self.linear(bundle)
+        return bundle
+
+    def forward(self, node):
+        h = node.data["h"]
+        c = node.data["c"]
+        bundle = self.concat(h, c)
+        bundle = F.normalize(bundle, p=2, dim=1)
+        if self.activation:
+            bundle = self.activation(bundle)
+        return {"h": bundle}