[Misc] Black auto fix. (#4651)

Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>

examples/pytorch/GNN-FiLM/data_loader.py

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

examples/pytorch/GNN-FiLM/main.py

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

examples/pytorch/GNN-FiLM/utils.py

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

examples/pytorch/ggnn/data_utils.py

@@ -3,49 +3,100 @@ Data utils for processing bAbI datasets
 """
 
 import os
+import string
 
+import torch
 from torch.utils.data import DataLoader
+
 import dgl
-import torch
-import string
-from dgl.data.utils import download, get_download_dir, _get_dgl_url, extract_archive
+from dgl.data.utils import (
+    _get_dgl_url,
+    download,
+    extract_archive,
+    get_download_dir,
+)
 
 
 def get_babi_dataloaders(batch_size, train_size=50, task_id=4, q_type=0):
     _download_babi_data()
 
-    node_dict = dict(zip(list(string.ascii_uppercase), range(len(string.ascii_uppercase))))
+    node_dict = dict(
+        zip(list(string.ascii_uppercase), range(len(string.ascii_uppercase)))
+    )
 
     if task_id == 4:
-        edge_dict = {'n': 0, 's': 1, 'w': 2, 'e': 3}
+        edge_dict = {"n": 0, "s": 1, "w": 2, "e": 3}
         reverse_edge = {}
-        return _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, '04')
+        return _ns_dataloader(
+            train_size,
+            q_type,
+            batch_size,
+            node_dict,
+            edge_dict,
+            reverse_edge,
+            "04",
+        )
     elif task_id == 15:
-        edge_dict = {'is': 0, 'has_fear': 1}
+        edge_dict = {"is": 0, "has_fear": 1}
         reverse_edge = {}
-        return _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, '15')
+        return _ns_dataloader(
+            train_size,
+            q_type,
+            batch_size,
+            node_dict,
+            edge_dict,
+            reverse_edge,
+            "15",
+        )
     elif task_id == 16:
-        edge_dict = {'is': 0, 'has_color': 1}
+        edge_dict = {"is": 0, "has_color": 1}
         reverse_edge = {0: 0}
-        return _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, '16')
+        return _ns_dataloader(
+            train_size,
+            q_type,
+            batch_size,
+            node_dict,
+            edge_dict,
+            reverse_edge,
+            "16",
+        )
     elif task_id == 18:
-        edge_dict = {'>': 0, '<': 1}
-        label_dict = {'false': 0, 'true': 1}
+        edge_dict = {">": 0, "<": 1}
+        label_dict = {"false": 0, "true": 1}
         reverse_edge = {0: 1, 1: 0}
-        return _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_dict, reverse_edge, '18')
+        return _gc_dataloader(
+            train_size,
+            q_type,
+            batch_size,
+            node_dict,
+            edge_dict,
+            label_dict,
+            reverse_edge,
+            "18",
+        )
     elif task_id == 19:
-        edge_dict = {'n': 0, 's': 1, 'w': 2, 'e': 3, '<end>': 4}
+        edge_dict = {"n": 0, "s": 1, "w": 2, "e": 3, "<end>": 4}
         reverse_edge = {0: 1, 1: 0, 2: 3, 3: 2}
         max_seq_length = 2
-        return _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, reverse_edge, '19', max_seq_length)
-
-
-def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, path):
+        return _path_finding_dataloader(
+            train_size,
+            batch_size,
+            node_dict,
+            edge_dict,
+            reverse_edge,
+            "19",
+            max_seq_length,
+        )
+
+
+def _ns_dataloader(
+    train_size, q_type, batch_size, node_dict, edge_dict, reverse_edge, path
+):
     def _collate_fn(batch):
         graphs = []
         labels = []
         for d in batch:
-            edges = d['edges']
+            edges = d["edges"]
 
             node_ids = []
             for s, e, t in edges:
@@ -55,12 +106,12 @@ def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse
                     node_ids.append(t)
             g = dgl.DGLGraph()
             g.add_nodes(len(node_ids))
-            g.ndata['node_id'] = torch.tensor(node_ids, dtype=torch.long)
+            g.ndata["node_id"] = torch.tensor(node_ids, dtype=torch.long)
 
             nid2idx = dict(zip(node_ids, list(range(len(node_ids)))))
 
             # convert label to node index
-            label = d['eval'][2]
+            label = d["eval"][2]
             label_idx = nid2idx[label]
             labels.append(label_idx)
 
@@ -71,19 +122,26 @@ def _ns_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, reverse
                 if e in reverse_edge:
                     g.add_edge(nid2idx[t], nid2idx[s])
                     edge_types.append(reverse_edge[e])
-            g.edata['type'] = torch.tensor(edge_types, dtype=torch.long)
+            g.edata["type"] = torch.tensor(edge_types, dtype=torch.long)
             annotation = torch.zeros(len(node_ids), dtype=torch.long)
-            annotation[nid2idx[d['eval'][0]]] = 1
-            g.ndata['annotation'] = annotation.unsqueeze(-1)
+            annotation[nid2idx[d["eval"][0]]] = 1
+            g.ndata["annotation"] = annotation.unsqueeze(-1)
             graphs.append(g)
         batch_graph = dgl.batch(graphs)
         labels = torch.tensor(labels, dtype=torch.long)
         return batch_graph, labels
 
     def _get_dataloader(data, shuffle):
-        return DataLoader(dataset=data, batch_size=batch_size, shuffle=shuffle, collate_fn=_collate_fn)
-
-    train_set, dev_set, test_sets = _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type)
+        return DataLoader(
+            dataset=data,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            collate_fn=_collate_fn,
+        )
+
+    train_set, dev_set, test_sets = _convert_ns_dataset(
+        train_size, node_dict, edge_dict, path, q_type
+    )
     train_dataloader = _get_dataloader(train_set, True)
     dev_dataloader = _get_dataloader(dev_set, False)
     test_dataloaders = []
@@ -100,26 +158,36 @@ def _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type):
     def convert(file):
         dataset = []
         d = dict()
-        with open(file, 'r') as f:
+        with open(file, "r") as f:
             for i, line in enumerate(f.readlines()):
                 line = line.strip().split()
-                if line[0] == '1' and len(d) > 0:
+                if line[0] == "1" and len(d) > 0:
                     d = dict()
-                if line[1] == 'eval':
+                if line[1] == "eval":
                     # (src, edge, label)
-                    d['eval'] = (node_dict[line[2]], edge_dict[line[3]], node_dict[line[4]])
-                    if d['eval'][1] == q_type:
+                    d["eval"] = (
+                        node_dict[line[2]],
+                        edge_dict[line[3]],
+                        node_dict[line[4]],
+                    )
+                    if d["eval"][1] == q_type:
                         dataset.append(d)
                         if len(dataset) >= total_num:
                             break
                 else:
-                    if 'edges' not in d:
-                        d['edges'] = []
-                    d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]]))
+                    if "edges" not in d:
+                        d["edges"] = []
+                    d["edges"].append(
+                        (
+                            node_dict[line[1]],
+                            edge_dict[line[2]],
+                            node_dict[line[3]],
+                        )
+                    )
         return dataset
 
     download_dir = get_download_dir()
-    filename = os.path.join(download_dir, 'babi_data', path, 'data.txt')
+    filename = os.path.join(download_dir, "babi_data", path, "data.txt")
     data = convert(filename)
 
     assert len(data) == total_num
@@ -128,18 +196,27 @@ def _convert_ns_dataset(train_size, node_dict, edge_dict, path, q_type):
     dev_set = data[950:1000]
     test_sets = []
     for i in range(10):
-        test = data[1000 * (i + 1): 1000 * (i + 2)]
+        test = data[1000 * (i + 1) : 1000 * (i + 2)]
         test_sets.append(test)
 
     return train_set, dev_set, test_sets
 
 
-def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_dict, reverse_edge, path):
+def _gc_dataloader(
+    train_size,
+    q_type,
+    batch_size,
+    node_dict,
+    edge_dict,
+    label_dict,
+    reverse_edge,
+    path,
+):
     def _collate_fn(batch):
         graphs = []
         labels = []
         for d in batch:
-            edges = d['edges']
+            edges = d["edges"]
 
             node_ids = []
             for s, e, t in edges:
@@ -149,11 +226,11 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d
                     node_ids.append(t)
             g = dgl.DGLGraph()
             g.add_nodes(len(node_ids))
-            g.ndata['node_id'] = torch.tensor(node_ids, dtype=torch.long)
+            g.ndata["node_id"] = torch.tensor(node_ids, dtype=torch.long)
 
             nid2idx = dict(zip(node_ids, list(range(len(node_ids)))))
 
-            labels.append(d['eval'][-1])
+            labels.append(d["eval"][-1])
 
             edge_types = []
             for s, e, t in edges:
@@ -162,20 +239,27 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d
                 if e in reverse_edge:
                     g.add_edge(nid2idx[t], nid2idx[s])
                     edge_types.append(reverse_edge[e])
-            g.edata['type'] = torch.tensor(edge_types, dtype=torch.long)
+            g.edata["type"] = torch.tensor(edge_types, dtype=torch.long)
             annotation = torch.zeros([len(node_ids), 2], dtype=torch.long)
-            annotation[nid2idx[d['eval'][0]]][0] = 1
-            annotation[nid2idx[d['eval'][2]]][1] = 1
-            g.ndata['annotation'] = annotation
+            annotation[nid2idx[d["eval"][0]]][0] = 1
+            annotation[nid2idx[d["eval"][2]]][1] = 1
+            g.ndata["annotation"] = annotation
             graphs.append(g)
         batch_graph = dgl.batch(graphs)
         labels = torch.tensor(labels, dtype=torch.long)
         return batch_graph, labels
 
     def _get_dataloader(data, shuffle):
-        return DataLoader(dataset=data, batch_size=batch_size, shuffle=shuffle, collate_fn=_collate_fn)
-
-    train_set, dev_set, test_sets = _convert_gc_dataset(train_size, node_dict, edge_dict, label_dict, path, q_type)
+        return DataLoader(
+            dataset=data,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            collate_fn=_collate_fn,
+        )
+
+    train_set, dev_set, test_sets = _convert_gc_dataset(
+        train_size, node_dict, edge_dict, label_dict, path, q_type
+    )
     train_dataloader = _get_dataloader(train_set, True)
     dev_dataloader = _get_dataloader(dev_set, False)
     test_dataloaders = []
@@ -186,33 +270,46 @@ def _gc_dataloader(train_size, q_type, batch_size, node_dict, edge_dict, label_d
     return train_dataloader, dev_dataloader, test_dataloaders
 
 
-def _convert_gc_dataset(train_size, node_dict, edge_dict, label_dict, path, q_type):
+def _convert_gc_dataset(
+    train_size, node_dict, edge_dict, label_dict, path, q_type
+):
     total_num = 11000
 
     def convert(file):
         dataset = []
         d = dict()
-        with open(file, 'r') as f:
+        with open(file, "r") as f:
             for i, line in enumerate(f.readlines()):
                 line = line.strip().split()
-                if line[0] == '1' and len(d) > 0:
+                if line[0] == "1" and len(d) > 0:
                     d = dict()
-                if line[1] == 'eval':
+                if line[1] == "eval":
                     # (src, edge, label)
-                    if 'eval' not in d:
-                        d['eval'] = (node_dict[line[2]], edge_dict[line[3]], node_dict[line[4]], label_dict[line[5]])
-                        if d['eval'][1] == q_type:
+                    if "eval" not in d:
+                        d["eval"] = (
+                            node_dict[line[2]],
+                            edge_dict[line[3]],
+                            node_dict[line[4]],
+                            label_dict[line[5]],
+                        )
+                        if d["eval"][1] == q_type:
                             dataset.append(d)
                             if len(dataset) >= total_num:
                                 break
                 else:
-                    if 'edges' not in d:
-                        d['edges'] = []
-                    d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]]))
+                    if "edges" not in d:
+                        d["edges"] = []
+                    d["edges"].append(
+                        (
+                            node_dict[line[1]],
+                            edge_dict[line[2]],
+                            node_dict[line[3]],
+                        )
+                    )
         return dataset
 
     download_dir = get_download_dir()
-    filename = os.path.join(download_dir, 'babi_data', path, 'data.txt')
+    filename = os.path.join(download_dir, "babi_data", path, "data.txt")
     data = convert(filename)
 
     assert len(data) == total_num
@@ -221,19 +318,27 @@ def _convert_gc_dataset(train_size, node_dict, edge_dict, label_dict, path, q_ty
     dev_set = data[950:1000]
     test_sets = []
     for i in range(10):
-        test = data[1000 * (i + 1): 1000 * (i + 2)]
+        test = data[1000 * (i + 1) : 1000 * (i + 2)]
         test_sets.append(test)
 
     return train_set, dev_set, test_sets
 
 
-def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, reverse_edge, path, max_seq_length):
+def _path_finding_dataloader(
+    train_size,
+    batch_size,
+    node_dict,
+    edge_dict,
+    reverse_edge,
+    path,
+    max_seq_length,
+):
     def _collate_fn(batch):
         graphs = []
         ground_truths = []
         seq_lengths = []
         for d in batch:
-            edges = d['edges']
+            edges = d["edges"]
 
             node_ids = []
             for s, e, t in edges:
@@ -243,12 +348,14 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever
                     node_ids.append(t)
             g = dgl.DGLGraph()
             g.add_nodes(len(node_ids))
-            g.ndata['node_id'] = torch.tensor(node_ids, dtype=torch.long)
+            g.ndata["node_id"] = torch.tensor(node_ids, dtype=torch.long)
 
             nid2idx = dict(zip(node_ids, list(range(len(node_ids)))))
 
-            truth = d['seq_out'] + [edge_dict['<end>']] * (max_seq_length - len(d['seq_out']))
-            seq_len = len(d['seq_out'])
+            truth = d["seq_out"] + [edge_dict["<end>"]] * (
+                max_seq_length - len(d["seq_out"])
+            )
+            seq_len = len(d["seq_out"])
             ground_truths.append(truth)
             seq_lengths.append(seq_len)
 
@@ -259,11 +366,11 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever
                 if e in reverse_edge:
                     g.add_edge(nid2idx[t], nid2idx[s])
                     edge_types.append(reverse_edge[e])
-            g.edata['type'] = torch.tensor(edge_types, dtype=torch.long)
+            g.edata["type"] = torch.tensor(edge_types, dtype=torch.long)
             annotation = torch.zeros([len(node_ids), 2], dtype=torch.long)
-            annotation[nid2idx[d['eval'][0]]][0] = 1
-            annotation[nid2idx[d['eval'][1]]][1] = 1
-            g.ndata['annotation'] = annotation
+            annotation[nid2idx[d["eval"][0]]][0] = 1
+            annotation[nid2idx[d["eval"][1]]][1] = 1
+            g.ndata["annotation"] = annotation
             graphs.append(g)
         batch_graph = dgl.batch(graphs)
         ground_truths = torch.tensor(ground_truths, dtype=torch.long)
@@ -271,9 +378,16 @@ def _path_finding_dataloader(train_size, batch_size, node_dict, edge_dict, rever
         return batch_graph, ground_truths, seq_lengths
 
     def _get_dataloader(data, shuffle):
-        return DataLoader(dataset=data, batch_size=batch_size, shuffle=shuffle, collate_fn=_collate_fn)
-
-    train_set, dev_set, test_sets = _convert_path_finding(train_size, node_dict, edge_dict, path)
+        return DataLoader(
+            dataset=data,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            collate_fn=_collate_fn,
+        )
+
+    train_set, dev_set, test_sets = _convert_path_finding(
+        train_size, node_dict, edge_dict, path
+    )
     train_dataloader = _get_dataloader(train_set, True)
     dev_dataloader = _get_dataloader(dev_set, False)
     test_dataloaders = []
@@ -290,29 +404,35 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
     def convert(file):
         dataset = []
         d = dict()
-        with open(file, 'r') as f:
+        with open(file, "r") as f:
             for line in f.readlines():
                 line = line.strip().split()
-                if line[0] == '1' and len(d) > 0:
+                if line[0] == "1" and len(d) > 0:
                     d = dict()
-                if line[1] == 'eval':
+                if line[1] == "eval":
                     # (src, edge, label)
-                    d['eval'] = (node_dict[line[3]], node_dict[line[4]])
-                    d['seq_out'] = []
-                    seq_out = line[5].split(',')
+                    d["eval"] = (node_dict[line[3]], node_dict[line[4]])
+                    d["seq_out"] = []
+                    seq_out = line[5].split(",")
                     for e in seq_out:
-                        d['seq_out'].append(edge_dict[e])
+                        d["seq_out"].append(edge_dict[e])
                     dataset.append(d)
                     if len(dataset) >= total_num:
                         break
                 else:
-                    if 'edges' not in d:
-                        d['edges'] = []
-                    d['edges'].append((node_dict[line[1]], edge_dict[line[2]], node_dict[line[3]]))
+                    if "edges" not in d:
+                        d["edges"] = []
+                    d["edges"].append(
+                        (
+                            node_dict[line[1]],
+                            edge_dict[line[2]],
+                            node_dict[line[3]],
+                        )
+                    )
         return dataset
 
     download_dir = get_download_dir()
-    filename = os.path.join(download_dir, 'babi_data', path, 'data.txt')
+    filename = os.path.join(download_dir, "babi_data", path, "data.txt")
     data = convert(filename)
 
     assert len(data) == total_num
@@ -321,7 +441,7 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
     dev_set = data[950:1000]
     test_sets = []
     for i in range(10):
-        test = data[1000 * (i + 1): 1000 * (i + 2)]
+        test = data[1000 * (i + 1) : 1000 * (i + 2)]
         test_sets.append(test)
 
     return train_set, dev_set, test_sets
@@ -329,11 +449,11 @@ def _convert_path_finding(train_size, node_dict, edge_dict, path):
 
 def _download_babi_data():
     download_dir = get_download_dir()
-    zip_file_path = os.path.join(download_dir, 'babi_data.zip')
+    zip_file_path = os.path.join(download_dir, "babi_data.zip")
 
-    data_url = _get_dgl_url('models/ggnn_babi_data.zip')
+    data_url = _get_dgl_url("models/ggnn_babi_data.zip")
     download(data_url, path=zip_file_path)
 
-    extract_dir = os.path.join(download_dir, 'babi_data')
+    extract_dir = os.path.join(download_dir, "babi_data")
     if not os.path.exists(extract_dir):
         extract_archive(zip_file_path, extract_dir)

examples/pytorch/ggnn/ggnn_gc.py

 """
 Gated Graph Neural Network module for graph classification tasks
 """
-from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
 import torch
 from torch import nn
 
+from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
+
 
 class GraphClsGGNN(nn.Module):
-    def __init__(self,
-                 annotation_size,
-                 out_feats,
-                 n_steps,
-                 n_etypes,
-                 num_cls):
+    def __init__(self, annotation_size, out_feats, n_steps, n_etypes, num_cls):
         super(GraphClsGGNN, self).__init__()
 
         self.annotation_size = annotation_size
         self.out_feats = out_feats
 
-        self.ggnn = GatedGraphConv(in_feats=out_feats,
+        self.ggnn = GatedGraphConv(
+            in_feats=out_feats,
             out_feats=out_feats,
             n_steps=n_steps,
-                                   n_etypes=n_etypes)
+            n_etypes=n_etypes,
+        )
 
         pooling_gate_nn = nn.Linear(annotation_size + out_feats, 1)
         self.pooling = GlobalAttentionPooling(pooling_gate_nn)
@@ -30,16 +28,18 @@ class GraphClsGGNN(nn.Module):
         self.loss_fn = nn.CrossEntropyLoss()
 
     def forward(self, graph, labels=None):
-        etypes = graph.edata.pop('type')
-        annotation = graph.ndata.pop('annotation').float()
+        etypes = graph.edata.pop("type")
+        annotation = graph.ndata.pop("annotation").float()
 
         assert annotation.size()[-1] == self.annotation_size
 
         node_num = graph.number_of_nodes()
 
-        zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size],
+        zero_pad = torch.zeros(
+            [node_num, self.out_feats - self.annotation_size],
             dtype=torch.float,
-                               device=annotation.device)
+            device=annotation.device,
+        )
 
         h1 = torch.cat([annotation, zero_pad], -1)
         out = self.ggnn(graph, h1, etypes)

examples/pytorch/ggnn/ggnn_ns.py

 """
 Gated Graph Neural Network module for node selection tasks
 """
-from dgl.nn.pytorch import GatedGraphConv
 import torch
 from torch import nn
+
 import dgl
+from dgl.nn.pytorch import GatedGraphConv
 
 
 class NodeSelectionGGNN(nn.Module):
-    def __init__(self,
-                 annotation_size,
-                 out_feats,
-                 n_steps,
-                 n_etypes):
+    def __init__(self, annotation_size, out_feats, n_steps, n_etypes):
         super(NodeSelectionGGNN, self).__init__()
 
         self.annotation_size = annotation_size
         self.out_feats = out_feats
 
-        self.ggnn = GatedGraphConv(in_feats=out_feats,
+        self.ggnn = GatedGraphConv(
+            in_feats=out_feats,
             out_feats=out_feats,
             n_steps=n_steps,
-                                   n_etypes=n_etypes)
+            n_etypes=n_etypes,
+        )
 
         self.output_layer = nn.Linear(annotation_size + out_feats, 1)
         self.loss_fn = nn.CrossEntropyLoss()
 
     def forward(self, graph, labels=None):
-        etypes = graph.edata.pop('type')
-        annotation = graph.ndata.pop('annotation').float()
+        etypes = graph.edata.pop("type")
+        annotation = graph.ndata.pop("annotation").float()
 
         assert annotation.size()[-1] == self.annotation_size
 
         node_num = graph.number_of_nodes()
 
-        zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size],
+        zero_pad = torch.zeros(
+            [node_num, self.out_feats - self.annotation_size],
             dtype=torch.float,
-                               device=annotation.device)
+            device=annotation.device,
+        )
 
         h1 = torch.cat([annotation, zero_pad], -1)
         out = self.ggnn(graph, h1, etypes)
 
-        all_logits = self.output_layer(torch.cat([out, annotation], -1)).squeeze(-1)
-        graph.ndata['logits'] = all_logits
+        all_logits = self.output_layer(
+            torch.cat([out, annotation], -1)
+        ).squeeze(-1)
+        graph.ndata["logits"] = all_logits
 
         batch_g = dgl.unbatch(graph)
 
@@ -50,7 +53,7 @@ class NodeSelectionGGNN(nn.Module):
         if labels is not None:
             loss = 0.0
         for i, g in enumerate(batch_g):
-            logits = g.ndata['logits']
+            logits = g.ndata["logits"]
             preds.append(torch.argmax(logits))
             if labels is not None:
                 logits = logits.unsqueeze(0)

examples/pytorch/ggnn/ggsnn.py

@@ -2,42 +2,49 @@
 Gated Graph Sequence Neural Network for sequence outputs
 """
 
-from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
 import torch
-from torch import nn
 import torch.nn.functional as F
+from torch import nn
+
+from dgl.nn.pytorch import GatedGraphConv, GlobalAttentionPooling
 
 
 class GGSNN(nn.Module):
-    def __init__(self,
+    def __init__(
+        self,
         annotation_size,
         out_feats,
         n_steps,
         n_etypes,
         max_seq_length,
-                 num_cls):
+        num_cls,
+    ):
         super(GGSNN, self).__init__()
 
         self.annotation_size = annotation_size
         self.out_feats = out_feats
         self.max_seq_length = max_seq_length
 
-        self.ggnn = GatedGraphConv(in_feats=out_feats,
+        self.ggnn = GatedGraphConv(
+            in_feats=out_feats,
             out_feats=out_feats,
             n_steps=n_steps,
-                                   n_etypes=n_etypes)
+            n_etypes=n_etypes,
+        )
 
-        self.annotation_out_layer = nn.Linear(annotation_size + out_feats, annotation_size)
+        self.annotation_out_layer = nn.Linear(
+            annotation_size + out_feats, annotation_size
+        )
 
         pooling_gate_nn = nn.Linear(annotation_size + out_feats, 1)
         self.pooling = GlobalAttentionPooling(pooling_gate_nn)
 
         self.output_layer = nn.Linear(annotation_size + out_feats, num_cls)
-        self.loss_fn = nn.CrossEntropyLoss(reduction='none')
+        self.loss_fn = nn.CrossEntropyLoss(reduction="none")
 
     def forward(self, graph, seq_lengths, ground_truth=None):
-        etypes = graph.edata.pop('type')
-        annotation = graph.ndata.pop('annotation').float()
+        etypes = graph.edata.pop("type")
+        annotation = graph.ndata.pop("annotation").float()
 
         assert annotation.size()[-1] == self.annotation_size
 
@@ -45,9 +52,11 @@ class GGSNN(nn.Module):
 
         all_logits = []
         for _ in range(self.max_seq_length):
-            zero_pad = torch.zeros([node_num, self.out_feats - self.annotation_size],
+            zero_pad = torch.zeros(
+                [node_num, self.out_feats - self.annotation_size],
                 dtype=torch.float,
-                                   device=annotation.device)
+                device=annotation.device,
+            )
 
             h1 = torch.cat([annotation.detach(), zero_pad], -1)
             out = self.ggnn(graph, h1, etypes)
@@ -71,14 +80,18 @@ def sequence_loss(logits, ground_truth, seq_length=None):
     def sequence_mask(length):
         max_length = logits.size(1)
         batch_size = logits.size(0)
-        range_tensor = torch.arange(0, max_length, dtype=seq_length.dtype, device=seq_length.device)
-        range_tensor = torch.stack([range_tensor]*batch_size, 0)
+        range_tensor = torch.arange(
+            0, max_length, dtype=seq_length.dtype, device=seq_length.device
+        )
+        range_tensor = torch.stack([range_tensor] * batch_size, 0)
 
-        expanded_length = torch.stack([length]*max_length, -1)
+        expanded_length = torch.stack([length] * max_length, -1)
         mask = (range_tensor < expanded_length).float()
         return mask
 
-    loss = nn.CrossEntropyLoss(reduction='none')(logits.permute((0, 2, 1)), ground_truth)
+    loss = nn.CrossEntropyLoss(reduction="none")(
+        logits.permute((0, 2, 1)), ground_truth
+    )
 
     if seq_length is None:
         loss = loss.mean()

examples/pytorch/ggnn/train_gc.py

@@ -3,34 +3,38 @@ Training and testing for graph classification tasks in bAbI
 """
 
 import argparse
+
+import numpy as np
+import torch
 from data_utils import get_babi_dataloaders
 from ggnn_gc import GraphClsGGNN
 from torch.optim import Adam
-import torch
-import numpy as np
 
 
 def main(args):
     out_feats = {18: 3}
     n_etypes = {18: 2}
 
-    train_dataloader, dev_dataloader, test_dataloaders = \
-        get_babi_dataloaders(batch_size=args.batch_size,
+    train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
+        batch_size=args.batch_size,
         train_size=args.train_num,
         task_id=args.task_id,
-                             q_type=args.question_id)
+        q_type=args.question_id,
+    )
 
-    model = GraphClsGGNN(annotation_size=2,
+    model = GraphClsGGNN(
+        annotation_size=2,
         out_feats=out_feats[args.task_id],
         n_steps=5,
         n_etypes=n_etypes[args.task_id],
-                         num_cls=2)
+        num_cls=2,
+    )
     opt = Adam(model.parameters(), lr=args.lr)
 
-    print(f'Task {args.task_id}, question_id {args.question_id}')
+    print(f"Task {args.task_id}, question_id {args.question_id}")
 
-    print(f'Training set size: {len(train_dataloader.dataset)}')
-    print(f'Dev set size: {len(dev_dataloader.dataset)}')
+    print(f"Training set size: {len(train_dataloader.dataset)}")
+    print(f"Dev set size: {len(dev_dataloader.dataset)}")
 
     # training and dev stage
     for epoch in range(args.epochs):
@@ -42,7 +46,7 @@ def main(args):
             loss.backward()
             opt.step()
             if epoch % 20 == 0:
-                print(f'Epoch {epoch}, batch {i} loss: {loss.data}')
+                print(f"Epoch {epoch}, batch {i} loss: {loss.data}")
 
         if epoch % 20 != 0:
             continue
@@ -62,7 +66,7 @@ def main(args):
 
     # test stage
     for i, dataloader in enumerate(test_dataloaders):
-        print(f'Test set {i} size: {len(dataloader.dataset)}')
+        print(f"Test set {i} size: {len(dataloader.dataset)}")
 
     test_acc_list = []
     for dataloader in test_dataloaders:
@@ -83,23 +87,29 @@ def main(args):
     test_acc_mean = np.mean(test_acc_list)
     test_acc_std = np.std(test_acc_list)
 
-    print(f'Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}')
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Gated Graph Neural Networks for graph classification tasks in bAbI')
-    parser.add_argument('--task_id', type=int, default=18,
-                        help='task id from 1 to 20')
-    parser.add_argument('--question_id', type=int, default=0,
-                        help='question id for each task')
-    parser.add_argument('--train_num', type=int, default=950,
-                        help='Number of training examples')
-    parser.add_argument('--batch_size', type=int, default=50,
-                        help='batch size')
-    parser.add_argument('--lr', type=float, default=1e-3,
-                        help='learning rate')
-    parser.add_argument('--epochs', type=int, default=200,
-                        help='number of training epochs')
+    print(
+        f"Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}"
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Gated Graph Neural Networks for graph classification tasks in bAbI"
+    )
+    parser.add_argument(
+        "--task_id", type=int, default=18, help="task id from 1 to 20"
+    )
+    parser.add_argument(
+        "--question_id", type=int, default=0, help="question id for each task"
+    )
+    parser.add_argument(
+        "--train_num", type=int, default=950, help="Number of training examples"
+    )
+    parser.add_argument("--batch_size", type=int, default=50, help="batch size")
+    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
+    parser.add_argument(
+        "--epochs", type=int, default=200, help="number of training epochs"
+    )
 
     args = parser.parse_args()
 

examples/pytorch/ggnn/train_ns.py

@@ -3,34 +3,38 @@ Training and testing for node selection tasks in bAbI
 """
 
 import argparse
+import time
+
+import numpy as np
+import torch
 from data_utils import get_babi_dataloaders
 from ggnn_ns import NodeSelectionGGNN
 from torch.optim import Adam
-import torch
-import numpy as np
-import time
 
 
 def main(args):
     out_feats = {4: 4, 15: 5, 16: 6}
     n_etypes = {4: 4, 15: 2, 16: 2}
 
-    train_dataloader, dev_dataloader, test_dataloaders = \
-        get_babi_dataloaders(batch_size=args.batch_size,
+    train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
+        batch_size=args.batch_size,
         train_size=args.train_num,
         task_id=args.task_id,
-                             q_type=args.question_id)
+        q_type=args.question_id,
+    )
 
-    model = NodeSelectionGGNN(annotation_size=1,
+    model = NodeSelectionGGNN(
+        annotation_size=1,
         out_feats=out_feats[args.task_id],
         n_steps=5,
-                              n_etypes=n_etypes[args.task_id])
+        n_etypes=n_etypes[args.task_id],
+    )
     opt = Adam(model.parameters(), lr=args.lr)
 
-    print(f'Task {args.task_id}, question_id {args.question_id}')
+    print(f"Task {args.task_id}, question_id {args.question_id}")
 
-    print(f'Training set size: {len(train_dataloader.dataset)}')
-    print(f'Dev set size: {len(dev_dataloader.dataset)}')
+    print(f"Training set size: {len(train_dataloader.dataset)}")
+    print(f"Dev set size: {len(dev_dataloader.dataset)}")
 
     # training and dev stage
     for epoch in range(args.epochs):
@@ -41,7 +45,7 @@ def main(args):
             opt.zero_grad()
             loss.backward()
             opt.step()
-            print(f'Epoch {epoch}, batch {i} loss: {loss.data}')
+            print(f"Epoch {epoch}, batch {i} loss: {loss.data}")
 
         dev_preds = []
         dev_labels = []
@@ -49,7 +53,9 @@ def main(args):
         for g, labels in dev_dataloader:
             with torch.no_grad():
                 preds = model(g)
-                preds = torch.tensor(preds, dtype=torch.long).data.numpy().tolist()
+                preds = (
+                    torch.tensor(preds, dtype=torch.long).data.numpy().tolist()
+                )
                 labels = labels.data.numpy().tolist()
                 dev_preds += preds
                 dev_labels += labels
@@ -59,7 +65,7 @@ def main(args):
 
     # test stage
     for i, dataloader in enumerate(test_dataloaders):
-        print(f'Test set {i} size: {len(dataloader.dataset)}')
+        print(f"Test set {i} size: {len(dataloader.dataset)}")
 
     test_acc_list = []
     for dataloader in test_dataloaders:
@@ -69,7 +75,9 @@ def main(args):
         for g, labels in dataloader:
             with torch.no_grad():
                 preds = model(g)
-                preds = torch.tensor(preds, dtype=torch.long).data.numpy().tolist()
+                preds = (
+                    torch.tensor(preds, dtype=torch.long).data.numpy().tolist()
+                )
                 labels = labels.data.numpy().tolist()
                 test_preds += preds
                 test_labels += labels
@@ -80,23 +88,29 @@ def main(args):
     test_acc_mean = np.mean(test_acc_list)
     test_acc_std = np.std(test_acc_list)
 
-    print(f'Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}')
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Gated Graph Neural Networks for node selection tasks in bAbI')
-    parser.add_argument('--task_id', type=int, default=16,
-                        help='task id from 1 to 20')
-    parser.add_argument('--question_id', type=int, default=1,
-                        help='question id for each task')
-    parser.add_argument('--train_num', type=int, default=50,
-                        help='Number of training examples')
-    parser.add_argument('--batch_size', type=int, default=10,
-                        help='batch size')
-    parser.add_argument('--lr', type=float, default=1e-3,
-                        help='learning rate')
-    parser.add_argument('--epochs', type=int, default=100,
-                        help='number of training epochs')
+    print(
+        f"Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}"
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Gated Graph Neural Networks for node selection tasks in bAbI"
+    )
+    parser.add_argument(
+        "--task_id", type=int, default=16, help="task id from 1 to 20"
+    )
+    parser.add_argument(
+        "--question_id", type=int, default=1, help="question id for each task"
+    )
+    parser.add_argument(
+        "--train_num", type=int, default=50, help="Number of training examples"
+    )
+    parser.add_argument("--batch_size", type=int, default=10, help="batch size")
+    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
+    parser.add_argument(
+        "--epochs", type=int, default=100, help="number of training epochs"
+    )
 
     args = parser.parse_args()
 

examples/pytorch/ggnn/train_path_finding.py

@@ -4,35 +4,39 @@ Here we take task 19 'Path Finding' as an example
 """
 
 import argparse
+
+import numpy as np
+import torch
 from data_utils import get_babi_dataloaders
 from ggsnn import GGSNN
 from torch.optim import Adam
-import torch
-import numpy as np
 
 
 def main(args):
     out_feats = {19: 6}
     n_etypes = {19: 4}
 
-    train_dataloader, dev_dataloader, test_dataloaders = \
-        get_babi_dataloaders(batch_size=args.batch_size,
+    train_dataloader, dev_dataloader, test_dataloaders = get_babi_dataloaders(
+        batch_size=args.batch_size,
         train_size=args.train_num,
         task_id=args.task_id,
-                             q_type=-1)
+        q_type=-1,
+    )
 
-    model = GGSNN(annotation_size=2,
+    model = GGSNN(
+        annotation_size=2,
         out_feats=out_feats[args.task_id],
         n_steps=5,
         n_etypes=n_etypes[args.task_id],
         max_seq_length=2,
-                  num_cls=5)
+        num_cls=5,
+    )
     opt = Adam(model.parameters(), lr=args.lr)
 
-    print(f'Task {args.task_id}')
+    print(f"Task {args.task_id}")
 
-    print(f'Training set size: {len(train_dataloader.dataset)}')
-    print(f'Dev set size: {len(dev_dataloader.dataset)}')
+    print(f"Training set size: {len(train_dataloader.dataset)}")
+    print(f"Dev set size: {len(dev_dataloader.dataset)}")
 
     # training and dev stage
     for epoch in range(args.epochs):
@@ -44,7 +48,7 @@ def main(args):
             loss.backward()
             opt.step()
             if epoch % 20 == 0:
-                print(f'Epoch {epoch}, batch {i} loss: {loss.data}')
+                print(f"Epoch {epoch}, batch {i} loss: {loss.data}")
 
         if epoch % 20 != 0:
             continue
@@ -65,7 +69,7 @@ def main(args):
 
     # test stage
     for i, dataloader in enumerate(test_dataloaders):
-        print(f'Test set {i} size: {len(dataloader.dataset)}')
+        print(f"Test set {i} size: {len(dataloader.dataset)}")
 
     test_acc_list = []
     for dataloader in test_dataloaders:
@@ -87,22 +91,27 @@ def main(args):
     test_acc_mean = np.mean(test_acc_list)
     test_acc_std = np.std(test_acc_list)
 
-    print(f'Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}')
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Gated Graph Sequence Neural Networks for sequential output tasks in '
-                                                 'bAbI')
-    parser.add_argument('--task_id', type=int, default=19,
-                        help='task id from 1 to 20')
-    parser.add_argument('--train_num', type=int, default=250,
-                        help='Number of training examples')
-    parser.add_argument('--batch_size', type=int, default=10,
-                        help='batch size')
-    parser.add_argument('--lr', type=float, default=1e-3,
-                        help='learning rate')
-    parser.add_argument('--epochs', type=int, default=200,
-                        help='number of training epochs')
+    print(
+        f"Mean of accuracy in 10 test datasets: {test_acc_mean}, std: {test_acc_std}"
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Gated Graph Sequence Neural Networks for sequential output tasks in "
+        "bAbI"
+    )
+    parser.add_argument(
+        "--task_id", type=int, default=19, help="task id from 1 to 20"
+    )
+    parser.add_argument(
+        "--train_num", type=int, default=250, help="Number of training examples"
+    )
+    parser.add_argument("--batch_size", type=int, default=10, help="batch size")
+    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
+    parser.add_argument(
+        "--epochs", type=int, default=200, help="number of training epochs"
+    )
 
     args = parser.parse_args()
 

examples/pytorch/gin/train.py

+import argparse
+
 import numpy as np
 import torch
 import torch.nn as nn
-import torch.optim as optim
 import torch.nn.functional as F
-from torch.utils.data.sampler import SubsetRandomSampler
+import torch.optim as optim
 from sklearn.model_selection import StratifiedKFold
+from torch.utils.data.sampler import SubsetRandomSampler
+
 from dgl.data import GINDataset
 from dgl.dataloading import GraphDataLoader
 from dgl.nn.pytorch.conv import GINConv
 from dgl.nn.pytorch.glob import SumPooling
-import argparse
+
 
 class MLP(nn.Module):
     """Construct two-layer MLP-type aggreator for GIN model"""
+
     def __init__(self, input_dim, hidden_dim, output_dim):
         super().__init__()
         self.linears = nn.ModuleList()
@@ -26,6 +30,7 @@ class MLP(nn.Module):
         h = F.relu(self.batch_norm(self.linears[0](h)))
         return self.linears[1](h)
 
+
 class GIN(nn.Module):
     def __init__(self, input_dim, hidden_dim, output_dim):
         super().__init__()
@@ -38,7 +43,9 @@ class GIN(nn.Module):
                 mlp = MLP(input_dim, hidden_dim, hidden_dim)
             else:
                 mlp = MLP(hidden_dim, hidden_dim, hidden_dim)
-            self.ginlayers.append(GINConv(mlp, learn_eps=False)) # set to True if learning epsilon
+            self.ginlayers.append(
+                GINConv(mlp, learn_eps=False)
+            )  # set to True if learning epsilon
             self.batch_norms.append(nn.BatchNorm1d(hidden_dim))
         # linear functions for graph sum poolings of output of each layer
         self.linear_prediction = nn.ModuleList()
@@ -48,7 +55,9 @@ class GIN(nn.Module):
             else:
                 self.linear_prediction.append(nn.Linear(hidden_dim, output_dim))
         self.drop = nn.Dropout(0.5)
-        self.pool = SumPooling() # change to mean readout (AvgPooling) on social network datasets
+        self.pool = (
+            SumPooling()
+        )  # change to mean readout (AvgPooling) on social network datasets
 
     def forward(self, g, h):
         # list of hidden representation at each layer (including the input layer)
@@ -65,6 +74,7 @@ class GIN(nn.Module):
             score_over_layer += self.drop(self.linear_prediction[i](pooled_h))
         return score_over_layer
 
+
 def split_fold10(labels, fold_idx=0):
     skf = StratifiedKFold(n_splits=10, shuffle=True, random_state=0)
     idx_list = []
@@ -73,6 +83,7 @@ def split_fold10(labels, fold_idx=0):
     train_idx, valid_idx = idx_list[fold_idx]
     return train_idx, valid_idx
 
+
 def evaluate(dataloader, device, model):
     model.eval()
     total = 0
@@ -80,7 +91,7 @@ def evaluate(dataloader, device, model):
     for batched_graph, labels in dataloader:
         batched_graph = batched_graph.to(device)
         labels = labels.to(device)
-        feat = batched_graph.ndata.pop('attr')
+        feat = batched_graph.ndata.pop("attr")
         total += len(labels)
         logits = model(batched_graph, feat)
         _, predicted = torch.max(logits, 1)
@@ -88,6 +99,7 @@ def evaluate(dataloader, device, model):
     acc = 1.0 * total_correct / total
     return acc
 
+
 def train(train_loader, val_loader, device, model):
     # loss function, optimizer and scheduler
     loss_fcn = nn.CrossEntropyLoss()
@@ -101,7 +113,7 @@ def train(train_loader, val_loader, device, model):
         for batch, (batched_graph, labels) in enumerate(train_loader):
             batched_graph = batched_graph.to(device)
             labels = labels.to(device)
-            feat = batched_graph.ndata.pop('attr')
+            feat = batched_graph.ndata.pop("attr")
             logits = model(batched_graph, feat)
             loss = loss_fcn(logits, labels)
             optimizer.zero_grad()
@@ -111,28 +123,46 @@ def train(train_loader, val_loader, device, model):
         scheduler.step()
         train_acc = evaluate(train_loader, device, model)
         valid_acc = evaluate(val_loader, device, model)
-        print("Epoch {:05d} | Loss {:.4f} | Train Acc. {:.4f} | Validation Acc. {:.4f} "
-              . format(epoch, total_loss / (batch + 1), train_acc, valid_acc))
+        print(
+            "Epoch {:05d} | Loss {:.4f} | Train Acc. {:.4f} | Validation Acc. {:.4f} ".format(
+                epoch, total_loss / (batch + 1), train_acc, valid_acc
+            )
+        )
+
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--dataset', type=str, default="MUTAG",
-                        choices=['MUTAG', 'PTC', 'NCI1', 'PROTEINS'],
-                        help='name of dataset (default: MUTAG)')
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="MUTAG",
+        choices=["MUTAG", "PTC", "NCI1", "PROTEINS"],
+        help="name of dataset (default: MUTAG)",
+    )
     args = parser.parse_args()
-    print(f'Training with DGL built-in GINConv module with a fixed epsilon = 0')
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Training with DGL built-in GINConv module with a fixed epsilon = 0")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     # load and split dataset
-    dataset = GINDataset(args.dataset, self_loop=True, degree_as_nlabel=False) # add self_loop and disable one-hot encoding for input features
+    dataset = GINDataset(
+        args.dataset, self_loop=True, degree_as_nlabel=False
+    )  # add self_loop and disable one-hot encoding for input features
     labels = [l for _, l in dataset]
     train_idx, val_idx = split_fold10(labels)
 
     # create dataloader
-    train_loader = GraphDataLoader(dataset, sampler=SubsetRandomSampler(train_idx),
-                                   batch_size=128, pin_memory=torch.cuda.is_available())
-    val_loader = GraphDataLoader(dataset, sampler=SubsetRandomSampler(val_idx),
-                                 batch_size=128, pin_memory=torch.cuda.is_available())
+    train_loader = GraphDataLoader(
+        dataset,
+        sampler=SubsetRandomSampler(train_idx),
+        batch_size=128,
+        pin_memory=torch.cuda.is_available(),
+    )
+    val_loader = GraphDataLoader(
+        dataset,
+        sampler=SubsetRandomSampler(val_idx),
+        batch_size=128,
+        pin_memory=torch.cuda.is_available(),
+    )
 
     # create GIN model
     in_size = dataset.dim_nfeats
@@ -140,5 +170,5 @@ if __name__ == '__main__':
     model = GIN(in_size, 16, out_size).to(device)
 
     # model training/validating
-    print('Training...')
+    print("Training...")
     train(train_loader, val_loader, device, model)

examples/pytorch/grace/aug.py

 # Data augmentation on graphs via edge dropping and feature masking
 
-import torch as th
 import numpy as np
+import torch as th
+
 import dgl
 
+
 def aug(graph, x, feat_drop_rate, edge_mask_rate):
     n_node = graph.num_nodes()
 
@@ -21,15 +23,18 @@ def aug(graph, x, feat_drop_rate, edge_mask_rate):
 
     return ng, feat
 
+
 def drop_feature(x, drop_prob):
-    drop_mask = th.empty((x.size(1),),
-                        dtype=th.float32,
-                        device=x.device).uniform_(0, 1) < drop_prob
+    drop_mask = (
+        th.empty((x.size(1),), dtype=th.float32, device=x.device).uniform_(0, 1)
+        < drop_prob
+    )
     x = x.clone()
     x[:, drop_mask] = 0
 
     return x
 
+
 def mask_edge(graph, mask_prob):
     E = graph.num_edges()
 

examples/pytorch/grace/dataset.py

-from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
+from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
+
 
 def load(name):
-    if name == 'cora':
+    if name == "cora":
         dataset = CoraGraphDataset()
-    elif name == 'citeseer':
+    elif name == "citeseer":
         dataset = CiteseerGraphDataset()
-    elif name == 'pubmed':
+    elif name == "pubmed":
         dataset = PubmedGraphDataset()
 
     graph = dataset[0]
 
-    train_mask = graph.ndata.pop('train_mask')
-    test_mask = graph.ndata.pop('test_mask')
+    train_mask = graph.ndata.pop("train_mask")
+    test_mask = graph.ndata.pop("test_mask")
 
-    feat = graph.ndata.pop('feat')
-    labels = graph.ndata.pop('label')
+    feat = graph.ndata.pop("feat")
+    labels = graph.ndata.pop("label")
 
     return graph, feat, labels, train_mask, test_mask

examples/pytorch/grace/eval.py

-'''
+"""
 Code adapted from https://github.com/CRIPAC-DIG/GRACE
 Linear evaluation on learned node embeddings
-'''
+"""
 
-import numpy as np
 import functools
 
-from sklearn.metrics import f1_score
+import numpy as np
 from sklearn.linear_model import LogisticRegression
-from sklearn.model_selection import train_test_split, GridSearchCV
+from sklearn.metrics import f1_score
+from sklearn.model_selection import GridSearchCV, train_test_split
 from sklearn.multiclass import OneVsRestClassifier
-from sklearn.preprocessing import normalize, OneHotEncoder
+from sklearn.preprocessing import OneHotEncoder, normalize
 
 
 def repeat(n_times):
@@ -22,8 +22,9 @@ def repeat(n_times):
             for key in results[0].keys():
                 values = [r[key] for r in results]
                 statistics[key] = {
-                    'mean': np.mean(values),
-                    'std': np.std(values)}
+                    "mean": np.mean(values),
+                    "std": np.std(values),
+                }
             print_statistics(statistics, f.__name__)
             return statistics
 
@@ -41,41 +42,49 @@ def prob_to_one_hot(y_pred):
 
 
 def print_statistics(statistics, function_name):
-    print(f'(E) | {function_name}:', end=' ')
+    print(f"(E) | {function_name}:", end=" ")
     for i, key in enumerate(statistics.keys()):
-        mean = statistics[key]['mean']
-        std = statistics[key]['std']
-        print(f'{key}={mean:.4f}+-{std:.4f}', end='')
+        mean = statistics[key]["mean"]
+        std = statistics[key]["std"]
+        print(f"{key}={mean:.4f}+-{std:.4f}", end="")
         if i != len(statistics.keys()) - 1:
-            print(',', end=' ')
+            print(",", end=" ")
         else:
             print()
 
 
 @repeat(3)
-def label_classification(embeddings, y, train_mask, test_mask, split='random', ratio=0.1):
+def label_classification(
+    embeddings, y, train_mask, test_mask, split="random", ratio=0.1
+):
     X = embeddings.detach().cpu().numpy()
     Y = y.detach().cpu().numpy()
     Y = Y.reshape(-1, 1)
-    onehot_encoder = OneHotEncoder(categories='auto').fit(Y)
+    onehot_encoder = OneHotEncoder(categories="auto").fit(Y)
     Y = onehot_encoder.transform(Y).toarray().astype(np.bool)
 
-    X = normalize(X, norm='l2')
+    X = normalize(X, norm="l2")
 
-    if split == 'random':
-        X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=1 - ratio)
-    elif split == 'public':
+    if split == "random":
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, Y, test_size=1 - ratio
+        )
+    elif split == "public":
         X_train = X[train_mask]
         X_test = X[test_mask]
         y_train = Y[train_mask]
         y_test = Y[test_mask]
 
-    logreg = LogisticRegression(solver='liblinear')
+    logreg = LogisticRegression(solver="liblinear")
     c = 2.0 ** np.arange(-10, 10)
 
-    clf = GridSearchCV(estimator=OneVsRestClassifier(logreg),
-                       param_grid=dict(estimator__C=c), n_jobs=8, cv=5,
-                       verbose=0)
+    clf = GridSearchCV(
+        estimator=OneVsRestClassifier(logreg),
+        param_grid=dict(estimator__C=c),
+        n_jobs=8,
+        cv=5,
+        verbose=0,
+    )
     clf.fit(X_train, y_train)
 
     y_pred = clf.predict_proba(X_test)
@@ -84,7 +93,4 @@ def label_classification(embeddings, y, train_mask, test_mask, split='random', r
     micro = f1_score(y_test, y_pred, average="micro")
     macro = f1_score(y_test, y_pred, average="macro")
 
-    return {
-        'F1Mi': micro,
-        'F1Ma': macro
-    }
+    return {"F1Mi": micro, "F1Ma": macro}

examples/pytorch/grace/main.py

 import argparse
-from model import Grace
-from aug import aug
-from dataset import load
+import warnings
 
 import numpy as np
 import torch as th
 import torch.nn as nn
-
+from aug import aug
+from dataset import load
 from eval import label_classification
-import warnings
+from model import Grace
 
-warnings.filterwarnings('ignore')
+warnings.filterwarnings("ignore")
 
 
 def count_parameters(model):
-    return sum([np.prod(p.size()) for p in model.parameters() if p.requires_grad])
+    return sum(
+        [np.prod(p.size()) for p in model.parameters() if p.requires_grad]
+    )
 
 
 parser = argparse.ArgumentParser()
-parser.add_argument('--dataname', type=str, default='cora')
-parser.add_argument('--gpu', type=int, default=0)
-parser.add_argument('--split', type=str, default='random')
-
-parser.add_argument('--epochs', type=int, default=500, help='Number of training periods.')
-parser.add_argument('--lr', type=float, default=0.001, help='Learning rate.')
-parser.add_argument('--wd', type=float, default=1e-5, help='Weight decay.')
-parser.add_argument('--temp', type=float, default=1.0, help='Temperature.')
-
-parser.add_argument('--act_fn', type=str, default='relu')
-
-parser.add_argument("--hid_dim", type=int, default=256, help='Hidden layer dim.')
-parser.add_argument("--out_dim", type=int, default=256, help='Output layer dim.')
-
-parser.add_argument("--num_layers", type=int, default=2, help='Number of GNN layers.')
-parser.add_argument('--der1', type=float, default=0.2, help='Drop edge ratio of the 1st augmentation.')
-parser.add_argument('--der2', type=float, default=0.2, help='Drop edge ratio of the 2nd augmentation.')
-parser.add_argument('--dfr1', type=float, default=0.2, help='Drop feature ratio of the 1st augmentation.')
-parser.add_argument('--dfr2', type=float, default=0.2, help='Drop feature ratio of the 2nd augmentation.')
+parser.add_argument("--dataname", type=str, default="cora")
+parser.add_argument("--gpu", type=int, default=0)
+parser.add_argument("--split", type=str, default="random")
+
+parser.add_argument(
+    "--epochs", type=int, default=500, help="Number of training periods."
+)
+parser.add_argument("--lr", type=float, default=0.001, help="Learning rate.")
+parser.add_argument("--wd", type=float, default=1e-5, help="Weight decay.")
+parser.add_argument("--temp", type=float, default=1.0, help="Temperature.")
+
+parser.add_argument("--act_fn", type=str, default="relu")
+
+parser.add_argument(
+    "--hid_dim", type=int, default=256, help="Hidden layer dim."
+)
+parser.add_argument(
+    "--out_dim", type=int, default=256, help="Output layer dim."
+)
+
+parser.add_argument(
+    "--num_layers", type=int, default=2, help="Number of GNN layers."
+)
+parser.add_argument(
+    "--der1",
+    type=float,
+    default=0.2,
+    help="Drop edge ratio of the 1st augmentation.",
+)
+parser.add_argument(
+    "--der2",
+    type=float,
+    default=0.2,
+    help="Drop edge ratio of the 2nd augmentation.",
+)
+parser.add_argument(
+    "--dfr1",
+    type=float,
+    default=0.2,
+    help="Drop feature ratio of the 1st augmentation.",
+)
+parser.add_argument(
+    "--dfr2",
+    type=float,
+    default=0.2,
+    help="Drop feature ratio of the 2nd augmentation.",
+)
 
 args = parser.parse_args()
 
 if args.gpu != -1 and th.cuda.is_available():
-    args.device = 'cuda:{}'.format(args.gpu)
+    args.device = "cuda:{}".format(args.gpu)
 else:
-    args.device = 'cpu'
+    args.device = "cpu"
 
-if __name__ == '__main__':
+if __name__ == "__main__":
 
     # Step 1: Load hyperparameters =================================================================== #
     lr = args.lr
@@ -53,7 +82,7 @@ if __name__ == '__main__':
     out_dim = args.out_dim
 
     num_layers = args.num_layers
-    act_fn = ({'relu': nn.ReLU(), 'prelu': nn.PReLU()})[args.act_fn]
+    act_fn = ({"relu": nn.ReLU(), "prelu": nn.PReLU()})[args.act_fn]
 
     drop_edge_rate_1 = args.der1
     drop_edge_rate_2 = args.der2
@@ -71,7 +100,7 @@ if __name__ == '__main__':
     # Step 3: Create model =================================================================== #
     model = Grace(in_dim, hid_dim, out_dim, num_layers, act_fn, temp)
     model = model.to(args.device)
-    print(f'# params: {count_parameters(model)}')
+    print(f"# params: {count_parameters(model)}")
 
     optimizer = th.optim.Adam(model.parameters(), lr=lr, weight_decay=wd)
 
@@ -92,7 +121,7 @@ if __name__ == '__main__':
         loss.backward()
         optimizer.step()
 
-        print(f'Epoch={epoch:03d}, loss={loss.item():.4f}')
+        print(f"Epoch={epoch:03d}, loss={loss.item():.4f}")
 
     # Step 5: Linear evaluation ============================================================== #
     print("=== Final ===")
@@ -102,5 +131,7 @@ if __name__ == '__main__':
     feat = feat.to(args.device)
     embeds = model.get_embedding(graph, feat)
 
-    '''Evaluation Embeddings  '''
-    label_classification(embeds, labels, train_mask, test_mask, split=args.split)
+    """Evaluation Embeddings  """
+    label_classification(
\ No newline at end of file
+        embeds, labels, train_mask, test_mask, split=args.split
+    )

examples/pytorch/grace/model.py

 import torch as th
 import torch.nn as nn
 import torch.nn.functional as F
+
 from dgl.nn import GraphConv
 
 
 # Multi-layer Graph Convolutional Networks
 class GCN(nn.Module):
-    def __init__(self, in_dim, out_dim, act_fn, num_layers = 2):
+    def __init__(self, in_dim, out_dim, act_fn, num_layers=2):
         super(GCN, self).__init__()
 
         assert num_layers >= 2
@@ -26,6 +27,7 @@ class GCN(nn.Module):
 
         return feat
 
+
 # Multi-layer(2-layer) Perceptron
 class MLP(nn.Module):
     def __init__(self, in_dim, out_dim):
@@ -56,6 +58,7 @@ class Grace(nn.Module):
     temp: float
         Temperature constant.
     """
+
     def __init__(self, in_dim, hid_dim, out_dim, num_layers, act_fn, temp):
         super(Grace, self).__init__()
         self.encoder = GCN(in_dim, hid_dim, act_fn, num_layers)

examples/pytorch/grand/main.py

 import argparse
+import warnings
+
 import numpy as np
 import torch as th
-import torch.optim as optim
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.optim as optim
+from model import GRAND
 
 import dgl
-from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
+from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
 
-from model import GRAND
+warnings.filterwarnings("ignore")
 
-import warnings
-warnings.filterwarnings('ignore')
 
 def argument():
 
-    parser = argparse.ArgumentParser(description='GRAND')
+    parser = argparse.ArgumentParser(description="GRAND")
 
     # data source params
-    parser.add_argument('--dataname', type=str, default='cora', help='Name of dataset.')
+    parser.add_argument(
+        "--dataname", type=str, default="cora", help="Name of dataset."
+    )
     # cuda params
-    parser.add_argument('--gpu', type=int, default=-1, help='GPU index. Default: -1, using CPU.')
+    parser.add_argument(
+        "--gpu", type=int, default=-1, help="GPU index. Default: -1, using CPU."
+    )
     # training params
-    parser.add_argument('--epochs', type=int, default=200, help='Training epochs.')
-    parser.add_argument('--early_stopping', type=int, default=200, help='Patient epochs to wait before early stopping.')
-    parser.add_argument('--lr', type=float, default=0.01, help='Learning rate.')
-    parser.add_argument('--weight_decay', type=float, default=5e-4, help='L2 reg.')
+    parser.add_argument(
+        "--epochs", type=int, default=200, help="Training epochs."
+    )
+    parser.add_argument(
+        "--early_stopping",
+        type=int,
+        default=200,
+        help="Patient epochs to wait before early stopping.",
+    )
+    parser.add_argument("--lr", type=float, default=0.01, help="Learning rate.")
+    parser.add_argument(
+        "--weight_decay", type=float, default=5e-4, help="L2 reg."
+    )
     # model params
-    parser.add_argument("--hid_dim", type=int, default=32, help='Hidden layer dimensionalities.')
-    parser.add_argument('--dropnode_rate', type=float, default=0.5,
-                        help='Dropnode rate (1 - keep probability).')
-    parser.add_argument('--input_droprate', type=float, default=0.0,
-                    help='dropout rate of input layer')
-    parser.add_argument('--hidden_droprate', type=float, default=0.0,
-                    help='dropout rate of hidden layer')
-    parser.add_argument('--order', type=int, default=8, help='Propagation step')
-    parser.add_argument('--sample', type=int, default=4, help='Sampling times of dropnode')
-    parser.add_argument('--tem', type=float, default=0.5, help='Sharpening temperature')
-    parser.add_argument('--lam', type=float, default=1., help='Coefficient of consistency regularization')
-    parser.add_argument('--use_bn', action='store_true', default=False, help='Using Batch Normalization')
+    parser.add_argument(
+        "--hid_dim", type=int, default=32, help="Hidden layer dimensionalities."
+    )
+    parser.add_argument(
+        "--dropnode_rate",
+        type=float,
+        default=0.5,
+        help="Dropnode rate (1 - keep probability).",
+    )
+    parser.add_argument(
+        "--input_droprate",
+        type=float,
+        default=0.0,
+        help="dropout rate of input layer",
+    )
+    parser.add_argument(
+        "--hidden_droprate",
+        type=float,
+        default=0.0,
+        help="dropout rate of hidden layer",
+    )
+    parser.add_argument("--order", type=int, default=8, help="Propagation step")
+    parser.add_argument(
+        "--sample", type=int, default=4, help="Sampling times of dropnode"
+    )
+    parser.add_argument(
+        "--tem", type=float, default=0.5, help="Sharpening temperature"
+    )
+    parser.add_argument(
+        "--lam",
+        type=float,
+        default=1.0,
+        help="Coefficient of consistency regularization",
+    )
+    parser.add_argument(
+        "--use_bn",
+        action="store_true",
+        default=False,
+        help="Using Batch Normalization",
+    )
 
     args = parser.parse_args()
 
     # check cuda
     if args.gpu != -1 and th.cuda.is_available():
-        args.device = 'cuda:{}'.format(args.gpu)
+        args.device = "cuda:{}".format(args.gpu)
     else:
-        args.device = 'cpu'
+        args.device = "cpu"
 
     return args
 
+
 def consis_loss(logps, temp, lam):
     ps = [th.exp(p) for p in logps]
-    ps = th.stack(ps, dim = 2)
+    ps = th.stack(ps, dim=2)
 
-    avg_p = th.mean(ps, dim = 2)
-    sharp_p = (th.pow(avg_p, 1./temp) / th.sum(th.pow(avg_p, 1./temp), dim=1, keepdim=True)).detach()
+    avg_p = th.mean(ps, dim=2)
+    sharp_p = (
+        th.pow(avg_p, 1.0 / temp)
+        / th.sum(th.pow(avg_p, 1.0 / temp), dim=1, keepdim=True)
+    ).detach()
 
     sharp_p = sharp_p.unsqueeze(2)
-    loss = th.mean(th.sum(th.pow(ps - sharp_p, 2), dim = 1, keepdim=True))
+    loss = th.mean(th.sum(th.pow(ps - sharp_p, 2), dim=1, keepdim=True))
 
     loss = lam * loss
     return loss
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
 
     # Step 1: Prepare graph data and retrieve train/validation/test index ============================= #
     # Load from DGL dataset
     args = argument()
     print(args)
 
-    if args.dataname == 'cora':
+    if args.dataname == "cora":
         dataset = CoraGraphDataset()
-    elif args.dataname == 'citeseer':
+    elif args.dataname == "citeseer":
         dataset = CiteseerGraphDataset()
-    elif args.dataname == 'pubmed':
+    elif args.dataname == "pubmed":
         dataset = PubmedGraphDataset()
 
     graph = dataset[0]
@@ -86,33 +133,42 @@ if __name__ == '__main__':
     n_classes = dataset.num_classes
 
     # retrieve labels of ground truth
-    labels = graph.ndata.pop('label').to(device).long()
+    labels = graph.ndata.pop("label").to(device).long()
 
     # Extract node features
-    feats = graph.ndata.pop('feat').to(device)
+    feats = graph.ndata.pop("feat").to(device)
     n_features = feats.shape[-1]
 
     # retrieve masks for train/validation/test
-    train_mask = graph.ndata.pop('train_mask')
-    val_mask = graph.ndata.pop('val_mask')
-    test_mask = graph.ndata.pop('test_mask')
+    train_mask = graph.ndata.pop("train_mask")
+    val_mask = graph.ndata.pop("val_mask")
+    test_mask = graph.ndata.pop("test_mask")
 
     train_idx = th.nonzero(train_mask, as_tuple=False).squeeze().to(device)
     val_idx = th.nonzero(val_mask, as_tuple=False).squeeze().to(device)
     test_idx = th.nonzero(test_mask, as_tuple=False).squeeze().to(device)
 
     # Step 2: Create model =================================================================== #
-    model = GRAND(n_features, args.hid_dim, n_classes, args.sample, args.order,
-                  args.dropnode_rate, args.input_droprate, 
-                  args.hidden_droprate, args.use_bn)
-
+    model = GRAND(
+        n_features,
+        args.hid_dim,
+        n_classes,
+        args.sample,
+        args.order,
+        args.dropnode_rate,
+        args.input_droprate,
+        args.hidden_droprate,
+        args.use_bn,
+    )
 
     model = model.to(args.device)
     graph = graph.to(args.device)
 
     # Step 3: Create training components ===================================================== #
     loss_fn = nn.NLLLoss()
-    opt = optim.Adam(model.parameters(), lr = args.lr, weight_decay = args.weight_decay)
+    opt = optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+    )
 
     loss_best = np.inf
     acc_best = 0
@@ -120,7 +176,7 @@ if __name__ == '__main__':
     # Step 4: training epoches =============================================================== #
     for epoch in range(args.epochs):
 
-        ''' Training '''
+        """Training"""
         model.train()
 
         loss_sup = 0
@@ -130,56 +186,68 @@ if __name__ == '__main__':
         for k in range(args.sample):
             loss_sup += F.nll_loss(logits[k][train_idx], labels[train_idx])
 
-        loss_sup = loss_sup/args.sample
+        loss_sup = loss_sup / args.sample
 
         # calculate consistency loss
         loss_consis = consis_loss(logits, args.tem, args.lam)
 
         loss_train = loss_sup + loss_consis
-        acc_train = th.sum(logits[0][train_idx].argmax(dim=1) == labels[train_idx]).item() / len(train_idx)
+        acc_train = th.sum(
+            logits[0][train_idx].argmax(dim=1) == labels[train_idx]
+        ).item() / len(train_idx)
 
         # backward
         opt.zero_grad()
         loss_train.backward()
         opt.step()
 
-        ''' Validating '''
+        """ Validating """
         model.eval()
         with th.no_grad():
 
             val_logits = model(graph, feats, False)
 
             loss_val = F.nll_loss(val_logits[val_idx], labels[val_idx])
-            acc_val = th.sum(val_logits[val_idx].argmax(dim=1) == labels[val_idx]).item() / len(val_idx)
+            acc_val = th.sum(
+                val_logits[val_idx].argmax(dim=1) == labels[val_idx]
+            ).item() / len(val_idx)
 
             # Print out performance
-            print("In epoch {}, Train Acc: {:.4f} | Train Loss: {:.4f} ,Val Acc: {:.4f} | Val Loss: {:.4f}".
-              format(epoch, acc_train, loss_train.item(), acc_val, loss_val.item()))
+            print(
+                "In epoch {}, Train Acc: {:.4f} | Train Loss: {:.4f} ,Val Acc: {:.4f} | Val Loss: {:.4f}".format(
+                    epoch,
+                    acc_train,
+                    loss_train.item(),
+                    acc_val,
+                    loss_val.item(),
+                )
+            )
 
             # set early stopping counter
             if loss_val < loss_best or acc_val > acc_best:
                 if loss_val < loss_best:
                     best_epoch = epoch
-                    th.save(model.state_dict(), args.dataname +'.pkl')
+                    th.save(model.state_dict(), args.dataname + ".pkl")
                 no_improvement = 0
                 loss_best = min(loss_val, loss_best)
                 acc_best = max(acc_val, acc_best)
             else:
                 no_improvement += 1
                 if no_improvement == args.early_stopping:
-                    print('Early stopping.')
+                    print("Early stopping.")
                     break
 
     print("Optimization Finished!")
 
-    print('Loading {}th epoch'.format(best_epoch))
-    model.load_state_dict(th.load(args.dataname +'.pkl'))
+    print("Loading {}th epoch".format(best_epoch))
+    model.load_state_dict(th.load(args.dataname + ".pkl"))
 
-    ''' Testing '''
+    """ Testing """
     model.eval()
 
     test_logits = model(graph, feats, False)
-    test_acc = th.sum(test_logits[test_idx].argmax(dim=1) == labels[test_idx]).item() / len(test_idx)
+    test_acc = th.sum(
+        test_logits[test_idx].argmax(dim=1) == labels[test_idx]
+    ).item() / len(test_idx)
 
     print("Test Acc: {:.4f}".format(test_acc))
-

examples/pytorch/grand/model.py

 import numpy as np
 import torch as th
 import torch.nn as nn
-import dgl.function as fn
 import torch.nn.functional as F
 
+import dgl.function as fn
+
+
 def drop_node(feats, drop_rate, training):
 
     n = feats.shape[0]
@@ -11,20 +13,23 @@ def drop_node(feats, drop_rate, training):
 
     if training:
 
-        masks = th.bernoulli(1. - drop_rates).unsqueeze(1)
+        masks = th.bernoulli(1.0 - drop_rates).unsqueeze(1)
         feats = masks.to(feats.device) * feats
 
     else:
-        feats = feats * (1. - drop_rate)
+        feats = feats * (1.0 - drop_rate)
 
     return feats
 
+
 class MLP(nn.Module):
-    def __init__(self, nfeat, nhid, nclass, input_droprate, hidden_droprate, use_bn =False):
+    def __init__(
+        self, nfeat, nhid, nclass, input_droprate, hidden_droprate, use_bn=False
+    ):
         super(MLP, self).__init__()
 
-        self.layer1 = nn.Linear(nfeat, nhid, bias = True)
-        self.layer2 = nn.Linear(nhid, nclass, bias = True)
+        self.layer1 = nn.Linear(nfeat, nhid, bias=True)
+        self.layer2 = nn.Linear(nhid, nclass, bias=True)
 
         self.input_dropout = nn.Dropout(input_droprate)
         self.hidden_dropout = nn.Dropout(hidden_droprate)
@@ -52,7 +57,7 @@ class MLP(nn.Module):
 
 
 def GRANDConv(graph, feats, order):
-    '''
+    """
     Parameters
     -----------
     graph: dgl.Graph
@@ -61,27 +66,28 @@ def GRANDConv(graph, feats, order):
         Node features
     order: int
         Propagation Steps
-    '''
+    """
     with graph.local_scope():
 
-        ''' Calculate Symmetric normalized adjacency matrix   \hat{A} '''
+        """Calculate Symmetric normalized adjacency matrix   \hat{A}"""
         degs = graph.in_degrees().float().clamp(min=1)
         norm = th.pow(degs, -0.5).to(feats.device).unsqueeze(1)
 
-        graph.ndata['norm'] = norm
-        graph.apply_edges(fn.u_mul_v('norm', 'norm', 'weight')) 
+        graph.ndata["norm"] = norm
+        graph.apply_edges(fn.u_mul_v("norm", "norm", "weight"))
 
-        ''' Graph Conv '''
+        """ Graph Conv """
         x = feats
-        y = 0+feats
+        y = 0 + feats
 
         for i in range(order):
-            graph.ndata['h'] = x
-            graph.update_all(fn.u_mul_e('h', 'weight', 'm'), fn.sum('m', 'h'))
-            x = graph.ndata.pop('h')
+            graph.ndata["h"] = x
+            graph.update_all(fn.u_mul_e("h", "weight", "m"), fn.sum("m", "h"))
+            x = graph.ndata.pop("h")
             y.add_(x)
 
-    return y /(order + 1)
+    return y / (order + 1)
+
 
 class GRAND(nn.Module):
     r"""
@@ -108,16 +114,19 @@ class GRAND(nn.Module):
         If True, use batch normalization.
 
     """
-    def __init__(self,
+
+    def __init__(
+        self,
         in_dim,
         hid_dim,
         n_class,
-                 S = 1,
-                 K = 3,
+        S=1,
+        K=3,
         node_dropout=0.0,
-                 input_droprate = 0.0, 
-                 hidden_droprate = 0.0,
-                 batchnorm=False):
+        input_droprate=0.0,
+        hidden_droprate=0.0,
+        batchnorm=False,
+    ):
 
         super(GRAND, self).__init__()
         self.in_dim = in_dim
@@ -126,12 +135,14 @@ class GRAND(nn.Module):
         self.K = K
         self.n_class = n_class
 
-        self.mlp = MLP(in_dim, hid_dim, n_class, input_droprate, hidden_droprate, batchnorm)
+        self.mlp = MLP(
+            in_dim, hid_dim, n_class, input_droprate, hidden_droprate, batchnorm
+        )
 
         self.dropout = node_dropout
         self.node_dropout = nn.Dropout(node_dropout)
 
-    def forward(self, graph, feats, training = True):
+    def forward(self, graph, feats, training=True):
 
         X = feats
         S = self.S
@@ -141,13 +152,13 @@ class GRAND(nn.Module):
             for s in range(S):
                 drop_feat = drop_node(X, self.dropout, True)  # Drop node
                 feat = GRANDConv(graph, drop_feat, self.K)  # Graph Convolution
-                output_list.append(th.log_softmax(self.mlp(feat), dim=-1))  # Prediction
+                output_list.append(
+                    th.log_softmax(self.mlp(feat), dim=-1)
+                )  # Prediction
 
             return output_list
         else:  # Inference Mode
             drop_feat = drop_node(X, self.dropout, False)
             X = GRANDConv(graph, drop_feat, self.K)
 
-            return th.log_softmax(self.mlp(X), dim = -1)
-        
-        
+            return th.log_softmax(self.mlp(X), dim=-1)

examples/pytorch/graph_matching/examples.py

-from ged import graph_edit_distance
-import dgl
 import numpy as np
+from ged import graph_edit_distance
 
+import dgl
 
-src1 = [0, 1, 2, 3, 4, 5];
-dst1 = [1, 2, 3, 4, 5, 6];
+src1 = [0, 1, 2, 3, 4, 5]
+dst1 = [1, 2, 3, 4, 5, 6]
 
-src2 = [0, 1, 3, 4, 5];
-dst2 = [1, 2, 4, 5, 6];
+src2 = [0, 1, 3, 4, 5]
+dst2 = [1, 2, 4, 5, 6]
 
 
 G1 = dgl.DGLGraph((src1, dst1))
@@ -15,55 +15,73 @@ G2 = dgl.DGLGraph((src2, dst2))
 
 
 # Exact edit distance with astar search
-distance, node_mapping, edge_mapping = graph_edit_distance(G1, G1, algorithm='astar')
+distance, node_mapping, edge_mapping = graph_edit_distance(
+    G1, G1, algorithm="astar"
+)
 print(distance)  # 0.0
-distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='astar')
+distance, node_mapping, edge_mapping = graph_edit_distance(
+    G1, G2, algorithm="astar"
+)
 print(distance)  # 1.0
 
 # With user-input cost matrices
-node_substitution_cost = np.empty((G1.number_of_nodes(), G2.number_of_nodes()));
-G1_node_deletion_cost = np.empty(G1.number_of_nodes());
-G2_node_insertion_cost = np.empty(G2.number_of_nodes());
+node_substitution_cost = np.empty((G1.number_of_nodes(), G2.number_of_nodes()))
+G1_node_deletion_cost = np.empty(G1.number_of_nodes())
+G2_node_insertion_cost = np.empty(G2.number_of_nodes())
 
-edge_substitution_cost = np.empty((G1.number_of_edges(), G2.number_of_edges()));
-G1_edge_deletion_cost = np.empty(G1.number_of_edges());
-G2_edge_insertion_cost = np.empty(G2.number_of_edges());
+edge_substitution_cost = np.empty((G1.number_of_edges(), G2.number_of_edges()))
+G1_edge_deletion_cost = np.empty(G1.number_of_edges())
+G2_edge_insertion_cost = np.empty(G2.number_of_edges())
 
 # Node substitution cost of 0 when node-ids are same, else 1
-node_substitution_cost.fill(1.0);
+node_substitution_cost.fill(1.0)
 for i in range(G1.number_of_nodes()):
     for j in range(G2.number_of_nodes()):
-        node_substitution_cost[i,j] = 0.0;
+        node_substitution_cost[i, j] = 0.0
 
 # Node insertion/deletion cost of 1
-G1_node_deletion_cost.fill(1.0);
-G2_node_insertion_cost.fill(1.0);
+G1_node_deletion_cost.fill(1.0)
+G2_node_insertion_cost.fill(1.0)
 
 # Edge substitution cost of 0
-edge_substitution_cost.fill(0.0);
+edge_substitution_cost.fill(0.0)
 
 # Edge insertion/deletion cost of 0.5
-G1_edge_deletion_cost.fill(0.5);
-G2_edge_insertion_cost.fill(0.5);
+G1_edge_deletion_cost.fill(0.5)
+G2_edge_insertion_cost.fill(0.5)
 
-distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, \
-                                                          node_substitution_cost, edge_substitution_cost, \
-                                                          G1_node_deletion_cost, G2_node_insertion_cost, \
-                                                          G1_edge_deletion_cost, G2_edge_insertion_cost, \
-                                                          algorithm="astar")
+distance, node_mapping, edge_mapping = graph_edit_distance(
+    G1,
+    G2,
+    node_substitution_cost,
+    edge_substitution_cost,
+    G1_node_deletion_cost,
+    G2_node_insertion_cost,
+    G1_edge_deletion_cost,
+    G2_edge_insertion_cost,
+    algorithm="astar",
+)
 
-print(distance) #0.5
+print(distance)  # 0.5
 
 
 # Approximate edit distance with beam search, it is more than or equal to the exact edit distance
-distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='beam', max_beam_size=2)
+distance, node_mapping, edge_mapping = graph_edit_distance(
+    G1, G2, algorithm="beam", max_beam_size=2
+)
 print(distance)  # 3.0
 
 # Approximate edit distance with bipartite heuristic, it is more than or equal to the exact edit distance
-distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='bipartite')
-print(distance) # 9.0, can be different as multiple solutions possible for the intermediate LAP used in this approximation
+distance, node_mapping, edge_mapping = graph_edit_distance(
+    G1, G2, algorithm="bipartite"
+)
+print(
+    distance
+)  # 9.0, can be different as multiple solutions possible for the intermediate LAP used in this approximation
 
 
 # Approximate edit distance with hausdorff heuristic, it is less than or equal to the exact edit distance
-distance, node_mapping, edge_mapping = graph_edit_distance(G1, G2, algorithm='hausdorff')
+distance, node_mapping, edge_mapping = graph_edit_distance(
+    G1, G2, algorithm="hausdorff"
+)
 print(distance)  # 0.0

examples/pytorch/graphsage/advanced/model.py

+import sklearn.linear_model as lm
+import sklearn.metrics as skm
 import torch as th
-import torch.nn as nn
 import torch.functional as F
+import torch.nn as nn
+import tqdm
+
 import dgl
 import dgl.nn as dglnn
-import sklearn.linear_model as lm
-import sklearn.metrics as skm
-import tqdm
+
 
 class SAGE(nn.Module):
-    def __init__(self, in_feats, n_hidden, n_classes, n_layers, activation, dropout):
+    def __init__(
+        self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
+    ):
         super().__init__()
         self.init(in_feats, n_hidden, n_classes, n_layers, activation, dropout)
 
-    def init(self, in_feats, n_hidden, n_classes, n_layers, activation, dropout):
+    def init(
+        self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
+    ):
         self.n_layers = n_layers
         self.n_hidden = n_hidden
         self.n_classes = n_classes
         self.layers = nn.ModuleList()
         if n_layers > 1:
-            self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
+            self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, "mean"))
             for i in range(1, n_layers - 1):
-                self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
-            self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, 'mean'))
+                self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, "mean"))
+            self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
         else:
-            self.layers.append(dglnn.SAGEConv(in_feats, n_classes, 'mean'))
+            self.layers.append(dglnn.SAGEConv(in_feats, n_classes, "mean"))
         self.dropout = nn.Dropout(dropout)
         self.activation = activation
 
@@ -51,7 +57,10 @@ class SAGE(nn.Module):
         # on each layer are of course splitted in batches.
         # TODO: can we standardize this?
         for l, layer in enumerate(self.layers):
-            y = th.zeros(g.num_nodes(), self.n_hidden if l != len(self.layers) - 1 else self.n_classes)
+            y = th.zeros(
+                g.num_nodes(),
+                self.n_hidden if l != len(self.layers) - 1 else self.n_classes,
+            )
 
             sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
             dataloader = dgl.dataloading.DataLoader(
@@ -62,7 +71,8 @@ class SAGE(nn.Module):
                 batch_size=batch_size,
                 shuffle=False,
                 drop_last=False,
-                num_workers=num_workers)
+                num_workers=num_workers,
+            )
 
             for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
                 block = blocks[0]
@@ -79,6 +89,7 @@ class SAGE(nn.Module):
             x = y
         return y
 
+
 def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids):
     """
     Compute the accuracy of prediction given the labels.
@@ -94,10 +105,10 @@ def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids):
 
     emb = (emb - emb.mean(0, keepdims=True)) / emb.std(0, keepdims=True)
 
-    lr = lm.LogisticRegression(multi_class='multinomial', max_iter=10000)
+    lr = lm.LogisticRegression(multi_class="multinomial", max_iter=10000)
     lr.fit(emb[train_nids], train_labels)
 
     pred = lr.predict(emb)
-    f1_micro_eval = skm.f1_score(val_labels, pred[val_nids], average='micro')
-    f1_micro_test = skm.f1_score(test_labels, pred[test_nids], average='micro')
+    f1_micro_eval = skm.f1_score(val_labels, pred[val_nids], average="micro")
+    f1_micro_test = skm.f1_score(test_labels, pred[test_nids], average="micro")
     return f1_micro_eval, f1_micro_test