[Example] Add BGRL example (#4077)

* Add BGRL example

* Update README.md

* Update utils.py

* Update

* Update utils.py

examples/README.md

@@ -28,6 +28,9 @@ To quickly locate the examples of your interest, search for the tagged keywords
 - <a name='gatv2'></a> Brody et al. How Attentive are Graph Attention Networks? [Paper link](https://arxiv.org/abs/2105.14491).
     - Example code: [PyTorch](../examples/pytorch/gatv2)
     - Tags: graph attention, gat, gatv2, attention
+- <a name='bgrl'></a> Thakoor et al. Large-Scale Representation Learning on Graphs via Bootstrapping. [Paper link](https://arxiv.org/abs/2102.06514).
+    - Example code: [PyTorch](../examples/pytorch/bgrl)
+    - Tags: contrastive learning for node classification.
 
 ## 2020
 - <a name="eeg-gcnn"></a> Wagh et al. EEG-GCNN: Augmenting Electroencephalogram-based Neurological Disease Diagnosis using a Domain-guided Graph Convolutional Neural Network. [Paper link](http://proceedings.mlr.press/v136/wagh20a.html). 

examples/pytorch/bgrl/README.md

+# 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.
\ No newline at end of file

examples/pytorch/bgrl/eval_function.py

+import torch
+import numpy as np
+from sklearn import metrics
+from sklearn.multiclass import OneVsRestClassifier
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import GridSearchCV, ShuffleSplit, train_test_split
+from sklearn.preprocessing import OneHotEncoder, normalize
+
+
+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

+import os
+import dgl
+import copy
+import torch
+import numpy as np
+from tqdm import tqdm
+from torch.optim import AdamW
+from torch.nn.functional import cosine_similarity
+from utils import get_graph_drop_transform, CosineDecayScheduler, get_dataset
+from model import GCN, GraphSAGE_GCN, MLP_Predictor, BGRL, compute_representations
+from eval_function import fit_logistic_regression, fit_logistic_regression_preset_splits, fit_ppi_linear
+
+
+import warnings
+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.])
+    parser.add_argument('--feat_mask_p', type=float, nargs='+', default=[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

+import dgl
+import copy
+import torch
+from torch import nn
+from torch.nn.init import ones_, zeros_
+from torch.nn import BatchNorm1d, Parameter
+from dgl.nn.pytorch.conv import GraphConv, SAGEConv
+
+
+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. - 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

+import copy
+import torch
+import numpy as np
+from dgl.dataloading import GraphDataLoader
+from dgl.transforms import Compose, DropEdge, FeatMask, RowFeatNormalizer
+from dgl.data import CoauthorCSDataset, CoauthorPhysicsDataset, AmazonCoBuyPhotoDataset, AmazonCoBuyComputerDataset, PPIDataset, WikiCSDataset
+
+
+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.:
+        transforms.append(DropEdge(drop_edge_p))
+
+    # drop features
+    if feat_mask_p > 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.number_of_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
\ No newline at end of file