examples (#5323)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>

examples/mxnet/appnp/appnp.py

 import argparse
 import time
 
+import dgl
+
 import mxnet as mx
 import numpy as np
+from dgl.data import (
+    CiteseerGraphDataset,
+    CoraGraphDataset,
+    PubmedGraphDataset,
+    register_data_args,
+)
+from dgl.nn.mxnet.conv import APPNPConv
 from mxnet import gluon, nd
 from mxnet.gluon import nn
 
-import dgl
-from dgl.data import (CiteseerGraphDataset, CoraGraphDataset,
-                      PubmedGraphDataset, register_data_args)
-from dgl.nn.mxnet.conv import APPNPConv
-
 
 class APPNP(nn.Block):
     def __init__(

examples/mxnet/gat/train.py

@@ -11,17 +11,21 @@ Pytorch implementation: https://github.com/Diego999/pyGAT
 import argparse
 import time
 
+import dgl
+
 import mxnet as mx
 import networkx as nx
 import numpy as np
+from dgl.data import (
+    CiteseerGraphDataset,
+    CoraGraphDataset,
+    PubmedGraphDataset,
+    register_data_args,
+)
 from gat import GAT
 from mxnet import gluon
 from utils import EarlyStopping
 
-import dgl
-from dgl.data import (CiteseerGraphDataset, CoraGraphDataset,
-                      PubmedGraphDataset, register_data_args)
-
 
 def elu(data):
     return mx.nd.LeakyReLU(data, act_type="elu")
@@ -132,7 +136,6 @@ def main(args):
 
 
 if __name__ == "__main__":
-
     parser = argparse.ArgumentParser(description="GAT")
     register_data_args(parser)
     parser.add_argument(

examples/mxnet/gcn/gcn.py

@@ -5,11 +5,10 @@ References:
 - Paper: https://arxiv.org/abs/1609.02907
 - Code: https://github.com/tkipf/gcn
 """
-import mxnet as mx
-from mxnet import gluon
-
 import dgl
+import mxnet as mx
 from dgl.nn.mxnet import GraphConv
+from mxnet import gluon
 
 
 class GCN(gluon.Block):

examples/mxnet/gcn/gcn_concat.py

@@ -5,48 +5,44 @@ Code: https://github.com/tkipf/gcn
 GCN with batch processing
 """
 import argparse
-import numpy as np
 import time
-import mxnet as mx
-from mxnet import gluon
+
 import dgl
 import dgl.function as fn
-from dgl.data import register_data_args
-from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
+import mxnet as mx
+import numpy as np
+from dgl.data import (
+    CiteseerGraphDataset,
+    CoraGraphDataset,
+    PubmedGraphDataset,
+    register_data_args,
+)
+from mxnet import gluon
 
 
 class GCNLayer(gluon.Block):
-    def __init__(self,
-                 g,
-                 out_feats,
-                 activation,
-                 dropout):
+    def __init__(self, g, out_feats, activation, dropout):
         super(GCNLayer, self).__init__()
         self.g = g
         self.dense = gluon.nn.Dense(out_feats, activation)
         self.dropout = dropout
 
     def forward(self, h):
-        self.g.ndata['h'] = h * self.g.ndata['out_norm']
-        self.g.update_all(fn.copy_u(u='h', out='m'),
-                          fn.sum(msg='m', out='accum'))
-        accum = self.g.ndata.pop('accum')
-        accum = self.dense(accum * self.g.ndata['in_norm'])
+        self.g.ndata["h"] = h * self.g.ndata["out_norm"]
+        self.g.update_all(
+            fn.copy_u(u="h", out="m"), fn.sum(msg="m", out="accum")
+        )
+        accum = self.g.ndata.pop("accum")
+        accum = self.dense(accum * self.g.ndata["in_norm"])
         if self.dropout:
             accum = mx.nd.Dropout(accum, p=self.dropout)
-        h = self.g.ndata.pop('h')
-        h = mx.nd.concat(h / self.g.ndata['out_norm'], accum, dim=1)
+        h = self.g.ndata.pop("h")
+        h = mx.nd.concat(h / self.g.ndata["out_norm"], accum, dim=1)
         return h
 
 
 class GCN(gluon.Block):
-    def __init__(self,
-                 g,
-                 n_hidden,
-                 n_classes,
-                 n_layers,
-                 activation,
-                 dropout):
+    def __init__(self, g, n_hidden, n_classes, n_layers, activation, dropout):
         super(GCN, self).__init__()
         self.inp_layer = gluon.nn.Dense(n_hidden, activation)
         self.dropout = dropout
@@ -55,7 +51,6 @@ class GCN(gluon.Block):
             self.layers.add(GCNLayer(g, n_hidden, activation, dropout))
         self.out_layer = gluon.nn.Dense(n_classes)
 
-
     def forward(self, features):
         emb_inp = [features, self.inp_layer(features)]
         if self.dropout:
@@ -75,14 +70,14 @@ def evaluate(model, features, labels, mask):
 
 def main(args):
     # load and preprocess dataset
-    if args.dataset == 'cora':
+    if args.dataset == "cora":
         data = CoraGraphDataset()
-    elif args.dataset == 'citeseer':
+    elif args.dataset == "citeseer":
         data = CiteseerGraphDataset()
-    elif args.dataset == 'pubmed':
+    elif args.dataset == "pubmed":
         data = PubmedGraphDataset()
     else:
-        raise ValueError('Unknown dataset: {}'.format(args.dataset))
+        raise ValueError("Unknown dataset: {}".format(args.dataset))
 
     g = data[0]
     if args.gpu < 0:
@@ -93,45 +88,51 @@ def main(args):
         ctx = mx.gpu(args.gpu)
         g = g.to(ctx)
 
-    features = g.ndata['feat']
-    labels = mx.nd.array(g.ndata['label'], dtype="float32", ctx=ctx)
-    train_mask = g.ndata['train_mask']
-    val_mask = g.ndata['val_mask']
-    test_mask = g.ndata['test_mask']
+    features = g.ndata["feat"]
+    labels = mx.nd.array(g.ndata["label"], dtype="float32", ctx=ctx)
+    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_labels
     n_edges = data.graph.number_of_edges()
-    print("""----Data statistics------'
+    print(
+        """----Data statistics------'
       #Edges %d
       #Classes %d
       #Train samples %d
       #Val samples %d
-      #Test samples %d""" %
-          (n_edges, n_classes,
+      #Test samples %d"""
+        % (
+            n_edges,
+            n_classes,
             train_mask.sum().asscalar(),
             val_mask.sum().asscalar(),
-              test_mask.sum().asscalar()))
+            test_mask.sum().asscalar(),
+        )
+    )
 
     # add self loop
     if args.self_loop:
         g = dgl.remove_self_loop(g)
         g = dgl.add_self_loop(g)
     # normalization
-    in_degs = g.in_degrees().astype('float32')
-    out_degs = g.out_degrees().astype('float32')
+    in_degs = g.in_degrees().astype("float32")
+    out_degs = g.out_degrees().astype("float32")
     in_norm = mx.nd.power(in_degs, -0.5)
     out_norm = mx.nd.power(out_degs, -0.5)
     if cuda:
         in_norm = in_norm.as_in_context(ctx)
         out_norm = out_norm.as_in_context(ctx)
-    g.ndata['in_norm'] = mx.nd.expand_dims(in_norm, 1)
-    g.ndata['out_norm'] = mx.nd.expand_dims(out_norm, 1)
+    g.ndata["in_norm"] = mx.nd.expand_dims(in_norm, 1)
+    g.ndata["out_norm"] = mx.nd.expand_dims(out_norm, 1)
 
-    model = GCN(g,
+    model = GCN(
+        g,
         args.n_hidden,
         n_classes,
         args.n_layers,
-                'relu',
+        "relu",
         args.dropout,
     )
     model.initialize(ctx=ctx)
@@ -140,8 +141,11 @@ def main(args):
 
     # use optimizer
     print(model.collect_params())
-    trainer = gluon.Trainer(model.collect_params(), 'adam',
-            {'learning_rate': args.lr, 'wd': args.weight_decay})
+    trainer = gluon.Trainer(
+        model.collect_params(),
+        "adam",
+        {"learning_rate": args.lr, "wd": args.weight_decay},
+    )
 
     # initialize graph
     dur = []
@@ -160,36 +164,53 @@ def main(args):
         if epoch >= 3:
             dur.append(time.time() - t0)
             acc = evaluate(model, features, labels, val_mask)
-            print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
-                  "ETputs(KTEPS) {:.2f}". format(
-                epoch, np.mean(dur), loss.asscalar(), acc, n_edges / np.mean(dur) / 1000))
+            print(
+                "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
+                "ETputs(KTEPS) {:.2f}".format(
+                    epoch,
+                    np.mean(dur),
+                    loss.asscalar(),
+                    acc,
+                    n_edges / np.mean(dur) / 1000,
+                )
+            )
 
     # test set accuracy
     acc = evaluate(model, features, labels, test_mask)
     print("Test accuracy {:.2%}".format(acc))
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='GCN')
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="GCN")
     register_data_args(parser)
-    parser.add_argument("--dropout", type=float, default=0.5,
-            help="dropout probability")
-    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("--n-hidden", type=int, default=16,
-            help="number of hidden gcn units")
-    parser.add_argument("--n-layers", type=int, default=1,
-            help="number of hidden gcn layers")
-    parser.add_argument("--normalization",
-            choices=['sym','left'], default=None,
-            help="graph normalization types (default=None)")
-    parser.add_argument("--self-loop", action='store_true',
-            help="graph self-loop (default=False)")
-    parser.add_argument("--weight-decay", type=float, default=5e-4,
-            help="Weight for L2 loss")
+    parser.add_argument(
+        "--dropout", type=float, default=0.5, help="dropout probability"
+    )
+    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(
+        "--n-hidden", type=int, default=16, help="number of hidden gcn units"
+    )
+    parser.add_argument(
+        "--n-layers", type=int, default=1, help="number of hidden gcn layers"
+    )
+    parser.add_argument(
+        "--normalization",
+        choices=["sym", "left"],
+        default=None,
+        help="graph normalization types (default=None)",
+    )
+    parser.add_argument(
+        "--self-loop",
+        action="store_true",
+        help="graph self-loop (default=False)",
+    )
+    parser.add_argument(
+        "--weight-decay", type=float, default=5e-4, help="Weight for L2 loss"
+    )
     args = parser.parse_args()
 
     print(args)

examples/mxnet/gcn/train.py

@@ -2,14 +2,14 @@
 import argparse
 import time
 
+import dgl
+
 import mxnet as mx
 import numpy as np
+from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
 from gcn import GCN
 from mxnet import gluon
 
-import dgl
-from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
-
 # from gcn_mp import GCN
 # from gcn_spmv import GCN
 

examples/mxnet/gin/dataloader.py

@@ -4,13 +4,13 @@ MxNet compatible dataloader
 
 import math
 
+import dgl
+
 import numpy as np
 from mxnet import nd
 from mxnet.gluon.data import DataLoader, Sampler
 from sklearn.model_selection import StratifiedKFold
 
-import dgl
-
 
 class SubsetRandomSampler(Sampler):
     def __init__(self, indices):
@@ -52,7 +52,6 @@ class GraphDataLoader:
         fold_idx=0,
         split_ratio=0.7,
     ):
-
         self.shuffle = shuffle
         self.seed = seed
 

examples/mxnet/gin/gin.py

@@ -6,11 +6,11 @@ Author's implementation: https://github.com/weihua916/powerful-gnns
 """
 
 import mxnet as mx
-from mxnet import gluon, nd
-from mxnet.gluon import nn
 
 from dgl.nn.mxnet.conv import GINConv
 from dgl.nn.mxnet.glob import AvgPooling, MaxPooling, SumPooling
+from mxnet import gluon, nd
+from mxnet.gluon import nn
 
 
 class ApplyNodeFunc(nn.Block):

examples/mxnet/gin/main.py

@@ -3,14 +3,14 @@ from parser import Parser
 
 import mxnet as mx
 import numpy as np
-from dataloader import GraphDataLoader, collate
+from dataloader import collate, GraphDataLoader
+
+from dgl.data.gindt import GINDataset
 from gin import GIN
 from mxnet import gluon, nd
 from mxnet.gluon import nn
 from tqdm import tqdm
 
-from dgl.data.gindt import GINDataset
-
 
 def train(args, net, trainloader, trainer, criterion, epoch):
     running_loss = 0
@@ -71,7 +71,6 @@ def eval_net(args, net, dataloader, criterion):
 
 
 def main(args):
-
     # set up seeds, args.seed supported
     mx.random.seed(0)
     np.random.seed(seed=0)

examples/mxnet/graphsage/main.py

@@ -7,17 +7,21 @@ Simple reference implementation of GraphSAGE.
 import argparse
 import time
 
+import dgl
+
 import mxnet as mx
 import networkx as nx
 import numpy as np
+from dgl.data import (
+    CiteseerGraphDataset,
+    CoraGraphDataset,
+    PubmedGraphDataset,
+    register_data_args,
+)
+from dgl.nn.mxnet.conv import SAGEConv
 from mxnet import gluon, nd
 from mxnet.gluon import nn
 
-import dgl
-from dgl.data import (CiteseerGraphDataset, CoraGraphDataset,
-                      PubmedGraphDataset, register_data_args)
-from dgl.nn.mxnet.conv import SAGEConv
-
 
 class GraphSAGE(nn.Block):
     def __init__(

examples/mxnet/monet/citation.py

 import argparse
 import time
 
+import dgl
+
 import mxnet as mx
 import networkx as nx
 import numpy as np
+from dgl.data import (
+    CiteseerGraphDataset,
+    CoraGraphDataset,
+    PubmedGraphDataset,
+    register_data_args,
+)
+from dgl.nn.mxnet.conv import GMMConv
 from mxnet import gluon, nd
 from mxnet.gluon import nn
 
-import dgl
-from dgl.data import (CiteseerGraphDataset, CoraGraphDataset,
-                      PubmedGraphDataset, register_data_args)
-from dgl.nn.mxnet.conv import GMMConv
-
 
 class MoNet(nn.Block):
     def __init__(

examples/mxnet/rgcn/entity_classify.py

@@ -9,43 +9,66 @@ Difference compared to tkipf/relation-gcn
 """
 
 import argparse
-import numpy as np
 import time
+from functools import partial
+
+import dgl
 import mxnet as mx
-from mxnet import gluon
 import mxnet.ndarray as F
-import dgl
+import numpy as np
+from dgl.data.rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
 from dgl.nn.mxnet import RelGraphConv
-from functools import partial
-from dgl.data.rdf import AIFBDataset, MUTAGDataset, BGSDataset, AMDataset
 
 from model import BaseRGCN
+from mxnet import gluon
+
 
 class EntityClassify(BaseRGCN):
     def build_input_layer(self):
-        return RelGraphConv(self.num_nodes, self.h_dim, self.num_rels, "basis",
-                self.num_bases, activation=F.relu, self_loop=self.use_self_loop,
-                dropout=self.dropout)
+        return RelGraphConv(
+            self.num_nodes,
+            self.h_dim,
+            self.num_rels,
+            "basis",
+            self.num_bases,
+            activation=F.relu,
+            self_loop=self.use_self_loop,
+            dropout=self.dropout,
+        )
 
     def build_hidden_layer(self, idx):
-        return RelGraphConv(self.h_dim, self.h_dim, self.num_rels, "basis",
-                self.num_bases, activation=F.relu, self_loop=self.use_self_loop,
-                dropout=self.dropout)
+        return RelGraphConv(
+            self.h_dim,
+            self.h_dim,
+            self.num_rels,
+            "basis",
+            self.num_bases,
+            activation=F.relu,
+            self_loop=self.use_self_loop,
+            dropout=self.dropout,
+        )
 
     def build_output_layer(self):
-        return RelGraphConv(self.h_dim, self.out_dim, self.num_rels, "basis",
-                self.num_bases, activation=None,
-                self_loop=self.use_self_loop)
+        return RelGraphConv(
+            self.h_dim,
+            self.out_dim,
+            self.num_rels,
+            "basis",
+            self.num_bases,
+            activation=None,
+            self_loop=self.use_self_loop,
+        )
+
 
 def main(args):
     # load graph data
-    if args.dataset == 'aifb':
+    if args.dataset == "aifb":
         dataset = AIFBDataset()
-    elif args.dataset == 'mutag':
+    elif args.dataset == "mutag":
         dataset = MUTAGDataset()
-    elif args.dataset == 'bgs':
+    elif args.dataset == "bgs":
         dataset = BGSDataset()
-    elif args.dataset == 'am':
+    elif args.dataset == "am":
         dataset = AMDataset()
     else:
         raise ValueError()
@@ -56,27 +79,31 @@ def main(args):
     num_rels = len(hg.canonical_etypes)
     category = dataset.predict_category
     num_classes = dataset.num_classes
-    train_mask = hg.nodes[category].data.pop('train_mask')
-    test_mask = hg.nodes[category].data.pop('test_mask')
-    train_idx = mx.nd.array(np.nonzero(train_mask.asnumpy())[0], dtype='int64')
-    test_idx = mx.nd.array(np.nonzero(test_mask.asnumpy())[0], dtype='int64')
-    labels = mx.nd.array(hg.nodes[category].data.pop('labels'), dtype='int64')
+    train_mask = hg.nodes[category].data.pop("train_mask")
+    test_mask = hg.nodes[category].data.pop("test_mask")
+    train_idx = mx.nd.array(np.nonzero(train_mask.asnumpy())[0], dtype="int64")
+    test_idx = mx.nd.array(np.nonzero(test_mask.asnumpy())[0], dtype="int64")
+    labels = mx.nd.array(hg.nodes[category].data.pop("labels"), dtype="int64")
 
     # split dataset into train, validate, test
     if args.validation:
-        val_idx = train_idx[:len(train_idx) // 5]
-        train_idx = train_idx[len(train_idx) // 5:]
+        val_idx = train_idx[: len(train_idx) // 5]
+        train_idx = train_idx[len(train_idx) // 5 :]
     else:
         val_idx = train_idx
 
     # calculate norm for each edge type and store in edge
     for canonical_etype in hg.canonical_etypes:
-        u, v, eid = hg.all_edges(form='all', etype=canonical_etype)
+        u, v, eid = hg.all_edges(form="all", etype=canonical_etype)
         v = v.asnumpy()
-        _, inverse_index, count = np.unique(v, return_inverse=True, return_counts=True)
+        _, inverse_index, count = np.unique(
+            v, return_inverse=True, return_counts=True
+        )
         degrees = count[inverse_index]
         norm = np.ones(eid.shape[0]) / degrees
-        hg.edges[canonical_etype].data['norm'] = mx.nd.expand_dims(mx.nd.array(norm), axis=1)
+        hg.edges[canonical_etype].data["norm"] = mx.nd.expand_dims(
+            mx.nd.array(norm), axis=1
+        )
 
     # get target category id
     category_id = len(hg.ntypes)
@@ -84,20 +111,20 @@ def main(args):
         if ntype == category:
             category_id = i
 
-    g = dgl.to_homogeneous(hg, edata=['norm'])
+    g = dgl.to_homogeneous(hg, edata=["norm"])
     num_nodes = g.number_of_nodes()
     node_ids = mx.nd.arange(num_nodes)
-    edge_norm = g.edata['norm']
+    edge_norm = g.edata["norm"]
     edge_type = g.edata[dgl.ETYPE]
 
     # find out the target node ids in g
     node_tids = g.ndata[dgl.NTYPE]
-    loc = (node_tids == category_id)
-    loc = mx.nd.array(np.nonzero(loc.asnumpy())[0], dtype='int64')
+    loc = node_tids == category_id
+    loc = mx.nd.array(np.nonzero(loc.asnumpy())[0], dtype="int64")
     target_idx = node_ids[loc]
 
     # since the nodes are featureless, the input feature is then the node id.
-    feats = mx.nd.arange(num_nodes, dtype='int32')
+    feats = mx.nd.arange(num_nodes, dtype="int32")
 
     # check cuda
     use_cuda = args.gpu >= 0
@@ -113,7 +140,8 @@ def main(args):
         ctx = mx.cpu(0)
 
     # create model
-    model = EntityClassify(num_nodes,
+    model = EntityClassify(
+        num_nodes,
         args.n_hidden,
         num_classes,
         num_rels,
@@ -121,11 +149,16 @@ def main(args):
         num_hidden_layers=args.n_layers - 2,
         dropout=args.dropout,
         use_self_loop=args.use_self_loop,
-                           gpu_id=args.gpu)
+        gpu_id=args.gpu,
+    )
     model.initialize(ctx=ctx)
 
     # optimizer
-    trainer = gluon.Trainer(model.collect_params(), 'adam', {'learning_rate': args.lr, 'wd': args.l2norm})
+    trainer = gluon.Trainer(
+        model.collect_params(),
+        "adam",
+        {"learning_rate": args.lr, "wd": args.l2norm},
+    )
     loss_fcn = gluon.loss.SoftmaxCELoss(from_logits=False)
 
     # training loop
@@ -145,49 +178,88 @@ def main(args):
 
         forward_time.append(t1 - t0)
         backward_time.append(t2 - t1)
-        print("Epoch {:05d} | Train Forward Time(s) {:.4f} | Backward Time(s) {:.4f}".
-              format(epoch, forward_time[-1], backward_time[-1]))
+        print(
+            "Epoch {:05d} | Train Forward Time(s) {:.4f} | Backward Time(s) {:.4f}".format(
+                epoch, forward_time[-1], backward_time[-1]
+            )
+        )
 
-        train_acc = F.sum(mx.nd.cast(pred[train_idx].argmax(axis=1), 'int64') == labels[train_idx]).asscalar() / train_idx.shape[0]
-        val_acc = F.sum(mx.nd.cast(pred[val_idx].argmax(axis=1), 'int64')  == labels[val_idx]).asscalar() / len(val_idx)
-        print("Train Accuracy: {:.4f} | Validation Accuracy: {:.4f}".format(train_acc, val_acc))
+        train_acc = (
+            F.sum(
+                mx.nd.cast(pred[train_idx].argmax(axis=1), "int64")
+                == labels[train_idx]
+            ).asscalar()
+            / train_idx.shape[0]
+        )
+        val_acc = F.sum(
+            mx.nd.cast(pred[val_idx].argmax(axis=1), "int64") == labels[val_idx]
+        ).asscalar() / len(val_idx)
+        print(
+            "Train Accuracy: {:.4f} | Validation Accuracy: {:.4f}".format(
+                train_acc, val_acc
+            )
+        )
     print()
 
     logits = model.forward(g, feats, edge_type, edge_norm)
     logits = logits[target_idx]
-    test_acc = F.sum(mx.nd.cast(logits[test_idx].argmax(axis=1), 'int64')  == labels[test_idx]).asscalar() / len(test_idx)
+    test_acc = F.sum(
+        mx.nd.cast(logits[test_idx].argmax(axis=1), "int64") == labels[test_idx]
+    ).asscalar() / len(test_idx)
     print("Test Accuracy: {:.4f}".format(test_acc))
     print()
 
-    print("Mean forward time: {:4f}".format(np.mean(forward_time[len(forward_time) // 4:])))
-    print("Mean backward time: {:4f}".format(np.mean(backward_time[len(backward_time) // 4:])))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='RGCN')
-    parser.add_argument("--dropout", type=float, default=0,
-            help="dropout probability")
-    parser.add_argument("--n-hidden", type=int, default=16,
-            help="number of hidden units")
-    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-bases", type=int, default=-1,
-            help="number of filter weight matrices, default: -1 [use all]")
-    parser.add_argument("--n-layers", type=int, default=2,
-            help="number of propagation rounds")
-    parser.add_argument("-e", "--n-epochs", type=int, default=50,
-            help="number of training epochs")
-    parser.add_argument("-d", "--dataset", type=str, required=True,
-            help="dataset to use")
-    parser.add_argument("--l2norm", type=float, default=0,
-            help="l2 norm coef")
-    parser.add_argument("--use-self-loop", default=False, action='store_true',
-            help="include self feature as a special relation")
+    print(
+        "Mean forward time: {:4f}".format(
+            np.mean(forward_time[len(forward_time) // 4 :])
+        )
+    )
+    print(
+        "Mean backward time: {:4f}".format(
+            np.mean(backward_time[len(backward_time) // 4 :])
+        )
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="RGCN")
+    parser.add_argument(
+        "--dropout", type=float, default=0, help="dropout probability"
+    )
+    parser.add_argument(
+        "--n-hidden", type=int, default=16, help="number of hidden units"
+    )
+    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-bases",
+        type=int,
+        default=-1,
+        help="number of filter weight matrices, default: -1 [use all]",
+    )
+    parser.add_argument(
+        "--n-layers", type=int, default=2, help="number of propagation rounds"
+    )
+    parser.add_argument(
+        "-e",
+        "--n-epochs",
+        type=int,
+        default=50,
+        help="number of training epochs",
+    )
+    parser.add_argument(
+        "-d", "--dataset", type=str, required=True, help="dataset to use"
+    )
+    parser.add_argument("--l2norm", type=float, default=0, help="l2 norm coef")
+    parser.add_argument(
+        "--use-self-loop",
+        default=False,
+        action="store_true",
+        help="include self feature as a special relation",
+    )
     fp = parser.add_mutually_exclusive_group(required=False)
-    fp.add_argument('--validation', dest='validation', action='store_true')
-    fp.add_argument('--testing', dest='validation', action='store_false')
+    fp.add_argument("--validation", dest="validation", action="store_true")
+    fp.add_argument("--testing", dest="validation", action="store_false")
     parser.set_defaults(validation=True)
 
     args = parser.parse_args()

examples/mxnet/scenegraph/data/dataloader.py

 """DataLoader utils."""
+import dgl
 from gluoncv.data.batchify import Pad
 from mxnet import nd
 
-import dgl
-
 
 def dgl_mp_batchify_fn(data):
     if isinstance(data[0], tuple):

examples/mxnet/scenegraph/data/relation.py

@@ -8,13 +8,15 @@ import pickle
 import warnings
 from collections import Counter
 
+import dgl
+
 import mxnet as mx
 import numpy as np
 from gluoncv.data.base import VisionDataset
 from gluoncv.data.transforms.presets.rcnn import (
-    FasterRCNNDefaultTrainTransform, FasterRCNNDefaultValTransform)
-
-import dgl
+    FasterRCNNDefaultTrainTransform,
+    FasterRCNNDefaultValTransform,
+)
 
 
 class VGRelation(VisionDataset):

examples/mxnet/scenegraph/demo_reldn.py

@@ -5,7 +5,7 @@ import mxnet as mx
 from data import *
 from gluoncv.data.transforms import presets
 from gluoncv.utilz import download
-from model import RelDN, faster_rcnn_resnet101_v1d_custom
+from model import faster_rcnn_resnet101_v1d_custom, RelDN
 from utils import *
 
 import dgl

examples/mxnet/scenegraph/model/faster_rcnn.py

@@ -5,20 +5,26 @@ import os
 import warnings
 
 import mxnet as mx
-from mxnet import autograd
-from mxnet.gluon import nn
-from mxnet.gluon.contrib.nn import SyncBatchNorm
 
-from gluoncv.model_zoo.faster_rcnn.rcnn_target import RCNNTargetSampler, RCNNTargetGenerator
+from gluoncv.model_zoo.faster_rcnn.rcnn_target import (
+    RCNNTargetGenerator,
+    RCNNTargetSampler,
+)
 from gluoncv.model_zoo.rcnn import RCNN
 from gluoncv.model_zoo.rpn import RPN
 from gluoncv.nn.feature import FPNFeatureExpander
+from mxnet import autograd
+from mxnet.gluon import nn
+from mxnet.gluon.contrib.nn import SyncBatchNorm
 
-__all__ = ['FasterRCNN', 'get_faster_rcnn',
-           'faster_rcnn_resnet50_v1b_coco',
-           'faster_rcnn_resnet50_v1b_custom',
-           'faster_rcnn_resnet101_v1d_coco',
-           'faster_rcnn_resnet101_v1d_custom']
+__all__ = [
+    "FasterRCNN",
+    "get_faster_rcnn",
+    "faster_rcnn_resnet50_v1b_coco",
+    "faster_rcnn_resnet50_v1b_custom",
+    "faster_rcnn_resnet101_v1d_coco",
+    "faster_rcnn_resnet101_v1d_custom",
+]
 
 
 class FasterRCNN(RCNN):
@@ -161,22 +167,62 @@ class FasterRCNN(RCNN):
 
     """
 
-    def __init__(self, features, top_features, classes, box_features=None,
-                 short=600, max_size=1000, min_stage=4, max_stage=4, train_patterns=None,
-                 nms_thresh=0.3, nms_topk=400, post_nms=100,
-                 roi_mode='align', roi_size=(14, 14), strides=16, clip=None,
-                 rpn_channel=1024, base_size=16, scales=(8, 16, 32),
-                 ratios=(0.5, 1, 2), alloc_size=(128, 128), rpn_nms_thresh=0.7,
-                 rpn_train_pre_nms=12000, rpn_train_post_nms=2000, rpn_test_pre_nms=6000,
-                 rpn_test_post_nms=300, rpn_min_size=16, per_device_batch_size=1, num_sample=128,
-                 pos_iou_thresh=0.5, pos_ratio=0.25, max_num_gt=300, additional_output=False,
-                 force_nms=False, **kwargs):
+    def __init__(
+        self,
+        features,
+        top_features,
+        classes,
+        box_features=None,
+        short=600,
+        max_size=1000,
+        min_stage=4,
+        max_stage=4,
+        train_patterns=None,
+        nms_thresh=0.3,
+        nms_topk=400,
+        post_nms=100,
+        roi_mode="align",
+        roi_size=(14, 14),
+        strides=16,
+        clip=None,
+        rpn_channel=1024,
+        base_size=16,
+        scales=(8, 16, 32),
+        ratios=(0.5, 1, 2),
+        alloc_size=(128, 128),
+        rpn_nms_thresh=0.7,
+        rpn_train_pre_nms=12000,
+        rpn_train_post_nms=2000,
+        rpn_test_pre_nms=6000,
+        rpn_test_post_nms=300,
+        rpn_min_size=16,
+        per_device_batch_size=1,
+        num_sample=128,
+        pos_iou_thresh=0.5,
+        pos_ratio=0.25,
+        max_num_gt=300,
+        additional_output=False,
+        force_nms=False,
+        **kwargs
+    ):
         super(FasterRCNN, self).__init__(
-            features=features, top_features=top_features, classes=classes,
-            box_features=box_features, short=short, max_size=max_size,
-            train_patterns=train_patterns, nms_thresh=nms_thresh, nms_topk=nms_topk,
-            post_nms=post_nms, roi_mode=roi_mode, roi_size=roi_size, strides=strides, clip=clip,
-            force_nms=force_nms, **kwargs)
+            features=features,
+            top_features=top_features,
+            classes=classes,
+            box_features=box_features,
+            short=short,
+            max_size=max_size,
+            train_patterns=train_patterns,
+            nms_thresh=nms_thresh,
+            nms_topk=nms_topk,
+            post_nms=post_nms,
+            roi_mode=roi_mode,
+            roi_size=roi_size,
+            strides=strides,
+            clip=clip,
+            force_nms=force_nms,
+            **kwargs
+        )
         if rpn_train_post_nms > rpn_train_pre_nms:
             rpn_train_post_nms = rpn_train_pre_nms
         if rpn_test_post_nms > rpn_test_pre_nms:
@@ -187,27 +233,43 @@ class FasterRCNN(RCNN):
         self._max_stage = max_stage
         self.num_stages = max_stage - min_stage + 1
         if self.num_stages > 1:
-            assert len(scales) == len(strides) == self.num_stages, \
-                "The num_stages (%d) must match number of scales (%d) and strides (%d)" \
+            assert len(scales) == len(strides) == self.num_stages, (
+                "The num_stages (%d) must match number of scales (%d) and strides (%d)"
                 % (self.num_stages, len(scales), len(strides))
+            )
         self._batch_size = per_device_batch_size
         self._num_sample = num_sample
         self._rpn_test_post_nms = rpn_test_post_nms
-        self._target_generator = RCNNTargetGenerator(self.num_class, int(num_sample * pos_ratio),
-                                                     self._batch_size)
+        self._target_generator = RCNNTargetGenerator(
+            self.num_class, int(num_sample * pos_ratio), self._batch_size
+        )
         self._additional_output = additional_output
         with self.name_scope():
             self.rpn = RPN(
-                channels=rpn_channel, strides=strides, base_size=base_size,
-                scales=scales, ratios=ratios, alloc_size=alloc_size,
-                clip=clip, nms_thresh=rpn_nms_thresh, train_pre_nms=rpn_train_pre_nms,
-                train_post_nms=rpn_train_post_nms, test_pre_nms=rpn_test_pre_nms,
-                test_post_nms=rpn_test_post_nms, min_size=rpn_min_size,
-                multi_level=self.num_stages > 1, per_level_nms=False)
-            self.sampler = RCNNTargetSampler(num_image=self._batch_size,
-                                             num_proposal=rpn_train_post_nms, num_sample=num_sample,
-                                             pos_iou_thresh=pos_iou_thresh, pos_ratio=pos_ratio,
-                                             max_num_gt=max_num_gt)
+                channels=rpn_channel,
+                strides=strides,
+                base_size=base_size,
+                scales=scales,
+                ratios=ratios,
+                alloc_size=alloc_size,
+                clip=clip,
+                nms_thresh=rpn_nms_thresh,
+                train_pre_nms=rpn_train_pre_nms,
+                train_post_nms=rpn_train_post_nms,
+                test_pre_nms=rpn_test_pre_nms,
+                test_post_nms=rpn_test_post_nms,
+                min_size=rpn_min_size,
+                multi_level=self.num_stages > 1,
+                per_level_nms=False,
+            )
+            self.sampler = RCNNTargetSampler(
+                num_image=self._batch_size,
+                num_proposal=rpn_train_post_nms,
+                num_sample=num_sample,
+                pos_iou_thresh=pos_iou_thresh,
+                pos_ratio=pos_ratio,
+                max_num_gt=max_num_gt,
+            )
 
     @property
     def target_generator(self):
@@ -248,11 +310,21 @@ class FasterRCNN(RCNN):
 
         """
         super(FasterRCNN, self).reset_class(classes, reuse_weights)
-        self._target_generator = RCNNTargetGenerator(self.num_class, self.sampler._max_pos,
-                                                     self._batch_size)
+        self._target_generator = RCNNTargetGenerator(
+            self.num_class, self.sampler._max_pos, self._batch_size
+        )
 
-    def _pyramid_roi_feats(self, F, features, rpn_rois, roi_size, strides, roi_mode='align',
-                           roi_canonical_scale=224.0, eps=1e-6):
+    def _pyramid_roi_feats(
+        self,
+        F,
+        features,
+        rpn_rois,
+        roi_size,
+        strides,
+        roi_mode="align",
+        roi_canonical_scale=224.0,
+        eps=1e-6,
+    ):
         """Assign rpn_rois to specific FPN layers according to its area
            and then perform `ROIPooling` or `ROIAlign` to generate final
            region proposals aggregated features.
@@ -281,7 +353,9 @@ class FasterRCNN(RCNN):
         _, x1, y1, x2, y2 = F.split(rpn_rois, axis=-1, num_outputs=5)
         h = y2 - y1 + 1
         w = x2 - x1 + 1
-        roi_level = F.floor(4 + F.log2(F.sqrt(w * h) / roi_canonical_scale + eps))
+        roi_level = F.floor(
+            4 + F.log2(F.sqrt(w * h) / roi_canonical_scale + eps)
+        )
         roi_level = F.squeeze(F.clip(roi_level, self._min_stage, max_stage))
         # [2,2,..,3,3,...,4,4,...,5,5,...] ``Prohibit swap order here``
         # roi_level_sorted_args = F.argsort(roi_level, is_ascend=True)
@@ -289,23 +363,44 @@ class FasterRCNN(RCNN):
         # rpn_rois = F.take(rpn_rois, roi_level_sorted_args, axis=0)
         pooled_roi_feats = []
         for i, l in enumerate(range(self._min_stage, max_stage + 1)):
-            if roi_mode == 'pool':
+            if roi_mode == "pool":
                 # Pool features with all rois first, and then set invalid pooled features to zero,
                 # at last ele-wise add together to aggregate all features.
-                pooled_feature = F.ROIPooling(features[i], rpn_rois, roi_size, 1. / strides[i])
-                pooled_feature = F.where(roi_level == l, pooled_feature,
-                                         F.zeros_like(pooled_feature))
-            elif roi_mode == 'align':
-                if 'box_encode' in F.contrib.__dict__ and 'box_decode' in F.contrib.__dict__:
+                pooled_feature = F.ROIPooling(
+                    features[i], rpn_rois, roi_size, 1.0 / strides[i]
+                )
+                pooled_feature = F.where(
+                    roi_level == l, pooled_feature, F.zeros_like(pooled_feature)
+                )
+            elif roi_mode == "align":
+                if (
+                    "box_encode" in F.contrib.__dict__
+                    and "box_decode" in F.contrib.__dict__
+                ):
                     # TODO(jerryzcn): clean this up for once mx 1.6 is released.
-                    masked_rpn_rois = F.where(roi_level == l, rpn_rois, F.ones_like(rpn_rois) * -1.)
-                    pooled_feature = F.contrib.ROIAlign(features[i], masked_rpn_rois, roi_size,
-                                                        1. / strides[i], sample_ratio=2)
+                    masked_rpn_rois = F.where(
+                        roi_level == l, rpn_rois, F.ones_like(rpn_rois) * -1.0
+                    )
+                    pooled_feature = F.contrib.ROIAlign(
+                        features[i],
+                        masked_rpn_rois,
+                        roi_size,
+                        1.0 / strides[i],
+                        sample_ratio=2,
+                    )
                 else:
-                    pooled_feature = F.contrib.ROIAlign(features[i], rpn_rois, roi_size,
-                                                        1. / strides[i], sample_ratio=2)
-                    pooled_feature = F.where(roi_level == l, pooled_feature,
-                                             F.zeros_like(pooled_feature))
+                    pooled_feature = F.contrib.ROIAlign(
+                        features[i],
+                        rpn_rois,
+                        roi_size,
+                        1.0 / strides[i],
+                        sample_ratio=2,
+                    )
+                    pooled_feature = F.where(
+                        roi_level == l,
+                        pooled_feature,
+                        F.zeros_like(pooled_feature),
+                    )
             else:
                 raise ValueError("Invalid roi mode: {}".format(roi_mode))
             pooled_roi_feats.append(pooled_feature)
@@ -341,7 +436,9 @@ class FasterRCNN(RCNN):
         """
 
         def _split(x, axis, num_outputs, squeeze_axis):
-            x = F.split(x, axis=axis, num_outputs=num_outputs, squeeze_axis=squeeze_axis)
+            x = F.split(
+                x, axis=axis, num_outputs=num_outputs, squeeze_axis=squeeze_axis
+            )
             if isinstance(x, list):
                 return x
             else:
@@ -361,9 +458,16 @@ class FasterRCNN(RCNN):
                 rpn_box = m_rpn_box
                 self.nms_thresh = 1
             else:
-                rpn_score, rpn_box, raw_rpn_score, raw_rpn_box, anchors = \
-                    self.rpn(F.zeros_like(x), *feat)
-                rpn_box, samples, matches = self.sampler(rpn_box, rpn_score, gt_box)
+                (
+                    rpn_score,
+                    rpn_box,
+                    raw_rpn_score,
+                    raw_rpn_box,
+                    anchors,
+                ) = self.rpn(F.zeros_like(x), *feat)
+                rpn_box, samples, matches = self.sampler(
+                    rpn_box, rpn_score, gt_box
+                )
         else:
             if manual_rpn_box:
                 rpn_box = m_rpn_box
@@ -373,7 +477,11 @@ class FasterRCNN(RCNN):
 
         # create batchid for roi
         if not manual_rpn_box:
-            num_roi = self._num_sample if autograd.is_training() else self._rpn_test_post_nms
+            num_roi = (
+                self._num_sample
+                if autograd.is_training()
+                else self._rpn_test_post_nms
+            )
             batch_size = self._batch_size if autograd.is_training() else 1
         else:
             num_roi = m_rpn_box.shape[1]
@@ -383,20 +491,36 @@ class FasterRCNN(RCNN):
             roi_batchid = F.arange(0, batch_size)
             roi_batchid = F.repeat(roi_batchid, num_roi)
             # remove batch dim because ROIPooling require 2d input
-            rpn_roi = F.concat(*[roi_batchid.reshape((-1, 1)), rpn_box.reshape((-1, 4))], dim=-1)
+            rpn_roi = F.concat(
+                *[roi_batchid.reshape((-1, 1)), rpn_box.reshape((-1, 4))],
+                dim=-1
+            )
             rpn_roi = F.stop_gradient(rpn_roi)
 
         if self.num_stages > 1:
             # using FPN
-            pooled_feat = self._pyramid_roi_feats(F, feat, rpn_roi, self._roi_size,
-                                                  self._strides, roi_mode=self._roi_mode)
+            pooled_feat = self._pyramid_roi_feats(
+                F,
+                feat,
+                rpn_roi,
+                self._roi_size,
+                self._strides,
+                roi_mode=self._roi_mode,
+            )
         else:
             # ROI features
-            if self._roi_mode == 'pool':
-                pooled_feat = F.ROIPooling(feat[0], rpn_roi, self._roi_size, 1. / self._strides)
-            elif self._roi_mode == 'align':
-                pooled_feat = F.contrib.ROIAlign(feat[0], rpn_roi, self._roi_size,
-                                                 1. / self._strides, sample_ratio=2)
+            if self._roi_mode == "pool":
+                pooled_feat = F.ROIPooling(
+                    feat[0], rpn_roi, self._roi_size, 1.0 / self._strides
+                )
+            elif self._roi_mode == "align":
+                pooled_feat = F.contrib.ROIAlign(
+                    feat[0],
+                    rpn_roi,
+                    self._roi_size,
+                    1.0 / self._strides,
+                    sample_ratio=2,
+                )
             else:
                 raise ValueError("Invalid roi mode: {}".format(self._roi_mode))
 
@@ -413,30 +537,79 @@ class FasterRCNN(RCNN):
         # cls_pred (B * N, C) -> (B, N, C)
         cls_pred = cls_pred.reshape((batch_size, num_roi, self.num_class + 1))
         if manual_rpn_box:
-            spatial_feat = top_feat.mean(axis=1).reshape((-4, rpn_box.shape[0], rpn_box.shape[1], -3))
+            spatial_feat = top_feat.mean(axis=1).reshape(
+                (-4, rpn_box.shape[0], rpn_box.shape[1], -3)
+            )
             cls_ids, scores = self.cls_decoder(F.softmax(cls_pred, axis=-1))
             cls_ids = cls_ids.transpose((0, 2, 1)).reshape((0, 0, 0, 1))
             scores = scores.transpose((0, 2, 1)).reshape((0, 0, 0, 1))
-            cls_ids = _split(cls_ids, axis=0, num_outputs=batch_size, squeeze_axis=True)
-            scores = _split(scores, axis=0, num_outputs=batch_size, squeeze_axis=True)
+            cls_ids = _split(
+                cls_ids, axis=0, num_outputs=batch_size, squeeze_axis=True
+            )
+            scores = _split(
+                scores, axis=0, num_outputs=batch_size, squeeze_axis=True
+            )
             return cls_ids, scores, rpn_box, spatial_feat
 
         # no need to convert bounding boxes in training, just return
         if autograd.is_training():
-            cls_targets, box_targets, box_masks, indices = \
-                self._target_generator(rpn_box, samples, matches, gt_label, gt_box)
+            (
+                cls_targets,
+                box_targets,
+                box_masks,
+                indices,
+            ) = self._target_generator(
+                rpn_box, samples, matches, gt_label, gt_box
+            )
             box_feat = F.reshape(box_feat.expand_dims(0), (batch_size, -1, 0))
-            box_pred = self.box_predictor(F.concat(
-                *[F.take(F.slice_axis(box_feat, axis=0, begin=i, end=i + 1).squeeze(),
-                         F.slice_axis(indices, axis=0, begin=i, end=i + 1).squeeze())
-                  for i in range(batch_size)], dim=0))
+            box_pred = self.box_predictor(
+                F.concat(
+                    *[
+                        F.take(
+                            F.slice_axis(
+                                box_feat, axis=0, begin=i, end=i + 1
+                            ).squeeze(),
+                            F.slice_axis(
+                                indices, axis=0, begin=i, end=i + 1
+                            ).squeeze(),
+                        )
+                        for i in range(batch_size)
+                    ],
+                    dim=0
+                )
+            )
             # box_pred (B * N, C * 4) -> (B, N, C, 4)
             box_pred = box_pred.reshape((batch_size, -1, self.num_class, 4))
             if self._additional_output:
-                return (cls_pred, box_pred, rpn_box, samples, matches, raw_rpn_score, raw_rpn_box,
-                        anchors, cls_targets, box_targets, box_masks, top_feat, indices)
-            return (cls_pred, box_pred, rpn_box, samples, matches, raw_rpn_score, raw_rpn_box,
-                    anchors, cls_targets, box_targets, box_masks, indices)
+                return (
+                    cls_pred,
+                    box_pred,
+                    rpn_box,
+                    samples,
+                    matches,
+                    raw_rpn_score,
+                    raw_rpn_box,
+                    anchors,
+                    cls_targets,
+                    box_targets,
+                    box_masks,
+                    top_feat,
+                    indices,
+                )
+            return (
+                cls_pred,
+                box_pred,
+                rpn_box,
+                samples,
+                matches,
+                raw_rpn_score,
+                raw_rpn_box,
+                anchors,
+                cls_targets,
+                box_targets,
+                box_masks,
+                indices,
+            )
 
         box_pred = self.box_predictor(box_feat)
         # box_pred (B * N, C * 4) -> (B, N, C, 4)
@@ -450,12 +623,20 @@ class FasterRCNN(RCNN):
         box_pred = box_pred.transpose((0, 2, 1, 3))
 
         # rpn_boxes (B, N, 4) -> B * (1, N, 4)
-        rpn_boxes = _split(rpn_box, axis=0, num_outputs=batch_size, squeeze_axis=False)
+        rpn_boxes = _split(
+            rpn_box, axis=0, num_outputs=batch_size, squeeze_axis=False
+        )
         # cls_ids, scores (B, C, N, 1) -> B * (C, N, 1)
-        cls_ids = _split(cls_ids, axis=0, num_outputs=batch_size, squeeze_axis=True)
-        scores = _split(scores, axis=0, num_outputs=batch_size, squeeze_axis=True)
+        cls_ids = _split(
+            cls_ids, axis=0, num_outputs=batch_size, squeeze_axis=True
+        )
+        scores = _split(
+            scores, axis=0, num_outputs=batch_size, squeeze_axis=True
+        )
         # box_preds (B, C, N, 4) -> B * (C, N, 4)
-        box_preds = _split(box_pred, axis=0, num_outputs=batch_size, squeeze_axis=True)
+        box_preds = _split(
+            box_pred, axis=0, num_outputs=batch_size, squeeze_axis=True
+        )
 
         # per batch predict, nms, each class has topk outputs
         results = []
@@ -465,8 +646,14 @@ class FasterRCNN(RCNN):
             # ind = mx.nd.array(list(range(sizes[1])))
             ind = mx.nd.linspace(0, 999, 1000)
             ind = mx.nd.repeat(ind, repeats=sizes[0])
-            ind = ind.reshape(sizes[1], sizes[0]).transpose((1, 0)).expand_dims(axis=2)
-        for rpn_box, cls_id, score, box_pred in zip(rpn_boxes, cls_ids, scores, box_preds):
+            ind = (
+                ind.reshape(sizes[1], sizes[0])
+                .transpose((1, 0))
+                .expand_dims(axis=2)
+            )
+        for rpn_box, cls_id, score, box_pred in zip(
+            rpn_boxes, cls_ids, scores, box_preds
+        ):
             # box_pred (C, N, 4) rpn_box (1, N, 4) -> bbox (C, N, 4)
             bbox = self.box_decoder(box_pred, rpn_box)
             if self._additional_output:
@@ -480,8 +667,15 @@ class FasterRCNN(RCNN):
                 res = res.reshape((1, -1, 0))
             # res (C, self.nms_topk, 6)
             res = F.contrib.box_nms(
-                res, overlap_thresh=self.nms_thresh, topk=self.nms_topk, valid_thresh=0.001,
-                id_index=0, score_index=1, coord_start=2, force_suppress=self.force_nms)
+                res,
+                overlap_thresh=self.nms_thresh,
+                topk=self.nms_topk,
+                valid_thresh=0.001,
+                id_index=0,
+                score_index=1,
+                coord_start=2,
+                force_suppress=self.force_nms,
+            )
             # res (C * self.nms_topk, 6)
             res = res.reshape((-3, 0))
             results.append(res)
@@ -493,12 +687,21 @@ class FasterRCNN(RCNN):
         bboxes = F.slice_axis(result, axis=-1, begin=2, end=6)
         if self._additional_output:
             feat_ind = F.slice_axis(result, axis=-1, begin=6, end=7)
-            spatial_feat = top_feat.mean(axis=1).expand_dims(0).reshape(batch_size, 0, -1)
+            spatial_feat = (
+                top_feat.mean(axis=1).expand_dims(0).reshape(batch_size, 0, -1)
+            )
             return ids, scores, bboxes, feat, feat_ind, spatial_feat
         return ids, scores, bboxes
 
-def get_faster_rcnn(name, dataset, pretrained=False, ctx=mx.cpu(),
-                    root=os.path.join('~', '.mxnet', 'models'), **kwargs):
+
+def get_faster_rcnn(
+    name,
+    dataset,
+    pretrained=False,
+    ctx=mx.cpu(),
+    root=os.path.join("~", ".mxnet", "models"),
+    **kwargs
+):
     r"""Utility function to return faster rcnn networks.
 
     Parameters
@@ -524,9 +727,14 @@ def get_faster_rcnn(name, dataset, pretrained=False, ctx=mx.cpu(),
     net = FasterRCNN(**kwargs)
     if pretrained:
         from gluoncv.model_zoo.model_store import get_model_file
-        full_name = '_'.join(('faster_rcnn', name, dataset))
-        net.load_parameters(get_model_file(full_name, tag=pretrained, root=root), ctx=ctx,
-                            ignore_extra=True, allow_missing=True)
+
+        full_name = "_".join(("faster_rcnn", name, dataset))
+        net.load_parameters(
+            get_model_file(full_name, tag=pretrained, root=root),
+            ctx=ctx,
+            ignore_extra=True,
+            allow_missing=True,
+        )
     else:
         for v in net.collect_params().values():
             try:
@@ -535,7 +743,10 @@ def get_faster_rcnn(name, dataset, pretrained=False, ctx=mx.cpu(),
                 pass
     return net
 
-def faster_rcnn_resnet50_v1b_coco(pretrained=False, pretrained_base=True, **kwargs):
+
+def faster_rcnn_resnet50_v1b_coco(
+    pretrained=False, pretrained_base=True, **kwargs
+):
     r"""Faster RCNN model from the paper
     "Ren, S., He, K., Girshick, R., & Sun, J. (2015). Faster r-cnn: Towards
     real-time object detection with region proposal networks"
@@ -558,34 +769,73 @@ def faster_rcnn_resnet50_v1b_coco(pretrained=False, pretrained_base=True, **kwar
     >>> model = get_faster_rcnn_resnet50_v1b_coco(pretrained=True)
     >>> print(model)
     """
-    from gluoncv.model_zoo.resnetv1b import resnet50_v1b
     from gluoncv.data import COCODetection
+    from gluoncv.model_zoo.resnetv1b import resnet50_v1b
+
     classes = COCODetection.CLASSES
     pretrained_base = False if pretrained else pretrained_base
-    base_network = resnet50_v1b(pretrained=pretrained_base, dilated=False,
-                                use_global_stats=True, **kwargs)
+    base_network = resnet50_v1b(
+        pretrained=pretrained_base,
+        dilated=False,
+        use_global_stats=True,
+        **kwargs
+    )
     features = nn.HybridSequential()
     top_features = nn.HybridSequential()
-    for layer in ['conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', 'layer3']:
+    for layer in [
+        "conv1",
+        "bn1",
+        "relu",
+        "maxpool",
+        "layer1",
+        "layer2",
+        "layer3",
+    ]:
         features.add(getattr(base_network, layer))
-    for layer in ['layer4']:
+    for layer in ["layer4"]:
         top_features.add(getattr(base_network, layer))
-    train_patterns = '|'.join(['.*dense', '.*rpn', '.*down(2|3|4)_conv', '.*layers(2|3|4)_conv'])
+    train_patterns = "|".join(
+        [".*dense", ".*rpn", ".*down(2|3|4)_conv", ".*layers(2|3|4)_conv"]
+    )
     return get_faster_rcnn(
-        name='resnet50_v1b', dataset='coco', pretrained=pretrained,
-        features=features, top_features=top_features, classes=classes,
-        short=800, max_size=1333, train_patterns=train_patterns,
-        nms_thresh=0.7, nms_topk=-1, post_nms=-1,
-        roi_mode='align', roi_size=(14, 14), strides=16, clip=4.14,
-        rpn_channel=1024, base_size=16, scales=(2, 4, 8, 16, 32),
-        ratios=(0.5, 1, 2), alloc_size=(128, 128), rpn_nms_thresh=0.7,
-        rpn_train_pre_nms=12000, rpn_train_post_nms=2000,
-        rpn_test_pre_nms=6000, rpn_test_post_nms=1000, rpn_min_size=1,
-        num_sample=128, pos_iou_thresh=0.5, pos_ratio=0.25,
-        max_num_gt=3000, **kwargs)
-
-def faster_rcnn_resnet50_v1b_custom(classes, transfer=None, pretrained_base=True,
-                                    pretrained=False, **kwargs):
+        name="resnet50_v1b",
+        dataset="coco",
+        pretrained=pretrained,
+        features=features,
+        top_features=top_features,
+        classes=classes,
+        short=800,
+        max_size=1333,
+        train_patterns=train_patterns,
+        nms_thresh=0.7,
+        nms_topk=-1,
+        post_nms=-1,
+        roi_mode="align",
+        roi_size=(14, 14),
+        strides=16,
+        clip=4.14,
+        rpn_channel=1024,
+        base_size=16,
+        scales=(2, 4, 8, 16, 32),
+        ratios=(0.5, 1, 2),
+        alloc_size=(128, 128),
+        rpn_nms_thresh=0.7,
+        rpn_train_pre_nms=12000,
+        rpn_train_post_nms=2000,
+        rpn_test_pre_nms=6000,
+        rpn_test_post_nms=1000,
+        rpn_min_size=1,
+        num_sample=128,
+        pos_iou_thresh=0.5,
+        pos_ratio=0.25,
+        max_num_gt=3000,
+        **kwargs
+    )
+
+
+def faster_rcnn_resnet50_v1b_custom(
+    classes, transfer=None, pretrained_base=True, pretrained=False, **kwargs
+):
     r"""Faster RCNN model with resnet50_v1b base network on custom dataset.
 
     Parameters
@@ -612,39 +862,85 @@ def faster_rcnn_resnet50_v1b_custom(classes, transfer=None, pretrained_base=True
         Hybrid faster RCNN network.
     """
     if pretrained:
-        warnings.warn("Custom models don't provide `pretrained` weights, ignored.")
+        warnings.warn(
+            "Custom models don't provide `pretrained` weights, ignored."
+        )
     if transfer is None:
         from gluoncv.model_zoo.resnetv1b import resnet50_v1b
-        base_network = resnet50_v1b(pretrained=pretrained_base, dilated=False,
-                                    use_global_stats=True, **kwargs)
+
+        base_network = resnet50_v1b(
+            pretrained=pretrained_base,
+            dilated=False,
+            use_global_stats=True,
+            **kwargs
+        )
         features = nn.HybridSequential()
         top_features = nn.HybridSequential()
-        for layer in ['conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', 'layer3']:
+        for layer in [
+            "conv1",
+            "bn1",
+            "relu",
+            "maxpool",
+            "layer1",
+            "layer2",
+            "layer3",
+        ]:
             features.add(getattr(base_network, layer))
-        for layer in ['layer4']:
+        for layer in ["layer4"]:
             top_features.add(getattr(base_network, layer))
-        train_patterns = '|'.join(['.*dense', '.*rpn', '.*down(2|3|4)_conv',
-                                   '.*layers(2|3|4)_conv'])
+        train_patterns = "|".join(
+            [".*dense", ".*rpn", ".*down(2|3|4)_conv", ".*layers(2|3|4)_conv"]
+        )
         return get_faster_rcnn(
-            name='resnet50_v1b', dataset='custom', pretrained=pretrained,
-            features=features, top_features=top_features, classes=classes,
-            short=600, max_size=1000, train_patterns=train_patterns,
-            nms_thresh=0.7, nms_topk=400, post_nms=100,
-            roi_mode='align', roi_size=(14, 14), strides=16, clip=4.14,
-            rpn_channel=1024, base_size=16, scales=(2, 4, 8, 16, 32),
-            ratios=(0.5, 1, 2), alloc_size=(128, 128), rpn_nms_thresh=0.7,
-            rpn_train_pre_nms=12000, rpn_train_post_nms=2000,
-            rpn_test_pre_nms=6000, rpn_test_post_nms=300, rpn_min_size=16,
-            num_sample=128, pos_iou_thresh=0.5, pos_ratio=0.25, max_num_gt=3000,
-            **kwargs)
+            name="resnet50_v1b",
+            dataset="custom",
+            pretrained=pretrained,
+            features=features,
+            top_features=top_features,
+            classes=classes,
+            short=600,
+            max_size=1000,
+            train_patterns=train_patterns,
+            nms_thresh=0.7,
+            nms_topk=400,
+            post_nms=100,
+            roi_mode="align",
+            roi_size=(14, 14),
+            strides=16,
+            clip=4.14,
+            rpn_channel=1024,
+            base_size=16,
+            scales=(2, 4, 8, 16, 32),
+            ratios=(0.5, 1, 2),
+            alloc_size=(128, 128),
+            rpn_nms_thresh=0.7,
+            rpn_train_pre_nms=12000,
+            rpn_train_post_nms=2000,
+            rpn_test_pre_nms=6000,
+            rpn_test_post_nms=300,
+            rpn_min_size=16,
+            num_sample=128,
+            pos_iou_thresh=0.5,
+            pos_ratio=0.25,
+            max_num_gt=3000,
+            **kwargs
+        )
     else:
         from gluoncv.model_zoo import get_model
-        net = get_model('faster_rcnn_resnet50_v1b_' + str(transfer), pretrained=True, **kwargs)
+
+        net = get_model(
+            "faster_rcnn_resnet50_v1b_" + str(transfer),
+            pretrained=True,
+            **kwargs
+        )
         reuse_classes = [x for x in classes if x in net.classes]
         net.reset_class(classes, reuse_weights=reuse_classes)
     return net
 
-def faster_rcnn_resnet101_v1d_coco(pretrained=False, pretrained_base=True, **kwargs):
+
+def faster_rcnn_resnet101_v1d_coco(
+    pretrained=False, pretrained_base=True, **kwargs
+):
     r"""Faster RCNN model from the paper
     "Ren, S., He, K., Girshick, R., & Sun, J. (2015). Faster r-cnn: Towards
     real-time object detection with region proposal networks"
@@ -666,34 +962,73 @@ def faster_rcnn_resnet101_v1d_coco(pretrained=False, pretrained_base=True, **kwa
     >>> model = get_faster_rcnn_resnet101_v1d_coco(pretrained=True)
     >>> print(model)
     """
-    from gluoncv.model_zoo.resnetv1b import resnet101_v1d
     from gluoncv.data import COCODetection
+    from gluoncv.model_zoo.resnetv1b import resnet101_v1d
+
     classes = COCODetection.CLASSES
     pretrained_base = False if pretrained else pretrained_base
-    base_network = resnet101_v1d(pretrained=pretrained_base, dilated=False,
-                                 use_global_stats=True, **kwargs)
+    base_network = resnet101_v1d(
+        pretrained=pretrained_base,
+        dilated=False,
+        use_global_stats=True,
+        **kwargs
+    )
     features = nn.HybridSequential()
     top_features = nn.HybridSequential()
-    for layer in ['conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', 'layer3']:
+    for layer in [
+        "conv1",
+        "bn1",
+        "relu",
+        "maxpool",
+        "layer1",
+        "layer2",
+        "layer3",
+    ]:
         features.add(getattr(base_network, layer))
-    for layer in ['layer4']:
+    for layer in ["layer4"]:
         top_features.add(getattr(base_network, layer))
-    train_patterns = '|'.join(['.*dense', '.*rpn', '.*down(2|3|4)_conv', '.*layers(2|3|4)_conv'])
+    train_patterns = "|".join(
+        [".*dense", ".*rpn", ".*down(2|3|4)_conv", ".*layers(2|3|4)_conv"]
+    )
     return get_faster_rcnn(
-        name='resnet101_v1d', dataset='coco', pretrained=pretrained,
-        features=features, top_features=top_features, classes=classes,
-        short=800, max_size=1333, train_patterns=train_patterns,
-        nms_thresh=0.5, nms_topk=-1, post_nms=100,
-        roi_mode='align', roi_size=(14, 14), strides=16, clip=4.14,
-        rpn_channel=1024, base_size=16, scales=(2, 4, 8, 16, 32),
-        ratios=(0.5, 1, 2), alloc_size=(128, 128), rpn_nms_thresh=0.7,
-        rpn_train_pre_nms=12000, rpn_train_post_nms=2000,
-        rpn_test_pre_nms=6000, rpn_test_post_nms=1000, rpn_min_size=1,
-        num_sample=128, pos_iou_thresh=0.5, pos_ratio=0.25, max_num_gt=3000,
-        **kwargs)
-
-def faster_rcnn_resnet101_v1d_custom(classes, transfer=None, pretrained_base=True,
-                                     pretrained=False, **kwargs):
+        name="resnet101_v1d",
+        dataset="coco",
+        pretrained=pretrained,
+        features=features,
+        top_features=top_features,
+        classes=classes,
+        short=800,
+        max_size=1333,
+        train_patterns=train_patterns,
+        nms_thresh=0.5,
+        nms_topk=-1,
+        post_nms=100,
+        roi_mode="align",
+        roi_size=(14, 14),
+        strides=16,
+        clip=4.14,
+        rpn_channel=1024,
+        base_size=16,
+        scales=(2, 4, 8, 16, 32),
+        ratios=(0.5, 1, 2),
+        alloc_size=(128, 128),
+        rpn_nms_thresh=0.7,
+        rpn_train_pre_nms=12000,
+        rpn_train_post_nms=2000,
+        rpn_test_pre_nms=6000,
+        rpn_test_post_nms=1000,
+        rpn_min_size=1,
+        num_sample=128,
+        pos_iou_thresh=0.5,
+        pos_ratio=0.25,
+        max_num_gt=3000,
+        **kwargs
+    )
+
+
+def faster_rcnn_resnet101_v1d_custom(
+    classes, transfer=None, pretrained_base=True, pretrained=False, **kwargs
+):
     r"""Faster RCNN model with resnet101_v1d base network on custom dataset.
 
     Parameters
@@ -717,31 +1052,69 @@ def faster_rcnn_resnet101_v1d_custom(classes, transfer=None, pretrained_base=Tru
         Hybrid faster RCNN network.
     """
     if pretrained:
-        warnings.warn("Custom models don't provide `pretrained` weights, ignored.")
+        warnings.warn(
+            "Custom models don't provide `pretrained` weights, ignored."
+        )
     if transfer is None:
         from gluoncv.model_zoo.resnetv1b import resnet101_v1d
-        base_network = resnet101_v1d(pretrained=pretrained_base, dilated=False,
-                                     use_global_stats=True, **kwargs)
+
+        base_network = resnet101_v1d(
+            pretrained=pretrained_base,
+            dilated=False,
+            use_global_stats=True,
+            **kwargs
+        )
         features = nn.HybridSequential()
         top_features = nn.HybridSequential()
-        for layer in ['conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', 'layer3']:
+        for layer in [
+            "conv1",
+            "bn1",
+            "relu",
+            "maxpool",
+            "layer1",
+            "layer2",
+            "layer3",
+        ]:
             features.add(getattr(base_network, layer))
-        for layer in ['layer4']:
+        for layer in ["layer4"]:
             top_features.add(getattr(base_network, layer))
-        train_patterns = '|'.join(['.*dense', '.*rpn', '.*down(2|3|4)_conv',
-                                   '.*layers(2|3|4)_conv'])
+        train_patterns = "|".join(
+            [".*dense", ".*rpn", ".*down(2|3|4)_conv", ".*layers(2|3|4)_conv"]
+        )
         return get_faster_rcnn(
-            name='resnet101_v1d', dataset='custom', pretrained=pretrained,
-            features=features, top_features=top_features, classes=classes,
-            short=600, max_size=1000, train_patterns=train_patterns,
-            nms_thresh=0.5, nms_topk=400, post_nms=100,
-            roi_mode='align', roi_size=(14, 14), strides=16, clip=4.14,
-            rpn_channel=1024, base_size=16, scales=(2, 4, 8, 16, 32),
-            ratios=(0.5, 1, 2), alloc_size=(128, 128), rpn_nms_thresh=0.7,
-            rpn_train_pre_nms=12000, rpn_train_post_nms=2000,
-            rpn_test_pre_nms=6000, rpn_test_post_nms=300, rpn_min_size=16,
-            num_sample=128, pos_iou_thresh=0.5, pos_ratio=0.25, max_num_gt=3000,
-            **kwargs)
+            name="resnet101_v1d",
+            dataset="custom",
+            pretrained=pretrained,
+            features=features,
+            top_features=top_features,
+            classes=classes,
+            short=600,
+            max_size=1000,
+            train_patterns=train_patterns,
+            nms_thresh=0.5,
+            nms_topk=400,
+            post_nms=100,
+            roi_mode="align",
+            roi_size=(14, 14),
+            strides=16,
+            clip=4.14,
+            rpn_channel=1024,
+            base_size=16,
+            scales=(2, 4, 8, 16, 32),
+            ratios=(0.5, 1, 2),
+            alloc_size=(128, 128),
+            rpn_nms_thresh=0.7,
+            rpn_train_pre_nms=12000,
+            rpn_train_post_nms=2000,
+            rpn_test_pre_nms=6000,
+            rpn_test_post_nms=300,
+            rpn_min_size=16,
+            num_sample=128,
+            pos_iou_thresh=0.5,
+            pos_ratio=0.25,
+            max_num_gt=3000,
+            **kwargs
+        )
     else:
         net = faster_rcnn_resnet101_v1d_coco(pretrained=True)
         reuse_classes = [x for x in classes if x in net.classes]

examples/mxnet/scenegraph/model/reldn.py

 import pickle
 
+import dgl
+
 import gluoncv as gcv
 import mxnet as mx
 import numpy as np
-from mxnet import nd
-from mxnet.gluon import nn
-
-import dgl
 from dgl.nn.mxnet import GraphConv
 from dgl.utils import toindex
+from mxnet import nd
+from mxnet.gluon import nn
 
 __all__ = ["RelDN"]
 

examples/mxnet/scenegraph/train_faster_rcnn.py

@@ -11,19 +11,26 @@ import gluoncv as gcv
 import mxnet as mx
 import numpy as np
 from data import *
-from gluoncv import data as gdata
-from gluoncv import utils as gutils
+from gluoncv import data as gdata, utils as gutils
 from gluoncv.data.batchify import Append, FasterRCNNTrainBatchify, Tuple
 from gluoncv.data.transforms.presets.rcnn import (
-    FasterRCNNDefaultTrainTransform, FasterRCNNDefaultValTransform)
+    FasterRCNNDefaultTrainTransform,
+    FasterRCNNDefaultValTransform,
+)
 from gluoncv.model_zoo import get_model
 from gluoncv.utils.metrics.coco_detection import COCODetectionMetric
-from gluoncv.utils.metrics.rcnn import (RCNNAccMetric, RCNNL1LossMetric,
-                                        RPNAccMetric, RPNL1LossMetric)
+from gluoncv.utils.metrics.rcnn import (
+    RCNNAccMetric,
+    RCNNL1LossMetric,
+    RPNAccMetric,
+    RPNL1LossMetric,
+)
 from gluoncv.utils.metrics.voc_detection import VOC07MApMetric
 from gluoncv.utils.parallel import Parallel, Parallelizable
-from model import (faster_rcnn_resnet50_v1b_custom,
-                   faster_rcnn_resnet101_v1d_custom)
+from model import (
+    faster_rcnn_resnet101_v1d_custom,
+    faster_rcnn_resnet50_v1b_custom,
+)
 from mxnet import autograd, gluon
 from mxnet.contrib import amp
 

examples/mxnet/scenegraph/train_freq_prior.py

@@ -28,6 +28,7 @@ use_overlap = args.overlap
 PATH_TO_DATASETS = os.path.expanduser(args.json_path)
 path_to_json = os.path.join(PATH_TO_DATASETS, "rel_annotations_train.json")
 
+
 # format in y1y2x1x2
 def with_overlap(boxA, boxB):
     xA = max(boxA[2], boxB[2])

examples/mxnet/scenegraph/train_reldn.py

@@ -7,7 +7,7 @@ import numpy as np
 from data import *
 from gluoncv.data.batchify import Pad
 from gluoncv.utils import makedirs
-from model import RelDN, faster_rcnn_resnet101_v1d_custom
+from model import faster_rcnn_resnet101_v1d_custom, RelDN
 from mxnet import gluon, nd
 from utils import *
 

examples/mxnet/scenegraph/utils/build_graph.py

+import dgl
 import numpy as np
 from mxnet import nd
 
-import dgl
-
 
 def bbox_improve(bbox):
     """bbox encoding"""