sinabs 3.0.4.dev25__py3-none-any.whl → 3.1.0__py3-none-any.whl

sinabs/backend/dynapcnn/dvs_layer.py

@@ -3,49 +3,26 @@ from typing import Optional, Tuple
 import torch.nn as nn
 
 from sinabs.layers import SumPool2d
+from sinabs.utils import expand_to_pair
 
 from .crop2d import Crop2d
 from .flipdims import FlipDims
 
 
-def expand_to_pair(value) -> (int, int):
-    """Expand a given value to a pair (tuple) if an int is passed.
-
-    Parameters
-    ----------
-    value:
-        int
-
-    Returns
-    -------
-    pair:
-        (int, int)
-    """
-    return (value, value) if isinstance(value, int) else value
-
-
 class DVSLayer(nn.Module):
     """DVSLayer representing the DVS pixel array on chip and/or the pre-processing. The order of
     processing is as follows MergePolarity -> Pool -> Cut -> Flip.
 
-    Parameters
-    ----------
-    input_shape;
-        Shape of input (height, width)
-    pool:
-        Sum pooling kernel size (height, width)
-    crop:
-        Crop the input to the given ROI ((top, bottom), (left, right))
-    merge_polarities:
-        If true, events from both polarities will be merged.
-    flip_x:
-        Flip the X axis
-    flip_y:
-        Flip the Y axis
-    swap_xy:
-        Swap X and Y dimensions
-    disable_pixel_array:
-        Disable the pixel array. This is useful if you want to use the DVS layer for input preprocessing.
+    Args:
+        input_shape: Shape of input (height, width).
+        pool: Sum pooling kernel size (height, width).
+        crop: Crop the input to the given ROI ((top, bottom), (left, right)).
+        merge_polarities: If true, events from both polarities will be merged.
+        flip_x: Flip the X axis.
+        flip_y: Flip the Y axis.
+        swap_xy: Swap X and Y dimensions.
+        disable_pixel_array: Disable the pixel array. This is useful if you want to use the
+            DVS layer for input preprocessing.
     """
 
     def __init__(
@@ -97,22 +74,15 @@ class DVSLayer(nn.Module):
     ) -> "DVSLayer":
         """Alternative factory method. Generate a DVSLayer from a set of torch layers.
 
-        Parameters
-        ----------
-        input_shape:
-            (channels, height, width)
-        pool_layer:
-            SumPool2d layer
-        crop_layer:
-            Crop2d layer
-        flip_layer:
-            FlipDims layer
-        disable_pixel_array:
-            Whether pixel array of new DVSLayer should be disabled.
-
-        Returns
-        -------
-        DVSLayer
+        Args:
+            input_shape: (channels, height, width).
+            pool_layer: SumPool2d layer.
+            crop_layer: Crop2d layer.
+            flip_layer: FlipDims layer.
+            disable_pixel_array: Whether pixel array of new DVSLayer should be disabled.
+
+        Returns:
+            DVSLayer
         """
         pool = (1, 1)
         crop = None
@@ -156,9 +126,8 @@ class DVSLayer(nn.Module):
     def input_shape_dict(self) -> dict:
         """The configuration dictionary for the input shape.
 
-        Returns
-        -------
-        dict
+        Returns:
+            dict
         """
         channel_count, input_size_y, input_size_x = self.input_shape
 
@@ -173,9 +142,8 @@ class DVSLayer(nn.Module):
     def get_output_shape_after_pooling(self) -> Tuple[int, int, int]:
         """Get the shape of data just after the pooling layer.
 
-        Returns
-        -------
-        (channel, height, width)
+        Returns:
+            (channel, height, width)
         """
         channel_count, input_size_y, input_size_x = self.input_shape
 
@@ -191,9 +159,8 @@ class DVSLayer(nn.Module):
     def get_output_shape_dict(self) -> dict:
         """Configuration dictionary for output shape.
 
-        Returns
-        -------
-        dict
+        Returns:
+            dict
         """
         (
             channel_count,
@@ -202,14 +169,13 @@ class DVSLayer(nn.Module):
         ) = self.get_output_shape_after_pooling()
 
         # Compute dims after cropping
-        if self.crop_layer is not None:
-            (
-                channel_count,
-                output_size_y,
-                output_size_x,
-            ) = self.crop_layer.get_output_shape(
-                (channel_count, output_size_y, output_size_x)
-            )
+        (
+            channel_count,
+            output_size_y,
+            output_size_x,
+        ) = self.crop_layer.get_output_shape(
+            (channel_count, output_size_y, output_size_x)
+        )
 
         # Compute dims after pooling
         return {
@@ -237,11 +203,13 @@ class DVSLayer(nn.Module):
         # Merge polarities
         if self.merge_polarities:
             data = data.sum(1, keepdim=True)
+
         # Pool
         out = self.pool_layer(data)
+
         # Crop
-        if self.crop_layer is not None:
-            out = self.crop_layer(out)
+        out = self.crop_layer(out)
+
         # Flip stuff
         out = self.flip_layer(out)
 
@@ -250,9 +218,8 @@ class DVSLayer(nn.Module):
     def get_pooling(self) -> Tuple[int, int]:
         """Pooling kernel shape.
 
-        Returns
-        -------
-        (ky, kx)
+        Returns:
+            (ky, kx)
         """
         return expand_to_pair(self.pool_layer.kernel_size)
 
@@ -260,26 +227,20 @@ class DVSLayer(nn.Module):
         """The coordinates for ROI. Note that this is not the same as crop parameter passed during
         the object construction.
 
-        Returns
-        -------
-        ((top, bottom), (left, right))
+        Returns:
+            ((top, bottom), (left, right))
         """
-        if self.crop_layer is not None:
-            _, h, w = self.get_output_shape_after_pooling()
-            return (
-                (self.crop_layer.top_crop, self.crop_layer.bottom_crop),
-                (self.crop_layer.left_crop, self.crop_layer.right_crop),
-            )
-        else:
-            _, output_size_y, output_size_x = self.get_output_shape()
-            return (0, output_size_y), (0, output_size_x)
+        _, h, w = self.get_output_shape_after_pooling()
+        return (
+            (self.crop_layer.top_crop, self.crop_layer.bottom_crop),
+            (self.crop_layer.left_crop, self.crop_layer.right_crop),
+        )
 
     def get_output_shape(self) -> Tuple[int, int, int]:
         """Output shape of the layer.
 
-        Returns
-        -------
-        (channel, height, width)
+        Returns:
+            (channel, height, width)
         """
         channel_count, input_size_y, input_size_x = self.input_shape
 
@@ -292,23 +253,21 @@ class DVSLayer(nn.Module):
         output_size_y = input_size_y // pooling[0]
 
         # Compute dims after cropping
-        if self.crop_layer is not None:
-            (
-                channel_count,
-                output_size_y,
-                output_size_x,
-            ) = self.crop_layer.get_output_shape(
-                (channel_count, output_size_y, output_size_x)
-            )
+        (
+            channel_count,
+            output_size_y,
+            output_size_x,
+        ) = self.crop_layer.get_output_shape(
+            (channel_count, output_size_y, output_size_x)
+        )
 
         return channel_count, output_size_y, output_size_x
 
     def get_flip_dict(self) -> dict:
         """Configuration dictionary for x, y flip.
 
-        Returns
-        -------
-        dict
+        Returns:
+            dict
         """
 
         return {"x": self.flip_layer.flip_x, "y": self.flip_layer.flip_y}
@@ -316,8 +275,7 @@ class DVSLayer(nn.Module):
     def get_swap_xy(self) -> bool:
         """True if XY has to be swapped.
 
-        Returns
-        -------
-        bool
+        Returns:
+            bool
         """
         return self.flip_layer.swap_xy

sinabs/backend/dynapcnn/dynapcnn_layer.py

@@ -1,40 +1,76 @@
 from copy import deepcopy
-from typing import Dict, Optional, Tuple, Union
-from warnings import warn
+from functools import partial
+from typing import List, Tuple
 
 import numpy as np
 import torch
 from torch import nn
 
-import sinabs.activation
 import sinabs.layers as sl
 
 from .discretize import discretize_conv_spike_
-from .dvs_layer import expand_to_pair
+
+# Define sum pooling functional as power-average pooling with power 1
+sum_pool2d = partial(nn.functional.lp_pool2d, norm_type=1)
+
+
+def convert_linear_to_conv(
+    lin: nn.Linear, input_shape: Tuple[int, int, int]
+) -> nn.Conv2d:
+    """Convert Linear layer to Conv2d.
+
+    Args:
+        lin (nn.Linear): linear layer to be converted.
+        input_shape (tuple): the tensor shape the layer expects.
+
+    Returns:
+        convolutional layer equivalent to `lin`.
+    """
+    in_chan, in_h, in_w = input_shape
+    if lin.in_features != in_chan * in_h * in_w:
+        raise ValueError(
+            "Shape of linear layer weight does not match provided input shape"
+        )
+
+    layer = nn.Conv2d(
+        in_channels=in_chan,
+        kernel_size=(in_h, in_w),
+        out_channels=lin.out_features,
+        padding=0,
+        bias=lin.bias is not None,
+    )
+
+    if lin.bias is not None:
+        layer.bias.data = lin.bias.data.clone().detach()
+
+    layer.weight.data = (
+        lin.weight.data.clone()
+        .detach()
+        .reshape((lin.out_features, in_chan, in_h, in_w))
+    )
+
+    return layer
 
 
 class DynapcnnLayer(nn.Module):
-    """Create a DynapcnnLayer object representing a dynapcnn layer.
-
-    Requires a convolutional layer, a sinabs spiking layer and an optional
-    pooling value. The layers are used in the order conv -> spike -> pool.
-
-    Parameters
-    ----------
-        conv: torch.nn.Conv2d or torch.nn.Linear
-            Convolutional or linear layer (linear will be converted to convolutional)
-        spk: sinabs.layers.IAFSqueeze
-            Sinabs IAF layer
-        in_shape: tuple of int
-            The input shape, needed to create dynapcnn configs if the network does not
-            contain an input layer. Convention: (features, height, width)
-        pool: int or None
-            Integer representing the sum pooling kernel and stride. If `None`, no
-            pooling will be applied.
-        discretize: bool
-            Whether to discretize parameters.
-        rescale_weights: int
-            Layer weights will be divided by this value.
+    """Create a DynapcnnLayer object representing a layer on DynapCNN or Speck.
+
+    Requires a convolutional layer, a sinabs spiking layer and a list of
+    pooling values. The layers are used in the order conv -> spike -> pool.
+
+    Attributes:
+        conv: torch.nn.Conv2d or torch.nn.Linear. Convolutional or linear layer.
+            Linear will be converted to convolutional.
+        spk (sinabs.layers.IAFSqueeze): Sinabs IAF layer.
+        in_shape (tuple of int):  The input shape, needed to create dynapcnn configs
+            if the network does not contain an input layer.
+            Convention: (features, height, width).
+        pool (List of integers): Each integer entry represents an output (destination
+            on chip) and whether pooling should be applied (values > 1) or not
+            (values equal to 1). The number of entries determines the number of tensors
+            the layer's forward method returns.
+        discretize (bool): Whether to discretize parameters.
+        rescale_weights (int): Layer weights will be multiplied by this value.
     """
 
     def __init__(
@@ -42,126 +78,136 @@ class DynapcnnLayer(nn.Module):
         conv: nn.Conv2d,
         spk: sl.IAFSqueeze,
         in_shape: Tuple[int, int, int],
-        pool: Optional[sl.SumPool2d] = None,
+        pool: List[int],
         discretize: bool = True,
         rescale_weights: int = 1,
     ):
         super().__init__()
 
-        self.input_shape = in_shape
+        self.in_shape = in_shape
+        self.pool = pool
+        self._discretize = discretize
+        self._rescale_weights = rescale_weights
 
+        if not isinstance(spk, sl.IAFSqueeze):
+            raise TypeError(
+                f"Unsupported spiking layer type {type(spk)}. "
+                "Only `IAFSqueeze` layers are supported."
+            )
         spk = deepcopy(spk)
+
+        # Convert `nn.Linear` to `nn.Conv2d`.
         if isinstance(conv, nn.Linear):
-            conv = self._convert_linear_to_conv(conv)
-            if spk.is_state_initialised():
-                # Expand dims
-                spk.v_mem = spk.v_mem.data.unsqueeze(-1).unsqueeze(-1)
+            conv = convert_linear_to_conv(conv, in_shape)
+            if spk.is_state_initialised() and (ndim := spk.v_mem.ndim) < 4:
+                for __ in range(4 - ndim):
+                    # Expand spatial dimensions
+                    spk.v_mem = spk.v_mem.data.unsqueeze(-1)
         else:
             conv = deepcopy(conv)
 
-        if rescale_weights != 1:
+        if self._rescale_weights != 1:
             # this has to be done after copying but before discretizing
-            conv.weight.data = (conv.weight / rescale_weights).clone().detach()
+            conv.weight.data = (conv.weight * self._rescale_weights).clone().detach()
+
+        # check if convolution kernel is a square.
+        if conv.kernel_size[0] != conv.kernel_size[1]:
+            raise ValueError(
+                "The kernel of a `nn.Conv2d` must have the same height and width."
+            )
+        for pool_size in pool:
+            if pool_size[0] != pool_size[1]:
+                raise ValueError("Only square pooling kernels are supported")
 
-        self.discretize = discretize
-        if discretize:
-            # int conversion is done while writing the config.
+        # int conversion is done while writing the config.
+        if self._discretize:
             conv, spk = discretize_conv_spike_(conv, spk, to_int=False)
 
-        self.conv_layer = conv
-        self.spk_layer = spk
-        if pool is not None:
-            if pool.kernel_size[0] != pool.kernel_size[1]:
-                raise ValueError("Only square kernels are supported")
-            self.pool_layer = deepcopy(pool)
-        else:
-            self.pool_layer = None
+        self.conv = conv
+        self.spk = spk
 
-    def _convert_linear_to_conv(self, lin: nn.Linear) -> nn.Conv2d:
-        """Convert Linear layer to Conv2d.
+    @property
+    def conv_layer(self):
+        return self.conv
 
-        Parameters
-        ----------
-            lin: nn.Linear
-                Linear layer to be converted
+    @property
+    def spk_layer(self):
+        return self.spk
 
-        Returns
-        -------
-            nn.Conv2d
-                Convolutional layer equivalent to `lin`.
-        """
+    @property
+    def discretize(self):
+        return self._discretize
 
-        in_chan, in_h, in_w = self.input_shape
+    @property
+    def rescale_weights(self):
+        return self._rescale_weights
 
-        if lin.in_features != in_chan * in_h * in_w:
-            raise ValueError("Shapes don't match.")
+    @property
+    def conv_out_shape(self):
+        return self._get_conv_output_shape()
 
-        layer = nn.Conv2d(
-            in_channels=in_chan,
-            kernel_size=(in_h, in_w),
-            out_channels=lin.out_features,
-            padding=0,
-            bias=lin.bias is not None,
-        )
+    def forward(self, x) -> List[torch.Tensor]:
+        """Torch forward pass.
 
-        if lin.bias is not None:
-            layer.bias.data = lin.bias.data.clone().detach()
+        ...
+        """
 
-        layer.weight.data = (
-            lin.weight.data.clone()
-            .detach()
-            .reshape((lin.out_features, in_chan, in_h, in_w))
-        )
+        returns = []
 
-        return layer
+        x = self.conv_layer(x)
+        x = self.spk_layer(x)
+
+        for pool in self.pool:
+            if pool == 1:
+                # no pooling is applied.
+                returns.append(x)
+            else:
+                # sum pooling of `(pool, pool)` is applied.
+                pool_out = sum_pool2d(x, kernel_size=pool)
+                returns.append(pool_out)
+
+        if len(returns) == 1:
+            return returns[0]
+        else:
+            return tuple(returns)
+
+    def zero_grad(self, set_to_none: bool = False) -> None:
+        """Call `zero_grad` method of spiking layer"""
+        return self.spk.zero_grad(set_to_none)
 
     def get_neuron_shape(self) -> Tuple[int, int, int]:
         """Return the output shape of the neuron layer.
 
-        Returns
-        -------
-        features, height, width
+        Returns:
+            conv_out_shape (tuple): formatted as (features, height, width).
         """
+        # same as the convolution's output.
+        return self._get_conv_output_shape()
 
-        def get_shape_after_conv(layer: nn.Conv2d, input_shape):
-            (ch_in, h_in, w_in) = input_shape
-            (kh, kw) = expand_to_pair(layer.kernel_size)
-            (pad_h, pad_w) = expand_to_pair(layer.padding)
-            (stride_h, stride_w) = expand_to_pair(layer.stride)
-
-            def out_len(in_len, k, s, p):
-                return (in_len - k + 2 * p) // s + 1
-
-            out_h = out_len(h_in, kh, stride_h, pad_h)
-            out_w = out_len(w_in, kw, stride_w, pad_w)
-            ch_out = layer.out_channels
-            return ch_out, out_h, out_w
+    def get_output_shape(self) -> List[Tuple[int, int, int]]:
+        """Return the output shapes of the layer, including pooling.
 
-        conv_out_shape = get_shape_after_conv(
-            self.conv_layer, input_shape=self.input_shape
-        )
-        return conv_out_shape
-
-    def get_output_shape(self) -> Tuple[int, int, int]:
+        Returns:
+            One entry per destination, each formatted as (features, height, width).
+        """
         neuron_shape = self.get_neuron_shape()
         # this is the actual output shape, including pooling
-        if self.pool_layer is not None:
-            pool = expand_to_pair(self.pool_layer.kernel_size)
-            return (
+        output_shape = []
+        for pool in self.pool:
+            output_shape.append(
                 neuron_shape[0],
-                neuron_shape[1] // pool[0],
-                neuron_shape[2] // pool[1],
+                neuron_shape[1] // pool,
+                neuron_shape[2] // pool,
             )
-        else:
-            return neuron_shape
+        return output_shape
 
     def summary(self) -> dict:
+        """Returns a summary of the convolution's/pooling's kernel sizes and the output shape of the spiking layer."""
+
         return {
-            "pool": (
-                None if self.pool_layer is None else list(self.pool_layer.kernel_size)
-            ),
+            "pool": (self.pool),
             "kernel": list(self.conv_layer.weight.data.shape),
-            "neuron": self.get_neuron_shape(),
+            "neuron": self._get_conv_output_shape(),  # neuron layer output has the same shape as the convolution layer ouput.
         }
 
     def memory_summary(self):
@@ -177,13 +223,18 @@ class DynapcnnLayer(nn.Module):
 
             N_{MT} = f \\cdot 2^{ \\lceil \\log_2\\left(f_y\\right) \\rceil + \\lceil \\log_2\\left(f_x\\right) \\rceil }
 
-        Returns
-        -------
-        A dictionary with keys kernel, neuron and bias and the corresponding memory sizes
+        Returns:
+            A dictionary with keys kernel, neuron and bias and the corresponding memory sizes
         """
         summary = self.summary()
         f, c, h, w = summary["kernel"]
-        f, neuron_height, neuron_width = self.get_neuron_shape()
+        (
+            f,
+            neuron_height,
+            neuron_width,
+        ) = (
+            self._get_conv_output_shape()
+        )  # neuron layer output has the same shape as the convolution layer ouput.
 
         return {
             "kernel": c * pow(2, np.ceil(np.log2(h * w)) + np.ceil(np.log2(f))),
@@ -192,13 +243,28 @@ class DynapcnnLayer(nn.Module):
             "bias": 0 if self.conv_layer.bias is None else len(self.conv_layer.bias),
         }
 
-    def forward(self, x):
-        """Torch forward pass."""
-        x = self.conv_layer(x)
-        x = self.spk_layer(x)
-        if self.pool_layer is not None:
-            x = self.pool_layer(x)
-        return x
+    def _get_conv_output_shape(self) -> Tuple[int, int, int]:
+        """Computes the output dimensions of `conv_layer`.
 
-    def zero_grad(self, set_to_none: bool = False) -> None:
-        return self.spk_layer.zero_grad(set_to_none)
+        Returns:
+            output dimensions (tuple): a tuple describing `(output channels, height, width)`.
+        """
+        # get the layer's parameters.
+
+        out_channels = self.conv_layer.out_channels
+        kernel_size = self.conv_layer.kernel_size
+        stride = self.conv_layer.stride
+        padding = self.conv_layer.padding
+        dilation = self.conv_layer.dilation
+
+        # compute the output height and width.
+        out_height = (
+            (self.in_shape[1] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1)
+            // stride[0]
+        ) + 1
+        out_width = (
+            (self.in_shape[2] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1)
+            // stride[1]
+        ) + 1
+
+        return (out_channels, out_height, out_width)