Add batch_size argument to normalization layers

CHANGELOG.md

@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Added an optional `batch_size` argument to `LayerNorm`, `GraphNorm`, `InstanceNorm`, `GraphSizeNorm` and `PairNorm` ([#7135](https://github.com/pyg-team/pytorch_geometric/pull/7135))
 - Improved code coverage ([#7093](https://github.com/pyg-team/pytorch_geometric/pull/7093))
 - Fix `numpy` incompatiblity when reading files for `Planetoid` datasets ([#7141](https://github.com/pyg-team/pytorch_geometric/pull/7141))
 - Added support for `Data.num_edges` for native `torch.sparse.Tensor` adjacency matrices ([#7104](https://github.com/pyg-team/pytorch_geometric/pull/7104))

torch_geometric/nn/norm/graph_norm.py

@@ -4,6 +4,7 @@
 from torch import Tensor
 
 from torch_geometric.nn.inits import ones, zeros
+from torch_geometric.typing import OptTensor
 from torch_geometric.utils import scatter
 
 
@@ -44,18 +45,23 @@ def reset_parameters(self):
         zeros(self.bias)
         ones(self.mean_scale)
 
-    def forward(self, x: Tensor, batch: Optional[Tensor] = None) -> Tensor:
+    def forward(self, x: Tensor, batch: OptTensor = None,
+                batch_size: Optional[int] = None) -> Tensor:
         r"""
         Args:
             x (torch.Tensor): The source tensor.
             batch (torch.Tensor, optional): The batch vector
                 :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
                 each element to a specific example. (default: :obj:`None`)
+            batch_size (int, optional): The number of examples :math:`B`.
+                Automatically calculated if not given. (default: :obj:`None`)
         """
         if batch is None:
             batch = x.new_zeros(x.size(0), dtype=torch.long)
+            batch_size = 1
 
-        batch_size = int(batch.max()) + 1
+        if batch_size is None:
+            batch_size = int(batch.max()) + 1
 
         mean = scatter(x, batch, 0, batch_size, reduce='mean')
         out = x - mean.index_select(0, batch) * self.mean_scale

torch_geometric/nn/norm/graph_size_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 import torch.nn as nn
 from torch import Tensor
@@ -18,18 +20,22 @@ class GraphSizeNorm(nn.Module):
     def __init__(self):
         super().__init__()
 
-    def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor:
+    def forward(self, x: Tensor, batch: OptTensor = None,
+                batch_size: Optional[int] = None) -> Tensor:
         r"""
         Args:
             x (torch.Tensor): The source tensor.
             batch (torch.Tensor, optional): The batch vector
                 :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
                 each element to a specific example. (default: :obj:`None`)
+            batch_size (int, optional): The number of examples :math:`B`.
+                Automatically calculated if not given. (default: :obj:`None`)
         """
         if batch is None:
             batch = torch.zeros(x.size(0), dtype=torch.long, device=x.device)
+            batch_size = 1
 
-        inv_sqrt_deg = degree(batch, dtype=x.dtype).pow(-0.5)
+        inv_sqrt_deg = degree(batch, batch_size, dtype=x.dtype).pow(-0.5)
         return x * inv_sqrt_deg.index_select(0, batch).view(-1, 1)
 
     def __repr__(self) -> str:

torch_geometric/nn/norm/instance_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn.modules.instancenorm import _InstanceNorm
@@ -50,13 +52,16 @@ def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
         super().reset_parameters()
 
-    def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor:
+    def forward(self, x: Tensor, batch: OptTensor = None,
+                batch_size: Optional[int] = None) -> Tensor:
         r"""
         Args:
             x (torch.Tensor): The source tensor.
             batch (torch.Tensor, optional): The batch vector
                 :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
                 each element to a specific example. (default: :obj:`None`)
+            batch_size (int, optional): The number of examples :math:`B`.
+                Automatically calculated if not given. (default: :obj:`None`)
         """
         if batch is None:
             out = F.instance_norm(
@@ -65,7 +70,8 @@ def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor:
                 or not self.track_running_stats, self.momentum, self.eps)
             return out.squeeze(0).t()
 
-        batch_size = int(batch.max()) + 1
+        if batch_size is None:
+            batch_size = int(batch.max()) + 1
 
         mean = var = unbiased_var = x  # Dummies.
 

torch_geometric/nn/norm/layer_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 import torch.nn.functional as F
 from torch import Tensor
@@ -62,21 +64,25 @@ def reset_parameters(self):
         ones(self.weight)
         zeros(self.bias)
 
-    def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor:
+    def forward(self, x: Tensor, batch: OptTensor = None,
+                batch_size: Optional[int] = None) -> Tensor:
         r"""
         Args:
             x (torch.Tensor): The source tensor.
             batch (torch.Tensor, optional): The batch vector
                 :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
                 each element to a specific example. (default: :obj:`None`)
+            batch_size (int, optional): The number of examples :math:`B`.
+                Automatically calculated if not given. (default: :obj:`None`)
         """
         if self.mode == 'graph':
             if batch is None:
                 x = x - x.mean()
                 out = x / (x.std(unbiased=False) + self.eps)
 
             else:
-                batch_size = int(batch.max()) + 1
+                if batch_size is None:
+                    batch_size = int(batch.max()) + 1
 
                 norm = degree(batch, batch_size, dtype=x.dtype).clamp_(min=1)
                 norm = norm.mul_(x.size(-1)).view(-1, 1)

torch_geometric/nn/norm/pair_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 from torch import Tensor
 
@@ -36,13 +38,16 @@ def __init__(self, scale: float = 1., scale_individually: bool = False,
         self.scale_individually = scale_individually
         self.eps = eps
 
-    def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor:
+    def forward(self, x: Tensor, batch: OptTensor = None,
+                batch_size: Optional[int] = None) -> Tensor:
         r"""
         Args:
             x (torch.Tensor): The source tensor.
             batch (torch.Tensor, optional): The batch vector
                 :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
                 each element to a specific example. (default: :obj:`None`)
+            batch_size (int, optional): The number of examples :math:`B`.
+                Automatically calculated if not given. (default: :obj:`None`)
         """
         scale = self.scale
 
@@ -55,13 +60,13 @@ def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor:
                 return scale * x / (self.eps + x.norm(2, -1, keepdim=True))
 
         else:
-            mean = scatter(x, batch, dim=0, reduce='mean')
+            mean = scatter(x, batch, dim=0, dim_size=batch_size, reduce='mean')
             x = x - mean.index_select(0, batch)
 
             if not self.scale_individually:
                 return scale * x / torch.sqrt(self.eps + scatter(
                     x.pow(2).sum(-1, keepdim=True), batch, dim=0,
-                    reduce='mean').index_select(0, batch))
+                    dim_size=batch_size, reduce='mean').index_select(0, batch))
             else:
                 return scale * x / (self.eps + x.norm(2, -1, keepdim=True))
 

torch_geometric/utils/nested.py

@@ -28,7 +28,7 @@ def to_nested_tensor(
             (default: :obj:`None`)
         ptr (torch.Tensor, optional): Alternative representation of
             :obj:`batch` in compressed format. (default: :obj:`None`)
-        batch_size (int, optional) The batch size :math:`B`.
+        batch_size (int, optional): The batch size :math:`B`.
             (default: :obj:`None`)
     """
     if ptr is not None: