Merge pull request #59 from LukasHedegaard/develop

Add support for GroupNorm and InstanceNorm

Author: LukasHedegaard

CHANGELOG.md

@@ -8,6 +8,11 @@ From v1.0.0 and on, the project will adherence strictly to Semantic Versioning.
 
 ## Unpublished
 
+## [1.1.1] - 2023-01-10
+
+### Added
+- Support for `GroupNorm` and `InstanceNorm`
+
 
 ## [1.1.0] - 2022-12-19
 

README.md

@@ -456,14 +456,19 @@ We support drop-in interoperability with with the following _torch.nn_ modules:
   - `nn.BatchNorm1d`
   - `nn.BatchNorm2d`
   - `nn.BatchNorm3d`
-  - `nn.LayerNorm`
+  - `nn.GroupNorm`,
+  - `nn.InstanceNorm1d` (affine=True, track_running_stats=True required)
+  - `nn.InstanceNorm2d` (affine=True, track_running_stats=True required)
+  - `nn.InstanceNorm3d` (affine=True, track_running_stats=True required)
+  - `nn.LayerNorm` (only non-temporal dimensions must be specified)
 
 </details>
 
 <details>
 <summary><b>Dropout</b></summary>
 
   - `nn.Dropout`
+  - `nn.Dropout1d`
   - `nn.Dropout2d`
   - `nn.Dropout3d`
   - `nn.AlphaDropout`

continual/__about__.py

@@ -1,9 +1,9 @@
 import time
 
-__version__ = "1.1.0"
+__version__ = "1.1.1"
 __author__ = "Lukas Hedegaard"
 __author_email__ = "[email protected]"
 __license__ = "Apache-2.0"
 __copyright__ = f'Copyright (c) 2021-{time.strftime("%Y")}, {__author__}'
 __homepage__ = "https://github.com/lukashedegaard/continual-inference"
-__docs__ = "Building blocks for Continual Inference Networks in PyTorch"
+__docs__ = "A Python library for Continual Inference Networks in PyTorch"

continual/convert.py

@@ -2,7 +2,7 @@
 
 from functools import wraps
 from types import FunctionType
-from typing import Callable, Type
+from typing import Callable, Type, Union
 
 from torch import Tensor, nn
 
@@ -150,6 +150,8 @@ def forward_with_callmode(*args, **kwargs):
     nn.BatchNorm1d,
     nn.BatchNorm2d,
     nn.BatchNorm3d,
+    nn.LayerNorm,
+    nn.GroupNorm,
     # >> Dropout modules
     nn.Dropout,
     nn.Dropout2d,
@@ -159,6 +161,28 @@ def forward_with_callmode(*args, **kwargs):
 }
 
 
+_circumvent_message = " to work with automatic conversion. You can circumvent this by wrapping the module in `co.forward_stepping(your_module)`. Note however, that this may break correspondence between forward and forward_step."
+
+
+def _instance_norm_condition(
+    module: Union[nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d]
+):
+    assert module.affine, (
+        f"{type(module)} must be specified with `affine==True`" + _circumvent_message
+    )
+    assert module.track_running_stats, (
+        f"{type(module)} must be specified with `track_running_stats==True`"
+        + _circumvent_message
+    )
+
+
+CONDITIONAL_MAPPING = {
+    nn.InstanceNorm1d: _instance_norm_condition,
+    nn.InstanceNorm2d: _instance_norm_condition,
+    nn.InstanceNorm3d: _instance_norm_condition,
+}
+
+
 class ModuleNotRegisteredError(Exception):
     ...
 
@@ -206,6 +230,10 @@ def continual(module: nn.Module) -> CoModule:
     if type(module) in NAIVE_MAPPING:
         return forward_stepping(module)
 
+    if type(module) in CONDITIONAL_MAPPING:
+        CONDITIONAL_MAPPING[type(module)](module)
+        return forward_stepping(module)
+
     assert type(module) in MODULE_MAPPING, (
         f"A registered conversion for {module} was not found. "
         "You can register a custom conversion as follows:"

tests/continual/test_norm.py

@@ -0,0 +1,84 @@
+import torch
+from torch import nn
+
+import continual as co
+
+
+def test_nn_norms():
+    S = 3
+
+    long_example_clip = torch.normal(mean=torch.zeros(10 * 3 * 3)).reshape(
+        (1, 1, 10, 3, 3)
+    )
+
+    b_norm = nn.BatchNorm3d(1)
+    b_norm.weight = nn.Parameter(3 * torch.ones_like(b_norm.weight))
+    b_norm.bias = nn.Parameter(1 * torch.ones_like(b_norm.bias))
+
+    i_norm = nn.InstanceNorm3d(2, affine=True, track_running_stats=True)
+    i_norm.weight = nn.Parameter(4 * torch.ones_like(i_norm.weight))
+    i_norm.bias = nn.Parameter(2 * torch.ones_like(i_norm.bias))
+
+    l_norm = nn.LayerNorm([S, S])  # NB: Doesn't work over temporal axis
+    l_norm.weight = nn.Parameter(5 * torch.ones_like(l_norm.weight))
+    l_norm.bias = nn.Parameter(3 * torch.ones_like(l_norm.bias))
+
+    g_norm = nn.GroupNorm(2, 2)
+    g_norm.weight = nn.Parameter(6 * torch.ones_like(g_norm.weight))
+    g_norm.bias = nn.Parameter(4 * torch.ones_like(g_norm.bias))
+
+    seq = nn.Sequential(
+        b_norm,
+        nn.Conv3d(
+            in_channels=1,
+            out_channels=2,
+            kernel_size=(5, S, S),
+            bias=True,
+            padding=(0, 1, 1),
+            padding_mode="zeros",
+        ),
+        i_norm,
+        l_norm,
+        g_norm,
+        nn.Conv3d(
+            in_channels=2,
+            out_channels=1,
+            kernel_size=(3, S, S),
+            bias=True,
+            padding=(0, 1, 1),
+            padding_mode="zeros",
+        ),
+        nn.MaxPool3d(kernel_size=(1, 2, 2)),
+    )
+    seq.eval()
+
+    coseq = co.Sequential.build_from(seq)
+    coseq.eval()
+
+    assert coseq.delay == (5 - 1) + (3 - 1)
+
+    # forward
+    output = seq.forward(long_example_clip)
+    co_output = coseq.forward(long_example_clip)
+    assert torch.allclose(output, co_output)
+
+    # forward_steps
+    co_output_firsts_0 = coseq.forward_steps(
+        long_example_clip[:, :, :-1], update_state=False
+    )
+    co_output_firsts = coseq.forward_steps(long_example_clip[:, :, :-1])
+    assert torch.allclose(co_output_firsts, co_output_firsts_0, atol=1e-7)
+    assert torch.allclose(co_output_firsts, output[:, :, :-1], atol=1e-7)
+
+    # forward_step
+    co_output_last_0 = coseq.forward_step(
+        long_example_clip[:, :, -1], update_state=False
+    )
+    co_output_last = coseq.forward_step(long_example_clip[:, :, -1])
+    assert torch.allclose(co_output_last, co_output_last_0, atol=1e-7)
+    assert torch.allclose(co_output_last, output[:, :, -1], atol=1e-7)
+
+    # Clean state can be used to restart seq computation
+    coseq.clean_state()
+    co_output_firsts = coseq.forward_steps(long_example_clip[:, :, :-1])
+    assert torch.allclose(co_output_firsts, output[:, :, :-1], atol=1e-7)