Merge branch 'master' into gcn

Author: rusty1s

CHANGELOG.md

@@ -6,12 +6,13 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ## [2.3.0] - 2023-MM-DD
 ### Added
 - Started adding `torch.sparse` support to PyG ([#5906](https://github.com/pyg-team/pytorch_geometric/pull/5906), [#5944](https://github.com/pyg-team/pytorch_geometric/pull/5944), [#6003](https://github.com/pyg-team/pytorch_geometric/pull/6003), [#6033](https://github.com/pyg-team/pytorch_geometric/pull/6033))
+- Add inputs_channels back in training benchmark ([#6154](https://github.com/pyg-team/pytorch_geometric/pull/6154))
 - Added support for dropping nodes in `utils.to_dense_batch` in case `max_num_nodes` is smaller than the number of nodes  ([#6124](https://github.com/pyg-team/pytorch_geometric/pull/6124))
 - Added the RandLA-Net architecture as an example ([#5117](https://github.com/pyg-team/pytorch_geometric/pull/5117))
 ### Changed
 - [Breaking Change] Moved PyTorch Lightning data modules to `torch_geometric.data.lightning` ([#6140](https://github.com/pyg-team/pytorch_geometric/pull/6140))
 - Make `torch_sparse` an optional dependency ([#6132](https://github.com/pyg-team/pytorch_geometric/pull/6132), [#6134](https://github.com/pyg-team/pytorch_geometric/pull/6134), [#6138](https://github.com/pyg-team/pytorch_geometric/pull/6138), [#6139](https://github.com/pyg-team/pytorch_geometric/pull/6139))
-- Optimized `utils.softmax` implementation ([#6113](https://github.com/pyg-team/pytorch_geometric/pull/6113))
+- Optimized `utils.softmax` implementation ([#6113](https://github.com/pyg-team/pytorch_geometric/pull/6113), [#6155](https://github.com/pyg-team/pytorch_geometric/pull/6155))
 - Optimized `topk` implementation for large enough graphs ([#6123](https://github.com/pyg-team/pytorch_geometric/pull/6123))
 ### Removed
 

benchmark/training/training_benchmark.py

@@ -80,6 +80,10 @@ def run(args: argparse.ArgumentParser) -> None:
         else:
             amp = torch.cpu.amp.autocast(enabled=args.bf16)
 
+        inputs_channels = data[
+            'paper'].num_features if dataset_name == 'ogbn-mag' \
+            else data.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} '
@@ -124,6 +128,7 @@ def run(args: argparse.ArgumentParser) -> None:
                               f'Sparse tensor={args.use_sparse_tensor}')
 
                         params = {
+                            'inputs_channels': inputs_channels,
                             'hidden_channels': hidden_channels,
                             'output_channels': num_classes,
                             'num_heads': args.num_heads,

test/loader/test_neighbor_loader.py

@@ -9,6 +9,7 @@
 import torch
 from torch_sparse import SparseTensor
 
+import torch_geometric.typing
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import GraphConv, to_hetero
@@ -20,12 +21,6 @@
 )
 from torch_geometric.utils import k_hop_subgraph
 
-try:
-    import pyg_lib  # noqa
-    _WITH_PYG_LIB = True
-except ImportError:
-    _WITH_PYG_LIB = False
-
 
 def get_edge_index(num_src_nodes, num_dst_nodes, num_edges, dtype=torch.int64):
     row = torch.randint(num_src_nodes, (num_edges, ), dtype=dtype)
@@ -44,7 +39,7 @@ def is_subset(subedge_index, edge_index, src_idx, dst_idx):
 @pytest.mark.parametrize('directed', [True])  # TODO re-enable undirected mode
 @pytest.mark.parametrize('dtype', [torch.int64, torch.int32])
 def test_homogeneous_neighbor_loader(directed, dtype):
-    if dtype != torch.int64 and not _WITH_PYG_LIB:
+    if dtype != torch.int64 and not torch_geometric.typing.WITH_PYG_LIB:
         return
 
     torch.manual_seed(12345)
@@ -83,7 +78,7 @@ def test_homogeneous_neighbor_loader(directed, dtype):
 @pytest.mark.parametrize('directed', [True])  # TODO re-enable undirected mode
 @pytest.mark.parametrize('dtype', [torch.int64, torch.int32])
 def test_heterogeneous_neighbor_loader(directed, dtype):
-    if dtype != torch.int64 and not _WITH_PYG_LIB:
+    if dtype != torch.int64 and not torch_geometric.typing.WITH_PYG_LIB:
         return
 
     torch.manual_seed(12345)

test/nn/norm/test_layer_norm.py

@@ -12,7 +12,7 @@ def test_layer_norm(affine, mode):
     batch = torch.zeros(100, dtype=torch.long)
 
     norm = LayerNorm(16, affine=affine, mode=mode)
-    assert norm.__repr__() == f'LayerNorm(16, mode={mode})'
+    assert norm.__repr__() == f'LayerNorm(16, affine={affine}, mode={mode})'
 
     if is_full_test():
         torch.jit.script(norm)

torch_geometric/nn/conv/rgcn_conv.py

@@ -8,20 +8,12 @@
 from torch_scatter import scatter
 from torch_sparse import SparseTensor, masked_select_nnz, matmul
 
+import torch_geometric.typing
 from torch_geometric.nn.conv import MessagePassing
-from torch_geometric.typing import Adj, OptTensor
+from torch_geometric.typing import Adj, OptTensor, pyg_lib
 
 from ..inits import glorot, zeros
 
-try:
-    from pyg_lib.ops import segment_matmul  # noqa
-    _WITH_PYG_LIB = True
-except ImportError:
-    _WITH_PYG_LIB = False
-
-    def segment_matmul(inputs: Tensor, ptr: Tensor, other: Tensor) -> Tensor:
-        raise NotImplementedError
-
 
 @torch.jit._overload
 def masked_edge_index(edge_index, edge_mask):
@@ -112,7 +104,6 @@ def __init__(
     ):
         kwargs.setdefault('aggr', aggr)
         super().__init__(node_dim=0, **kwargs)
-        self._WITH_PYG_LIB = _WITH_PYG_LIB
 
         if num_bases is not None and num_blocks is not None:
             raise ValueError('Can not apply both basis-decomposition and '
@@ -225,7 +216,7 @@ def forward(self, x: Union[OptTensor, Tuple[OptTensor, Tensor]],
                 out = out + h.contiguous().view(-1, self.out_channels)
 
         else:  # No regularization/Basis-decomposition ========================
-            if (self._WITH_PYG_LIB and self.num_bases is None
+            if (torch_geometric.typing.WITH_PYG_LIB and self.num_bases is None
                     and x_l.is_floating_point()
                     and isinstance(edge_index, Tensor)):
                 if not self.is_sorted:
@@ -264,7 +255,7 @@ def forward(self, x: Union[OptTensor, Tuple[OptTensor, Tensor]],
     def message(self, x_j: Tensor, edge_type_ptr: OptTensor) -> Tensor:
         if edge_type_ptr is not None:
             # TODO Re-weight according to edge type degree for `aggr=mean`.
-            return segment_matmul(x_j, edge_type_ptr, self.weight)
+            return pyg_lib.ops.segment_matmul(x_j, edge_type_ptr, self.weight)
 
         return x_j
 

torch_geometric/nn/dense/linear.py

@@ -7,16 +7,9 @@
 from torch import Tensor, nn
 from torch.nn.parameter import Parameter
 
+import torch_geometric.typing
 from torch_geometric.nn import inits
-
-try:
-    from pyg_lib.ops import segment_matmul  # noqa
-    _WITH_PYG_LIB = True
-except ImportError:
-    _WITH_PYG_LIB = False
-
-    def segment_matmul(inputs: Tensor, ptr: Tensor, other: Tensor) -> Tensor:
-        raise NotImplementedError
+from torch_geometric.typing import pyg_lib
 
 
 def is_uninitialized_parameter(x: Any) -> bool:
@@ -220,9 +213,7 @@ def __init__(self, in_channels: int, out_channels: int, num_types: int,
         self.is_sorted = is_sorted
         self.kwargs = kwargs
 
-        self._WITH_PYG_LIB = _WITH_PYG_LIB
-
-        if self._WITH_PYG_LIB:
+        if torch_geometric.typing.WITH_PYG_LIB:
             self.lins = None
             self.weight = torch.nn.Parameter(
                 torch.Tensor(num_types, in_channels, out_channels))
@@ -241,7 +232,7 @@ def __init__(self, in_channels: int, out_channels: int, num_types: int,
         self.reset_parameters()
 
     def reset_parameters(self):
-        if self._WITH_PYG_LIB:
+        if torch_geometric.typing.WITH_PYG_LIB:
             reset_weight_(self.weight, self.in_channels,
                           self.kwargs.get('weight_initializer', None))
             reset_weight_(self.bias, self.in_channels,
@@ -256,7 +247,7 @@ def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:
             x (Tensor): The input features.
             type_vec (LongTensor): A vector that maps each entry to a type.
         """
-        if self._WITH_PYG_LIB:
+        if torch_geometric.typing.WITH_PYG_LIB:
             assert self.weight is not None
 
             if not self.is_sorted:
@@ -266,7 +257,7 @@ def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:
 
             type_vec_ptr = torch.ops.torch_sparse.ind2ptr(
                 type_vec, self.num_types)
-            out = segment_matmul(x, type_vec_ptr, self.weight)
+            out = pyg_lib.ops.segment_matmul(x, type_vec_ptr, self.weight)
             if self.bias is not None:
                 out += self.bias[type_vec]
         else:

torch_geometric/nn/norm/layer_norm.py

@@ -42,6 +42,7 @@ def __init__(self, in_channels: int, eps: float = 1e-5,
 
         self.in_channels = in_channels
         self.eps = eps
+        self.affine = affine
         self.mode = mode
 
         if affine:
@@ -94,4 +95,4 @@ def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor:
 
     def __repr__(self):
         return (f'{self.__class__.__name__}({self.in_channels}, '
-                f'mode={self.mode})')
+                f'affine={self.affine}, mode={self.mode})')

torch_geometric/sampler/neighbor_sampler.py

@@ -4,6 +4,7 @@
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
 from torch_geometric.data import (
     Data,
     FeatureStore,
@@ -25,12 +26,6 @@
 from torch_geometric.sampler.utils import remap_keys, to_csc, to_hetero_csc
 from torch_geometric.typing import EdgeType, NodeType, NumNeighbors, OptTensor
 
-try:
-    import pyg_lib  # noqa
-    _WITH_PYG_LIB = True
-except ImportError:
-    _WITH_PYG_LIB = False
-
 
 class NeighborSampler(BaseSampler):
     r"""An implementation of an in-memory (heterogeneous) neighbor sampler used
@@ -227,7 +222,7 @@ def _sample(
         loaders."""
         # TODO(manan): remote backends only support heterogeneous graphs:
         if self.data_cls == 'custom' or issubclass(self.data_cls, HeteroData):
-            if _WITH_PYG_LIB:
+            if torch_geometric.typing.WITH_PYG_LIB:
                 # TODO (matthias) `return_edge_id` if edge features present
                 # TODO (matthias) Ideally, `seed` should inherit the type of
                 # `colptr_dict` and `row_dict`.
@@ -285,7 +280,7 @@ def _sample(
             )
 
         if issubclass(self.data_cls, Data):
-            if _WITH_PYG_LIB:
+            if torch_geometric.typing.WITH_PYG_LIB:
                 # TODO (matthias) `return_edge_id` if edge features present
                 # TODO (matthias) Ideally, `seed` should inherit the type of
                 # `colptr` and `row`.

torch_geometric/typing.py

@@ -1,8 +1,15 @@
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, Final, List, Optional, Tuple, Union
 
 import numpy as np
 from torch import Tensor
 
+try:
+    import pyg_lib  # noqa
+    WITH_PYG_LIB: Final[bool] = True
+except ImportError:
+    pyg_lib = object
+    WITH_PYG_LIB: Final[bool] = False
+
 try:
     from torch_sparse import SparseTensor
 except ImportError:

torch_geometric/utils/softmax.py

@@ -4,6 +4,8 @@
 from torch import Tensor
 from torch_scatter import gather_csr, segment_csr
 
+import torch_geometric.typing
+from torch_geometric.typing import pyg_lib
 from torch_geometric.utils import scatter
 
 from .num_nodes import maybe_num_nodes
@@ -67,8 +69,12 @@ def softmax(
         N = maybe_num_nodes(index, num_nodes)
         with torch.no_grad():
             src_max = scatter(src, index, dim, dim_size=N, reduce='max')
-            src_max = src_max.index_select(dim, index)
-        out = (src - src_max).exp()
+        if (torch_geometric.typing.WITH_PYG_LIB and src.dim() == 2
+                and (dim == 0 or dim == -2)):
+            out = pyg_lib.ops.sampled_sub(src, src_max, right_index=index)
+        else:
+            out = src - src_max.index_select(dim, index)
+        out = out.exp()
         out_sum = scatter(out, index, dim, dim_size=N, reduce='sum') + 1e-16
         out_sum = out_sum.index_select(dim, index)
     else: