[Benchmark] Adding inference benchmark suite

benchmark/inference/hetero_gat.py

@@ -0,0 +1,45 @@
+import torch
+from tqdm import tqdm
+
+from torch_geometric.nn import GATConv, to_hetero
+
+
+class HeteroGAT(torch.nn.Module):
+    def __init__(self, metadata, hidden_channels, num_layers, output_channels,
+                 num_heads):
+        super().__init__()
+        self.model = to_hetero(
+            GATForHetero(hidden_channels, num_layers, output_channels,
+                         num_heads), metadata)
+        self.training = False
+
+    @torch.no_grad()
+    def inference(self, loader, device, progress_bar=False):
+        self.model.eval()
+        if progress_bar:
+            loader = tqdm(loader)
+        for batch in loader:
+            batch = batch.to(device)
+            self.model(batch.x_dict, batch.edge_index_dict)
+
+
+class GATForHetero(torch.nn.Module):
+    def __init__(self, hidden_channels, num_layers, out_channels, heads):
+        super().__init__()
+        self.convs = torch.nn.ModuleList()
+        self.convs.append(
+            GATConv((-1, -1), hidden_channels, heads=heads,
+                    add_self_loops=False))
+        for _ in range(num_layers - 2):
+            self.convs.append(
+                GATConv((-1, -1), hidden_channels, heads=heads,
+                        add_self_loops=False))
+        self.convs.append(
+            GATConv((-1, -1), out_channels, heads=heads, add_self_loops=False))
+
+    def forward(self, x, edge_index):
+        for i, conv in enumerate(self.convs):
+            x = conv(x, edge_index)
+            if i < len(self.convs) - 1:
+                x = x.relu_()
+        return x

benchmark/inference/hetero_sage.py

@@ -0,0 +1,39 @@
+import torch
+from tqdm import tqdm
+
+from torch_geometric.nn import SAGEConv, to_hetero
+
+
+class HeteroGraphSAGE(torch.nn.Module):
+    def __init__(self, metadata, hidden_channels, num_layers, output_channels):
+        super().__init__()
+        self.model = to_hetero(
+            SAGEForHetero(hidden_channels, num_layers, output_channels),
+            metadata)
+        self.training = False
+
+    @torch.no_grad()
+    def inference(self, loader, device, progress_bar=False):
+        self.model.eval()
+        if progress_bar:
+            loader = tqdm(loader)
+        for batch in loader:
+            batch = batch.to(device)
+            self.model(batch.x_dict, batch.edge_index_dict)
+
+
+class SAGEForHetero(torch.nn.Module):
+    def __init__(self, hidden_channels, num_layers, out_channels):
+        super().__init__()
+        self.convs = torch.nn.ModuleList()
+        self.convs.append(SAGEConv((-1, -1), hidden_channels))
+        for i in range(num_layers - 2):
+            self.convs.append(SAGEConv((-1, -1), hidden_channels))
+        self.convs.append(SAGEConv((-1, -1), out_channels))
+
+    def forward(self, x, edge_index):
+        for i, conv in enumerate(self.convs):
+            x = conv(x, edge_index)
+            if i < len(self.convs) - 1:
+                x = x.relu_()
+        return x

benchmark/inference/inference_benchmark.py

@@ -0,0 +1,126 @@
+import argparse
+import copy
+from timeit import default_timer
+
+import torch
+from utils import get_dataset, get_model
+
+from torch_geometric.loader import NeighborLoader
+from torch_geometric.nn import PNAConv
+
+supported_sets = {
+    'ogbn-mag': ['rgat', 'rgcn'],
+    'reddit': ['edge_conv', 'gat', 'gcn', 'pna_conv'],
+    'ogbn-products': ['edge_conv', 'gat', 'gcn', 'pna_conv'],
+}
+
+
+def run(args: argparse.ArgumentParser) -> None:
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    progress_bar = True
+
+    print('BENCHMARK STARTS')
+    for dataset_name in args.datasets:
+        print(f'Dataset: {dataset_name}')
+        dataset = get_dataset(dataset_name, args.root)
+
+        hetero = True if dataset_name == 'ogbn-mag' else False
+        mask = ('paper', None) if hetero else None
+        degree = None
+
+        data = dataset[0].to(device)
+        inputs_channels = data.x_dict['paper'].size(
+            -1) if hetero else dataset.num_features
+
+        for model_name in args.models:
+            if model_name not in supported_sets[dataset_name]:
+                print(f'Configuration of {dataset_name} + {model_name} '
+                      f'not supported. Skipping.')
+                continue
+            print(f'Evaluation bench for {model_name}:')
+
+            for batch_size in args.eval_batch_sizes:
+                if not hetero:
+                    subgraph_loader = NeighborLoader(
+                        copy.copy(data),
+                        num_neighbors=[-1],
+                        input_nodes=mask,
+                        batch_size=batch_size,
+                        shuffle=False,
+                        num_workers=args.num_workers,
+                    )
+                    subgraph_loader.data.n_id = torch.arange(data.num_nodes)
+
+                for layers in args.num_layers:
+                    if hetero:
+                        subgraph_loader = NeighborLoader(
+                            copy.copy(data),
+                            num_neighbors=[args.hetero_num_neighbors] * layers,
+                            input_nodes=mask,
+                            batch_size=batch_size,
+                            shuffle=False,
+                            num_workers=args.num_workers,
+                        )
+                        subgraph_loader.data.n_id = torch.arange(
+                            data.num_nodes)
+
+                    for hidden_channels in args.num_hidden_channels:
+                        print(
+                            '-----------------------------------------------')
+                        print(f'Batch size={batch_size}, '
+                              f'Layers amount={layers}, '
+                              f'Hidden features size={hidden_channels}')
+                        params = {
+                            'inputs_channels': inputs_channels,
+                            'hidden_channels': hidden_channels,
+                            'output_channels': dataset.num_classes,
+                            'num_heads': args.num_heads,
+                            'num_layers': layers,
+                        }
+
+                        if model_name == 'pna_conv':
+                            if degree is None:
+                                degree = PNAConv.get_degree_histogram(
+                                    subgraph_loader)
+                                print(f'Calculated degree for {dataset_name}.')
+                            params['degree'] = degree
+
+                        model = get_model(
+                            model_name, params,
+                            metadata=data.metadata() if hetero else None)
+                        model = model.to(device)
+                        model.training = False
+
+                        start = default_timer()
+                        model.inference(subgraph_loader, device, progress_bar)
+                        stop = default_timer()
+                        print(f'Inference time={stop-start:.3f}\n')
+
+
+if __name__ == '__main__':
+    argparser = argparse.ArgumentParser('GNN inference benchmark')
+    argparser.add_argument('--datasets', nargs='+',
+                           default=['ogbn-mag', 'ogbn-products',
+                                    'reddit'], type=str)
+    argparser.add_argument(
+        '--models', nargs='+',
+        default=['edge_conv', 'gat', 'gcn', 'pna_conv', 'rgat',
+                 'rgcn'], type=str)
+    argparser.add_argument('--root', default='../../data', type=str)
+    argparser.add_argument('--eval-batch-sizes', nargs='+',
+                           default=[512, 1024, 2048, 4096, 8192], type=int)
+    argparser.add_argument('--num-layers', nargs='+', default=[2, 3], type=int)
+    argparser.add_argument('--num-hidden-channels', nargs='+',
+                           default=[64, 128, 256], type=int)
+    argparser.add_argument(
+        '--num-heads', default=2, type=int,
+        help='number of hidden attention heads, applies only for gat and rgat')
+    argparser.add_argument(
+        '--hetero-num-neighbors', default=-1, type=int,
+        help='number of neighbors to sample per layer for hetero workloads')
+    argparser.add_argument('--num-workers', default=2, type=int)
+
+    args = argparser.parse_args()
+
+    run(args)

benchmark/inference/utils.py

@@ -0,0 +1,72 @@
+import os.path as osp
+
+from hetero_gat import HeteroGAT
+from hetero_sage import HeteroGraphSAGE
+from ogb.nodeproppred import PygNodePropPredDataset
+
+import torch_geometric.transforms as T
+from torch_geometric.datasets import OGB_MAG, Reddit
+from torch_geometric.nn.models.basic_gnn import GAT, GCN, PNA, EdgeCNN
+
+models_dict = {
+    'edge_conv': EdgeCNN,
+    'gat': GAT,
+    'gcn': GCN,
+    'pna_conv': PNA,
+    'rgat': HeteroGAT,
+    'rgcn': HeteroGraphSAGE,
+}
+
+
+def get_dataset(name, root):
+    path = osp.dirname(osp.realpath(__file__))
+
+    if name == 'ogbn-mag':
+        transform = T.ToUndirected(merge=True)
+        dataset = OGB_MAG(root=osp.join(path, root, 'mag'),
+                          preprocess='metapath2vec', transform=transform)
+    elif name == 'ogbn-products':
+        dataset = PygNodePropPredDataset('ogbn-products',
+                                         root=osp.join(path, root, 'products'))
+    elif name == 'reddit':
+        dataset = Reddit(root=osp.join(path, root, 'reddit'))
+
+    return dataset
+
+
+def get_model(name, params, metadata=None):
+    try:
+        model_type = models_dict[name]
+    except KeyError:
+        print(f'Model {name} not supported!')
+
+    if name == 'rgat':
+        model = model_type(metadata, params['hidden_channels'],
+                           params['num_layers'], params['output_channels'],
+                           params['num_heads'])
+
+    elif name == 'rgcn':
+        model = model_type(metadata, params['hidden_channels'],
+                           params['num_layers'], params['output_channels'])
+
+    elif name == 'gat':
+        kwargs = {}
+        kwargs['heads'] = params['num_heads']
+        model = model_type(params['inputs_channels'],
+                           params['hidden_channels'], params['num_layers'],
+                           params['output_channels'], **kwargs)
+
+    elif name == 'pna_conv':
+        kwargs = {}
+        kwargs['aggregators'] = ['mean', 'min', 'max', 'std']
+        kwargs['scalers'] = ['identity', 'amplification', 'attenuation']
+        kwargs['deg'] = params['degree']
+        model = model_type(params['inputs_channels'],
+                           params['hidden_channels'], params['num_layers'],
+                           params['output_channels'], **kwargs)
+
+    else:
+        model = model_type(params['inputs_channels'],
+                           params['hidden_channels'], params['num_layers'],
+                           params['output_channels'])
+    return model

test/nn/conv/test_pna_conv.py

@@ -1,6 +1,8 @@
 import torch
 from torch_sparse import SparseTensor
 
+from torch_geometric.data import Data
+from torch_geometric.loader import DataLoader, NeighborLoader
 from torch_geometric.nn import PNAConv
 from torch_geometric.testing import is_full_test
 
@@ -32,3 +34,40 @@ def test_pna_conv():
         t = '(Tensor, SparseTensor, OptTensor) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
         assert torch.allclose(jit(x, adj.t()), out, atol=1e-6)
+
+
+def test_pna_conv_get_degree_histogram():
+    edge_index = torch.tensor([[0, 0, 0, 1, 1, 2, 3], [1, 2, 3, 2, 0, 0, 0]])
+    x = torch.randn(5, 16)
+    data = Data(x=x, edge_index=edge_index)
+    loader = NeighborLoader(
+        data,
+        num_neighbors=[-1],
+        input_nodes=None,
+        batch_size=5,
+        shuffle=False,
+    )
+    deg_hist = PNAConv.get_degree_histogram(loader)
+    deg_hist_ref = torch.tensor([1, 2, 1, 1])
+    assert torch.equal(deg_hist_ref, deg_hist)
+
+    edge_index_1 = torch.tensor([[0, 0, 0, 1, 1, 2, 3], [1, 2, 3, 2, 0, 0, 0]])
+    edge_index_2 = torch.tensor([[1, 1, 2, 2, 0, 3, 3], [2, 3, 3, 1, 1, 0, 2]])
+    edge_index_3 = torch.tensor([[1, 3, 2, 0, 0, 4, 2], [2, 0, 4, 1, 1, 0, 3]])
+    edge_index_4 = torch.tensor([[0, 1, 2, 4, 0, 1, 3], [2, 3, 3, 1, 1, 0, 2]])
+
+    x = torch.randn(5, 16)
+
+    data_1 = Data(x=x, edge_index=edge_index_1)  # deg_hist = [1, 1, 3]
+    data_2 = Data(x=x, edge_index=edge_index_2)  # deg_hist = [1, 2 ,1 ,1]
+    data_3 = Data(x=x, edge_index=edge_index_3)  # deg_hist = [0, 3, 2]
+    data_4 = Data(x=x, edge_index=edge_index_4)  # deg_hist = [1, 1, 3]
+
+    loader = DataLoader(
+        [data_1, data_2, data_3, data_4],
+        batch_size=1,
+        shuffle=False,
+    )
+    deg_hist = PNAConv.get_degree_histogram(loader)
+    deg_hist_ref = torch.tensor([3, 7, 9, 1])
+    assert torch.equal(deg_hist_ref, deg_hist)

torch_geometric/nn/conv/pna_conv.py

@@ -8,6 +8,7 @@
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.typing import Adj, OptTensor
+from torch_geometric.utils import degree
 
 from ..inits import reset
 
@@ -169,3 +170,19 @@ def __repr__(self):
         return (f'{self.__class__.__name__}({self.in_channels}, '
                 f'{self.out_channels}, towers={self.towers}, '
                 f'edge_dim={self.edge_dim})')
+
+    @staticmethod
+    def get_degree_histogram(loader) -> Tensor:
+        max_degree = 0
+        for data in loader:
+            d = degree(data.edge_index[1], num_nodes=data.num_nodes,
+                       dtype=torch.long)
+            max_degree = max(max_degree, int(d.max()))
+        # Compute the in-degree histogram tensor
+        deg_histogram = torch.zeros(max_degree + 1, dtype=torch.long)
+        for data in loader:
+            d = degree(data.edge_index[1], num_nodes=data.num_nodes,
+                       dtype=torch.long)
+            deg_histogram += torch.bincount(d, minlength=deg_histogram.numel())
+
+        return deg_histogram