Add normalization_resolver to resolve normalization layers

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.0.5] - 2022-MM-DD
 ### Added
+- Added support for `normalization_resolver` ([#4926](https://github.com/pyg-team/pytorch_geometric/pull/4926))
 - Added support for `follow_batch` for lists or dictionaries of tensors ([#4837](https://github.com/pyg-team/pytorch_geometric/pull/4837))
 - Added `Data.validate()` and `HeteroData.validate()` functionality ([#4885](https://github.com/pyg-team/pytorch_geometric/pull/4885))
 - Added `LinkNeighborLoader` support to `LightningDataModule` ([#4868](https://github.com/pyg-team/pytorch_geometric/pull/4868))

test/nn/test_resolver.py

@@ -5,6 +5,7 @@
 from torch_geometric.nn.resolver import (
     activation_resolver,
     aggregation_resolver,
+    normalization_resolver,
 )
 
 
@@ -34,3 +35,28 @@ def test_aggregation_resolver(aggr_tuple):
     aggr_module, aggr_repr = aggr_tuple
     assert isinstance(aggregation_resolver(aggr_module()), aggr_module)
     assert isinstance(aggregation_resolver(aggr_repr), aggr_module)
+
+
+@pytest.mark.parametrize('norm_tuple', [
+    (torch_geometric.nn.norm.BatchNorm, 'batchnorm'),
+    (torch_geometric.nn.norm.InstanceNorm, 'instancenorm'),
+    (torch_geometric.nn.norm.LayerNorm, 'layernorm'),
+    (torch_geometric.nn.norm.GraphNorm, 'graphnorm'),
+    (torch_geometric.nn.norm.GraphSizeNorm, 'graphsizenorm'),
+    (torch_geometric.nn.norm.PairNorm, 'pairnorm'),
+    (torch_geometric.nn.norm.MessageNorm, 'messagenorm'),
+    (torch_geometric.nn.norm.DiffGroupNorm, 'diffgroupnorm'),
+])
+def test_normalization_resolver(norm_tuple):
+    norm_module, norm_repr = norm_tuple
+    if norm_repr in {'graphsizenorm', 'pairnorm', 'messagenorm'}:
+        assert isinstance(normalization_resolver(norm_module()), norm_module)
+        assert isinstance(normalization_resolver(norm_repr), norm_module)
+    elif norm_repr == 'diffgroupnorm':
+        assert isinstance(normalization_resolver(norm_module(16, 4)),
+                          norm_module)
+        assert isinstance(normalization_resolver(norm_repr, 16, 4),
+                          norm_module)
+    else:
+        assert isinstance(normalization_resolver(norm_module(16)), norm_module)
+        assert isinstance(normalization_resolver(norm_repr, 16), norm_module)

torch_geometric/nn/resolver.py

@@ -61,6 +61,22 @@ def activation_resolver(query: Union[Any, str] = 'relu', *args, **kwargs):
     return resolver(acts, act_dict, query, base_cls, *args, **kwargs)
 
 
+# Normalization Resolver ######################################################
+
+
+def normalization_resolver(query: Union[Any, str], *args, **kwargs):
+    import torch
+
+    import torch_geometric.nn.norm as norm
+    base_cls = torch.nn.Module
+    norms = [
+        norm for norm in vars(norm).values()
+        if isinstance(norm, type) and issubclass(norm, base_cls)
+    ]
+    norm_dict = {}
+    return resolver(norms, norm_dict, query, base_cls, *args, **kwargs)
+
+
 # Aggregation Resolver ########################################################