Add option for BatchNorm to handle batches of size one

CHANGELOG.md

@@ -30,6 +30,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `BaseStorage.get()` functionality ([#5240](https://github.com/pyg-team/pytorch_geometric/pull/5240))
 - Added a test to confirm that `to_hetero` works with `SparseTensor` ([#5222](https://github.com/pyg-team/pytorch_geometric/pull/5222))
 ### Changed
+- Changed `BatchNorm` to allow for batches of size one during training ([#5530](https://github.com/pyg-team/pytorch_geometric/pull/5530))
 - Fixed a bug when applying several scalers with `PNAConv` ([#5514](https://github.com/pyg-team/pytorch_geometric/issues/5514))
 - Allow `.` in `ParameterDict` key names ([#5494](https://github.com/pyg-team/pytorch_geometric/pull/5494))
 - Renamed `drop_unconnected_nodes` to `drop_unconnected_node_types` and `drop_orig_edges` to `drop_orig_edge_types` in `AddMetapaths` ([#5490](https://github.com/pyg-team/pytorch_geometric/pull/5490))

test/nn/norm/test_batch_norm.py

@@ -17,3 +17,12 @@ def test_batch_norm(conf):
 
     out = norm(x)
     assert out.size() == (100, 16)
+
+    x = torch.randn(1, 16)
+    with pytest.raises(ValueError):
+        _ = norm(x)
+
+    norm = BatchNorm(16, affine=conf, track_running_stats=conf,
+                     allow_no_batch=True)
+    out = norm(x)
+    assert torch.allclose(out, x)

torch_geometric/nn/norm/batch_norm.py

@@ -30,18 +30,39 @@ class BatchNorm(torch.nn.Module):
             :obj:`False`, this module does not track such statistics and always
             uses batch statistics in both training and eval modes.
             (default: :obj:`True`)
+        allow_no_batch (bool, optional): If set to :obj:`True`, batches with
+            only a single element will work as though in training mode. That is
+            the running mean and variance will be used.
+            (default: :obj:`False`)
     """
-    def __init__(self, in_channels, eps=1e-5, momentum=0.1, affine=True,
-                 track_running_stats=True):
+    def __init__(self, in_channels: int, eps: float = 1e-5,
+                 momentum: float = 0.1, affine: bool = True,
+                 track_running_stats: bool = True,
+                 allow_no_batch: bool = False):
         super().__init__()
         self.module = torch.nn.BatchNorm1d(in_channels, eps, momentum, affine,
                                            track_running_stats)
+        self.allow_no_batch = allow_no_batch
+        self.in_channels = in_channels
 
     def reset_parameters(self):
         self.module.reset_parameters()
 
     def forward(self, x: Tensor) -> Tensor:
         """"""
+        if self.allow_no_batch and x.size(0) <= 1:
+            training = self.module.training
+            running_mean = self.module.running_mean
+            running_var = self.module.running_var
+            if running_mean is None or running_var is None:
+                self.module.running_var = torch.ones(self.in_channels)
+                self.module.running_mean = torch.zeros(self.in_channels)
+            self.module.eval()
+            out = self.module(x)
+            self.module.training = training
+            self.module.running_mean = running_mean
+            self.module.running_var = running_var
+            return out
         return self.module(x)
 
     def __repr__(self):