Merge pull request #34 from LukasHedegaard/develop

Receptive field, Phantom padding, Reshape, Lambda update and various fixes

Author: LukasHedegaard

CHANGELOG.md

@@ -10,6 +10,26 @@ From v1.0.0 and on, the project will adherence strictly to Semantic Versioning.
 ## [Unreleased]
 
 
+## [0.14.0]
+### Added
+- Added `phantom_padding` to `Residual`.
+- Added `receptive_field` property.
+- Added `Reshape` module.
+
+### Changed
+- Rename `forward_shrink` argument to `auto_shrink` in `Delay`.
+- Torch requirement to v1.9.
+- Replace `Lambda` unsqueeze_step with takes_time and new default to False.
+
+## Fixed
+- `padding` property in sequence.
+- `delay` property in sequence.
+- `strict` mode in `load_state_dict`.
+
+## Removed
+- Assertion error in `BroadcastReduce` for modules with different delays.
+
+
 ## [0.13.0]
 ### Added
 - Add `forward_shrink` option to `Delay` and `Residual`.

README.md

@@ -52,6 +52,10 @@ last = conv.forward_step(example[:, :, 4])
 
 assert torch.allclose(output[:, :, : conv.delay], firsts)
 assert torch.allclose(output[:, :, conv.delay], last)
+
+# Temporal properties
+assert conv.receptive_field == 3
+assert conv.delay == 2
 ```
 
 See the "Advanced Examples" section for additional examples..
@@ -211,6 +215,7 @@ Below is a list of the modules and utilities included in the library:
 
 - Other
     - `co.Delay` - Pure delay module (e.g. needed in residuals).
+    - `co.Reshape` - Reshape non-temporal dimensions.
     - `co.Lambda` - Lambda module which wraps any function.
     - `co.Add` - Adds a constant value.
     - `co.Multiply` - Multiplies with a constant factor.
@@ -361,7 +366,7 @@ inception_module = co.BroadcastReduce(
     ),
     co.Sequential(
         norm_relu(co.Conv3d(192, 16, kernel_size=1), 16),
-        norm_relu(co.Conv3d(16, 32, kernel_size=3, padding=1), 32),
+        norm_relu(co.Conv3d(16, 32, kernel_size=5, padding=2), 32),
     ),
     co.Sequential(
         co.MaxPool3d(kernel_size=(1, 3, 3), padding=(0, 1, 1), stride=1),

continual/__init__.py

@@ -26,4 +26,5 @@
     MaxPool3d,
 )
 from .ptflops import _register_ptflops  # noqa: F401
+from .shape import Reshape  # noqa: F401
 from .utils import flat_state_dict, load_state_dict, state_dict  # noqa: F401

continual/closure.py

@@ -9,16 +9,20 @@
 
 
 class Lambda(CoModule, nn.Module):
-    """Module wrapper for stateless functions
+    """Module wrapper for stateless functions.
 
     NB: Operations performed in a Lambda are not counted in `ptflops`
+
+    Args:
+        fn (Callable[[Tensor], Tensor]): Function to be called during forward.
+        takes_time (bool, optional): If True, `fn` recieves all steps, if False, it received one step and no time dimension. Defaults to False.
     """
 
-    def __init__(self, fn: Callable[[Tensor], Tensor], unsqueeze_step=True):
+    def __init__(self, fn: Callable[[Tensor], Tensor], takes_time=False):
         nn.Module.__init__(self)
         assert callable(fn), "The pased function should be callable."
         self.fn = fn
-        self.unsqueeze_step = unsqueeze_step
+        self.takes_time = takes_time
 
     def __repr__(self) -> str:
         s = self.fn.__name__
@@ -47,26 +51,27 @@ def __repr__(self) -> str:
         return f"Lambda({s})"
 
     def forward(self, input: Tensor) -> Tensor:
-        return self.fn(input)
+        if self.takes_time:
+            return self.fn(input)
+
+        return torch.stack(
+            [self.fn(input[:, :, t]) for t in range(input.shape[2])], dim=2
+        )
+
+    def forward_steps(self, input: Tensor, pad_end=False, update_state=True) -> Tensor:
+        return self.forward(input)
 
     def forward_step(self, input: Tensor, update_state=True) -> Tensor:
-        if self.unsqueeze_step:
+        if self.takes_time:
             input = input.unsqueeze(dim=2)
         output = self.fn(input)
-        if self.unsqueeze_step:
+        if self.takes_time:
             output = output.squeeze(dim=2)
         return output
 
-    def forward_steps(self, input: Tensor, pad_end=False, update_state=True) -> Tensor:
-        return self.fn(input)
-
-    @property
-    def delay(self) -> int:
-        return 0
-
     @staticmethod
-    def build_from(fn: Callable[[Tensor], Tensor]) -> "Lambda":
-        return Lambda(fn)
+    def build_from(fn: Callable[[Tensor], Tensor], takes_time=False) -> "Lambda":
+        return Lambda(fn, takes_time)
 
 
 def _multiply(x: Tensor, factor: Union[float, int, Tensor]):
@@ -76,7 +81,7 @@ def _multiply(x: Tensor, factor: Union[float, int, Tensor]):
 def Multiply(factor) -> Lambda:
     """Create Lambda with multiplication function"""
     fn = partial(_multiply, factor=factor)
-    return Lambda(fn)
+    return Lambda(fn, takes_time=True)
 
 
 def _add(x: Tensor, constant: Union[float, int, Tensor]):
@@ -86,7 +91,7 @@ def _add(x: Tensor, constant: Union[float, int, Tensor]):
 def Add(constant) -> Lambda:
     """Create Lambda with addition function"""
     fn = partial(_add, constant=constant)
-    return Lambda(fn)
+    return Lambda(fn, takes_time=True)
 
 
 def _unity(x: Tensor):
@@ -95,18 +100,18 @@ def _unity(x: Tensor):
 
 def Unity() -> Lambda:
     """Create Lambda with addition function"""
-    return Lambda(_unity)
+    return Lambda(_unity, takes_time=True)
 
 
 def Constant(constant: float):
-    return Lambda(lambda x: constant * torch.ones_like(x))
+    return Lambda(lambda x: constant * torch.ones_like(x), takes_time=True)
 
 
 def Zero() -> Lambda:
     """Create Lambda with zero output"""
-    return Lambda(torch.zeros_like)
+    return Lambda(torch.zeros_like, takes_time=True)
 
 
 def One() -> Lambda:
     """Create Lambda with zero output"""
-    return Lambda(torch.ones_like)
+    return Lambda(torch.ones_like, takes_time=True)