From ed66314fa8cd3ca2c635b4de65d45efffbb87338 Mon Sep 17 00:00:00 2001 From: Kale Kundert Date: Tue, 31 Oct 2023 21:55:06 -0400 Subject: [PATCH 1/6] Make the pointwise pooling code more modular. - Move the Gaussian blurring functionality (used by the "antialiased" layers) into it's own module. - Consolidate the logic in all of the pointwise `evaluate_output_shape()` methods into a single standalone function. - Implement `forward()` using the exact same object returned by `export()`, to the extent possible. Signed-off-by: Kale Kundert --- escnn/nn/modules/pooling/__init__.py | 30 +- escnn/nn/modules/pooling/gaussian_blur.py | 120 ++++ escnn/nn/modules/pooling/pointwise.py | 326 ++++++++++ .../modules/pooling/pointwise_adaptive_avg.py | 181 ------ .../modules/pooling/pointwise_adaptive_max.py | 186 ------ escnn/nn/modules/pooling/pointwise_avg.py | 310 +++++---- escnn/nn/modules/pooling/pointwise_avg_3d.py | 254 -------- escnn/nn/modules/pooling/pointwise_max.py | 364 +++-------- escnn/nn/modules/pooling/utils.py | 8 + test/nn/test_pooling.py | 598 +++++++++++++++++- 10 files changed, 1275 insertions(+), 1102 deletions(-) create mode 100644 escnn/nn/modules/pooling/gaussian_blur.py create mode 100644 escnn/nn/modules/pooling/pointwise.py delete mode 100644 escnn/nn/modules/pooling/pointwise_adaptive_avg.py delete mode 100644 escnn/nn/modules/pooling/pointwise_adaptive_max.py delete mode 100644 escnn/nn/modules/pooling/pointwise_avg_3d.py create mode 100644 escnn/nn/modules/pooling/utils.py diff --git a/escnn/nn/modules/pooling/__init__.py b/escnn/nn/modules/pooling/__init__.py index 567433b0..8121af9e 100644 --- a/escnn/nn/modules/pooling/__init__.py +++ b/escnn/nn/modules/pooling/__init__.py @@ -1,19 +1,25 @@ from .norm_max import NormMaxPool -from .pointwise_max import PointwiseMaxPool2D, PointwiseMaxPool -from .pointwise_max import PointwiseMaxPoolAntialiased2D, PointwiseMaxPoolAntialiased -from .pointwise_max import PointwiseMaxPool3D -from .pointwise_max import PointwiseMaxPoolAntialiased3D -from .pointwise_avg import PointwiseAvgPool, PointwiseAvgPool2D -from .pointwise_avg import PointwiseAvgPoolAntialiased, PointwiseAvgPoolAntialiased2D -from .pointwise_adaptive_avg import PointwiseAdaptiveAvgPool2D, PointwiseAdaptiveAvgPool -from .pointwise_adaptive_avg import PointwiseAdaptiveAvgPool3D -from .pointwise_adaptive_max import PointwiseAdaptiveMaxPool2D, PointwiseAdaptiveMaxPool -from .pointwise_adaptive_max import PointwiseAdaptiveMaxPool3D +from .pointwise_max import ( + PointwiseMaxPool2D, PointwiseMaxPool, + PointwiseMaxPool3D, -from .pointwise_avg_3d import PointwiseAvgPool3D -from .pointwise_avg_3d import PointwiseAvgPoolAntialiased3D + PointwiseAdaptiveMaxPool2D, PointwiseAdaptiveMaxPool, + PointwiseAdaptiveMaxPool3D, + PointwiseMaxPoolAntialiased2D, PointwiseMaxPoolAntialiased, + PointwiseMaxPoolAntialiased3D, +) +from .pointwise_avg import ( + PointwiseAvgPool2D, PointwiseAvgPool, + PointwiseAvgPool3D, + + PointwiseAdaptiveAvgPool2D, PointwiseAdaptiveAvgPool, + PointwiseAdaptiveAvgPool3D, + + PointwiseAvgPoolAntialiased2D, PointwiseAvgPoolAntialiased, + PointwiseAvgPoolAntialiased3D, +) __all__ = [ "NormMaxPool", diff --git a/escnn/nn/modules/pooling/gaussian_blur.py b/escnn/nn/modules/pooling/gaussian_blur.py new file mode 100644 index 00000000..aec75c8e --- /dev/null +++ b/escnn/nn/modules/pooling/gaussian_blur.py @@ -0,0 +1,120 @@ + +from .utils import get_nd_tuple + +import torch +import torch.nn.functional as F + +from typing import Optional, Union, Tuple + +_CONV = { + 2: F.conv2d, + 3: F.conv3d, +} + +class GaussianBlurND(torch.nn.Module): + + def __init__( + self, + d: int, + *, + sigma: float = 0.6, + stride: Union[int, Tuple[int, ...]] = 1, + rel_padding: Optional[Union[int, Tuple[int, ...]]] = None, + abs_padding: Optional[Union[int, Tuple[int, ...]]] = None, + channels: int, + + # 4 for max pooling, 3 for average pooling. + _kernel_size_factor: float, + ): + """ + Apply a Gaussian blur to the input. + + This is equivalent to a depth-wise convolution with a Gaussian filter. + + Args: + d (int): Dimensionality of the base space (2 for images, 3 for volumes) + + sigma (float): Standard deviation for the Gaussian blur filter. + + stride (int): The stride of the blur filter. + + abs_padding: Implicit zero padding to be added on all sides of the + input, without regard to the size of the filter. Note that the + size of the filter depends on *sigma* and *_kernel_size_factor*, + and in this padding mode, the shape of the output tensor + depends on the size of the filter. It is an error to specify + *abs_padding* and *rel_padding*. + + rel_padding: Implicit zero padding to be added on all sides of the + input, treating the filter as if it were 1x1 (or 1x1x1), no + matter what size it really is. This means that the shape of + the output tensor is independent of the filter size. It is an + error to specify *abs_padding* and *rel_padding*. + + channels (int): The channel dimension of the input. + + _kernel_size_factor (float): How big the Gaussian blur filter + should be, in terms of *sigma*. See the code for the exact + expression, but the basic idea is that larger filters are + needed for larger standard deviations. + """ + super().__init__() + + assert sigma > 0. + + if rel_padding is not None and abs_padding is not None: + raise ValueError("can't specify `rel_padding` and `max_padding`") + + self.d = d + self.sigma = sigma + self.stride = stride + self.kernel_size = 2 * int(round(_kernel_size_factor * sigma)) + 1 + + # Build the Gaussian smoothing filter + + grid = torch.meshgrid( + *[torch.arange(self.kernel_size)] * d, + indexing='ij', + ) + grid = torch.stack(grid, dim=-1) + + mean = (self.kernel_size - 1) / 2. + variance = sigma ** 2. + + # setting the dtype is needed, otherwise it becomes an integer tensor + r = torch.sum((grid - mean) ** 2., dim=-1, dtype=torch.get_default_dtype()) + + # Build the gaussian kernel + _filter = torch.exp(-r / (2 * variance)) + + # Normalize + _filter /= torch.sum(_filter) + + # The filter needs to be reshaped to be used in depthwise convolution + _filter = _filter\ + .view(1, 1, *[self.kernel_size]*d)\ + .repeat((channels, 1, *[1]*d)) + + self.register_buffer('filter', _filter) + + if abs_padding is not None: + self.padding = abs_padding + else: + half_kernel_size = (self.kernel_size - 1) // 2 + self.padding = tuple( + p + half_kernel_size + for p in get_nd_tuple(rel_padding or 0, d) + ) + + def __repr__(self): + return f'{self.__class__.__name__}(d={self.d}, sigma={self.sigma}, stride={self.stride}, padding={self.padding}, channels={self.filter.shape[1]})' + + def forward(self, x): + return _CONV[self.d]( + x, + self.filter, + stride=self.stride, + padding=self.padding, + groups=x.shape[1], + ) + diff --git a/escnn/nn/modules/pooling/pointwise.py b/escnn/nn/modules/pooling/pointwise.py new file mode 100644 index 00000000..a979b758 --- /dev/null +++ b/escnn/nn/modules/pooling/pointwise.py @@ -0,0 +1,326 @@ + +from escnn.nn import FieldType +from escnn.nn import GeometricTensor +from escnn.gspaces import GSpace + +from ..equivariant_module import EquivariantModule +from .gaussian_blur import GaussianBlurND +from .utils import get_nd_tuple + +import torch +import math +from collections import OrderedDict + +from typing import List, Tuple, Any, Union, Optional + +_MAX_POOLS = { + 2: torch.nn.MaxPool2d, + 3: torch.nn.MaxPool3d, +} + +class _PointwisePoolND(EquivariantModule): + + def __init__( + self, + in_type: FieldType, + d: int, + pool: torch.nn.Module, + ): + r""" + + Channel-wise max-pooling: each channel is treated independently. + This module works exactly as :class:`torch.nn.MaxPool2D` or :class:`torch.nn.MaxPool3D`, wrapping it in the + :class:`~escnn.nn.EquivariantModule` interface. + + Notice that not all representations support this kind of pooling. In general, only representations which support + pointwise non-linearities do. + + .. warning :: + Even if the input tensor has a `coords` attribute, the output of this module will not have one. + + Args: + in_type (FieldType): the input field type + d (int): dimensionality of the base space (2 for images, 3 for volumes) + kernel_size: the size of the window to take a max over + stride: the stride of the window. Default value is :attr:`kernel_size` + padding: implicit zero padding to be added on both sides + dilation: a parameter that controls the stride of elements in the window + ceil_mode: when True, will use ceil instead of floor to compute the output shape + + """ + + super().__init__() + + check_dimensions(in_type, d) + check_pointwise_ok(in_type) + + self.d = d + self.space = in_type.gspace + self.in_type = in_type + self.out_type = in_type + self.pool = pool + + def forward(self, input: GeometricTensor) -> GeometricTensor: + r""" + + Args: + input (GeometricTensor): the input feature map + + Returns: + the resulting feature map + + """ + + assert input.type == self.in_type + + output = self.pool(input.tensor) + + return GeometricTensor(output, self.out_type, coords=None) + + def evaluate_output_shape(self, input_shape: Tuple) -> Tuple: + return calc_pool_output_shape(self.d, self.pool, input_shape, self.out_type) + + def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: + # this kind of pooling is not really equivariant so we can not test + # equivariance + pass + + def export(self): + r""" + Export this module to a normal PyTorch :class:`torch.nn.MaxPool2d` module and set to "eval" mode. + + """ + + return self.pool.eval() + +class _PointwiseAvgPoolAntialiasedND(EquivariantModule): + + def __init__( + self, + in_type: FieldType, + d: int, + *, + sigma: float, + stride: Union[int, Tuple[int, int]], + padding: Optional[Union[int, Tuple[int, int]]], + ): + r""" + + Antialiased channel-wise average-pooling: each channel is treated independently. + It performs strided convolution with a Gaussian blur filter. + + The size of the filter is computed as 3 standard deviations of the Gaussian curve. + By default, padding is added such that input size is preserved if stride is 1. + + Based on `Making Convolutional Networks Shift-Invariant Again `_. + + .. warning :: + Even if the input tensor has a `coords` attribute, the output of this module will not have one. + + Args: + in_type (FieldType): the input field type + sigma (float): standard deviation for the Gaussian blur filter + stride: the stride of the window. + padding: additional zero padding to be added on both sides + + """ + + super().__init__() + + check_dimensions(in_type, d) + # Don't need to check that the representation is compatible with + # pointwise nonlinearities, see #65. + + self.d = d + self.space = in_type.gspace + self.in_type = in_type + self.out_type = in_type + + self.blur = GaussianBlurND( + d=d, + sigma=sigma, + stride=stride, + abs_padding=padding, + channels=in_type.size, + _kernel_size_factor=3, + ) + + def forward(self, input: GeometricTensor) -> GeometricTensor: + r""" + + Args: + input (GeometricTensor): the input feature map + + Returns: + the resulting feature map + + """ + + assert input.type == self.in_type + + output = self.blur(input.tensor) + + return GeometricTensor(output, self.out_type, coords=None) + + def evaluate_output_shape(self, input_shape: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]: + return calc_pool_output_shape(self.d, self.blur, input_shape, self.out_type) + + def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: + # this kind of pooling is not really equivariant so we can't test + # equivariance + pass + + def export(self): + return self.blur.eval() + +class _PointwiseMaxPoolAntialiasedND(EquivariantModule): + + def __init__( + self, + in_type: FieldType, + d: int, + *, + kernel_size: Union[int, Tuple[int, ...]], + stride: Optional[Union[int, Tuple[int, ...]]], + padding: Union[int, Tuple[int, ...]], + ceil_mode: bool, + sigma: float, + ): + r""" + + Anti-aliased version of channel-wise max-pooling (each channel is treated independently). + + The max over a neighborhood is performed pointwise withot downsampling. + Then, convolution with a gaussian blurring filter is performed before downsampling the feature map. + + Based on `Making Convolutional Networks Shift-Invariant Again `_. + + + Notice that not all representations support this kind of pooling. In general, only representations which support + pointwise non-linearities do. + + .. warning :: + Even if the input tensor has a `coords` attribute, the output of this module will not have one. + + Args: + in_type (FieldType): the input field type + d (int): dimensionality of the base space (2 for images, 3 for volumes) + kernel_size: the size of the window to take a max over + stride: the stride of the window. Default value is :attr:`kernel_size` + padding: implicit zero padding to be added on both sides + dilation: a parameter that controls the stride of elements in the window + ceil_mode: when ``True``, will use ceil instead of floor to compute the output shape + sigma (float): standard deviation for the Gaussian blur filter + + """ + + super().__init__() + + if ceil_mode: + from warnings import warn + warn("The `ceil_mode` argument doesn't do anything, because the stride for the pooling step is always 1.") + + check_dimensions(in_type, d) + check_pointwise_ok(in_type) + + self.d = d + self.space = in_type.gspace + self.in_type = in_type + self.out_type = in_type + + pool = _MAX_POOLS[d]( + kernel_size=kernel_size, + stride=1, + padding=padding, + ) + blur = GaussianBlurND( + d=d, + sigma=sigma, + stride=stride if stride is not None else kernel_size, + rel_padding=padding, + channels=in_type.size, + _kernel_size_factor=4, + ) + self.layers = torch.nn.Sequential( + OrderedDict([('pool', pool), ('blur', blur)]), + ) + + def forward(self, input: GeometricTensor) -> GeometricTensor: + r""" + + Args: + input (GeometricTensor): the input feature map + + Returns: + the resulting feature map + + """ + + assert input.type == self.in_type + + output = self.layers(input.tensor) + + return GeometricTensor(output, self.out_type, coords=None) + + def evaluate_output_shape(self, input_shape: Tuple) -> Tuple: + for layer in self.layers: + input_shape = calc_pool_output_shape(self.d, layer, input_shape, self.out_type) + + return input_shape + + def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: + # this kind of pooling is not really equivariant so we can not test + # equivariance + pass + + def export(self): + return self.layers.eval() + +def check_dimensions(in_type: FieldType, d: int) -> None: + assert d in [2, 3], f"Only dimensionality 2 or 3 are currently suported by 'd={d}' found" + + assert isinstance(in_type.gspace, GSpace) + assert in_type.gspace.dimensionality == d, (in_type.gspace.dimensionality, d) + +def check_pointwise_ok(in_type: FieldType) -> None: + for r in in_type.representations: + assert 'pointwise' in r.supported_nonlinearities, \ + f"""Error! Representation "{r.name}" does not support pointwise non-linearities + so it is not possible to pool each channel independently""" + +def calc_pool_output_shape( + d: int, + pool: Any, + input_shape: Tuple[int, ...], + out_type: FieldType, +) -> Tuple[int, ...]: + + assert len(input_shape) == 2 + d, (input_shape, d) + + if hasattr(pool, 'output_size'): + output_size = get_nd_tuple(pool.output_size, d) + return (input_shape[0], out_type.size, *output_size) + + b, c, *xyz_i = input_shape + + kernel_size = get_nd_tuple(pool.kernel_size, d) + padding = get_nd_tuple(pool.padding, d) + stride = get_nd_tuple(pool.stride, d) + dilation = get_nd_tuple(getattr(pool, 'dilation', 1), d) + + # See online docs for `torch.nn.MaxPool2D`. + xyz_o = tuple( + (xyz_i[i] + 2 * padding[i] - dilation[i] * (kernel_size[i] - 1) - 1) / stride[i] + 1 + for i in range(d) + ) + + ceil_mode = getattr(pool, 'ceil_mode', False) + round = math.ceil if ceil_mode else math.floor + + xyz_o = tuple( + int(round(xyz_o[i])) + for i in range(d) + ) + + return (b, out_type.size, *xyz_o) + diff --git a/escnn/nn/modules/pooling/pointwise_adaptive_avg.py b/escnn/nn/modules/pooling/pointwise_adaptive_avg.py deleted file mode 100644 index 8be19335..00000000 --- a/escnn/nn/modules/pooling/pointwise_adaptive_avg.py +++ /dev/null @@ -1,181 +0,0 @@ - - -from escnn.gspaces import * -from escnn.nn import FieldType -from escnn.nn import GeometricTensor - -from ..equivariant_module import EquivariantModule - -import torch -import torch.nn.functional as F - -from typing import List, Tuple, Any, Union - -__all__ = [ - "PointwiseAdaptiveAvgPool2D", "PointwiseAdaptiveAvgPool", - "PointwiseAdaptiveAvgPool3D", -] - - -class _PointwiseAdaptiveAvgPoolND(EquivariantModule): - - def __init__(self, - in_type: FieldType, - d: int, - output_size: Union[int, Tuple[int, int]] - ): - r""" - - Adaptive channel-wise average-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.AdaptiveAvgPool2D`, wrapping it in - the :class:`~escnn.nn.EquivariantModule` interface. - - Notice that not all representations support this kind of pooling. In general, only representations which support - pointwise non-linearities do. - - .. warning :: - Even if the input tensor has a `coords` attribute, the output of this module will not have one. - - Args: - in_type (FieldType): the input field type - output_size: the target output size of the image of the form H x W - - """ - - assert d in [2, 3], f"Only dimensionality 2 or 3 are currently suported by 'd={d}' found" - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == d, (in_type.gspace.dimensionality, d) - - for r in in_type.representations: - assert 'pointwise' in r.supported_nonlinearities, \ - f"""Error! Representation "{r.name}" does not support pointwise non-linearities - so it is not possible to pool each channel independently""" - - super(_PointwiseAdaptiveAvgPoolND, self).__init__() - - self.d = d - self.space = in_type.gspace - self.in_type = in_type - self.out_type = in_type - - if isinstance(output_size, int): - self.output_size = (output_size,) * self.d - else: - self.output_size = output_size - - assert isinstance(self.output_size, tuple) and len(self.output_size) == self.d, self.output_size - - def forward(self, input: GeometricTensor) -> GeometricTensor: - r""" - - Args: - input (GeometricTensor): the input feature map - - Returns: - the resulting feature map - - """ - - assert input.type == self.in_type - - # run the common avg-pooling - if self.d == 2: - output = F.adaptive_avg_pool2d(input.tensor, self.output_size) - elif self.d == 3: - output = F.adaptive_avg_pool3d(input.tensor, self.output_size) - else: - raise NotImplementedError - - # wrap the result in a GeometricTensor - return GeometricTensor(output, self.out_type, coords=None) - - def evaluate_output_shape(self, input_shape: Tuple) -> Tuple: - assert len(input_shape) == 2 + self.d - assert input_shape[1] == self.in_type.size - return (input_shape[0], self.out_type.size, *self.output_size) - - def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: - - # this kind of pooling is not really equivariant so we can not test equivariance - pass - - def export(self): - r""" - Export this module to a normal PyTorch :class:`torch.nn.AdaptiveAvgPool2d` or - :class:`torch.nn.AdaptiveAvgPool3d` module and set to "eval" mode. - - """ - - self.eval() - - if self.d == 2: - return torch.nn.AdaptiveAvgPool2d(self.output_size).eval() - elif self.d == 3: - return torch.nn.AdaptiveAvgPool3d(self.output_size).eval() - else: - raise NotImplementedError - - -class PointwiseAdaptiveAvgPool2D(_PointwiseAdaptiveAvgPoolND): - - def __init__(self, - in_type: FieldType, - output_size: Union[int, Tuple[int, int]] - ): - r""" - - Adaptive channel-wise average-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.AdaptiveAvgPool2D`, wrapping it in - the :class:`~escnn.nn.EquivariantModule` interface. - - Notice that not all representations support this kind of pooling. In general, only representations which support - pointwise non-linearities do. - - .. warning :: - Even if the input tensor has a `coords` attribute, the output of this module will not have one. - - Args: - in_type (FieldType): the input field type - output_size: the target output size of the image of the form H x W - - """ - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == 2 - - super(PointwiseAdaptiveAvgPool2D, self).__init__(in_type, 2, output_size) - - -class PointwiseAdaptiveAvgPool3D(_PointwiseAdaptiveAvgPoolND): - - def __init__(self, - in_type: FieldType, - output_size: Union[int, Tuple[int, int]] - ): - r""" - - Adaptive channel-wise average-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.AdaptiveAvgPool3D`, wrapping it in - the :class:`~escnn.nn.EquivariantModule` interface. - - Notice that not all representations support this kind of pooling. In general, only representations which support - pointwise non-linearities do. - - .. warning :: - Even if the input tensor has a `coords` attribute, the output of this module will not have one. - - Args: - in_type (FieldType): the input field type - output_size: the target output size of the volume of the form H x W x D - - """ - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == 3 - - super(PointwiseAdaptiveAvgPool3D, self).__init__(in_type, 3, output_size) - - -# for backward compatibility -PointwiseAdaptiveAvgPool = PointwiseAdaptiveAvgPool2D diff --git a/escnn/nn/modules/pooling/pointwise_adaptive_max.py b/escnn/nn/modules/pooling/pointwise_adaptive_max.py deleted file mode 100644 index 6c17f875..00000000 --- a/escnn/nn/modules/pooling/pointwise_adaptive_max.py +++ /dev/null @@ -1,186 +0,0 @@ - - -from escnn.gspaces import * -from escnn.nn import FieldType -from escnn.nn import GeometricTensor - -from ..equivariant_module import EquivariantModule - -import torch -import torch.nn.functional as F - -from typing import List, Tuple, Any, Union - -__all__ = [ - "PointwiseAdaptiveMaxPool2D", "PointwiseAdaptiveMaxPool", - "PointwiseAdaptiveMaxPool3D", -] - - -class _PointwiseAdaptiveMaxPoolND(EquivariantModule): - - def __init__(self, - in_type: FieldType, - d: int, - output_size: Union[int, Tuple] - ): - r""" - - Module that implements adaptive channel-wise max-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.AdaptiveMaxPool2D` or :class:`torch.nn.AdaptiveMaxPool3D`, - wrapping it in the :class:`~escnn.nn.EquivariantModule` interface. - - Notice that not all representations support this kind of pooling. In general, only representations which support - pointwise non-linearities do. - - .. warning :: - Even if the input tensor has a `coords` attribute, the output of this module will not have one. - - Args: - in_type (FieldType): the input field type - d (int): dimensionality of the base space (2 for images, 3 for volumes) - output_size: the target output size - - """ - - assert d in [2, 3], f"Only dimensionality 2 or 3 are currently suported by 'd={d}' found" - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == d, (in_type.gspace.dimensionality, d) - - for r in in_type.representations: - assert 'pointwise' in r.supported_nonlinearities, \ - f"""Error! Representation "{r}" does not support pointwise non-linearities - so it is not possible to pool each channel independently""" - - super(_PointwiseAdaptiveMaxPoolND, self).__init__() - - self.d = d - self.space = in_type.gspace - self.in_type = in_type - self.out_type = in_type - - if isinstance(output_size, int): - self.output_size = (output_size,) * self.d - else: - self.output_size = output_size - - assert isinstance(self.output_size, tuple) and len(self.output_size) == self.d, self.output_size - - def forward(self, input: GeometricTensor) -> GeometricTensor: - r""" - - Args: - input (GeometricTensor): the input feature map - - Returns: - the resulting feature map - - """ - - assert input.type == self.in_type - - # run the common max-pooling - if self.d == 2: - output = F.adaptive_max_pool2d(input.tensor, self.output_size) - elif self.d == 3: - output = F.adaptive_max_pool3d(input.tensor, self.output_size) - else: - raise NotImplementedError - - # wrap the result in a GeometricTensor - return GeometricTensor(output, self.out_type, coords=None) - - def evaluate_output_shape(self, input_shape: Tuple) -> Tuple: - assert len(input_shape) == 2 + self.d - assert input_shape[1] == self.in_type.size - return (input_shape[0], self.out_type.size, *self.output_size) - - def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: - - # this kind of pooling is not really equivariant so we can not test equivariance - pass - - def export(self): - r""" - Export this module to a normal PyTorch :class:`torch.nn.AdaptiveMaxPool2d` or - :class:`torch.nn.AdaptiveMaxPool3d` module and set to "eval" mode. - - """ - - self.eval() - - if self.d == 2: - return torch.nn.AdaptiveMaxPool2d(self.output_size).eval() - elif self.d == 3: - return torch.nn.AdaptiveMaxPool3d(self.output_size).eval() - else: - raise NotImplementedError - - -class PointwiseAdaptiveMaxPool2D(_PointwiseAdaptiveMaxPoolND): - - def __init__(self, - in_type: FieldType, - output_size: Union[int, Tuple[int, int]] - ): - r""" - - Module that implements adaptive channel-wise max-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.AdaptiveMaxPool2D`, wrapping it in the - :class:`~escnn.nn.EquivariantModule` interface. - - Notice that not all representations support this kind of pooling. In general, only representations which support - pointwise non-linearities do. - - .. warning :: - Even if the input tensor has a `coords` attribute, the output of this module will not have one. - - Args: - in_type (FieldType): the input field type - output_size: the target output size of the image of the form H x W - - """ - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == 2 - - super(PointwiseAdaptiveMaxPool2D, self).__init__( - in_type, 2, output_size - ) - - -class PointwiseAdaptiveMaxPool3D(_PointwiseAdaptiveMaxPoolND): - - def __init__(self, - in_type: FieldType, - output_size: Union[int, Tuple[int, int]] - ): - r""" - - Module that implements adaptive channel-wise max-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.AdaptiveMaxPool3D`, wrapping it in the - :class:`~escnn.nn.EquivariantModule` interface. - - Notice that not all representations support this kind of pooling. In general, only representations which support - pointwise non-linearities do. - - .. warning :: - Even if the input tensor has a `coords` attribute, the output of this module will not have one. - - Args: - in_type (FieldType): the input field type - output_size: the target output size of the image of the form H x W - - """ - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == 3 - - super(PointwiseAdaptiveMaxPool3D, self).__init__( - in_type, 3, output_size - ) - - -# for backward compatibility -PointwiseAdaptiveMaxPool = PointwiseAdaptiveMaxPool2D \ No newline at end of file diff --git a/escnn/nn/modules/pooling/pointwise_avg.py b/escnn/nn/modules/pooling/pointwise_avg.py index 29c92c05..1427bd48 100644 --- a/escnn/nn/modules/pooling/pointwise_avg.py +++ b/escnn/nn/modules/pooling/pointwise_avg.py @@ -1,17 +1,12 @@ -from escnn.gspaces import * from escnn.nn import FieldType -from escnn.nn import GeometricTensor -from ..equivariant_module import EquivariantModule +from .pointwise import _PointwisePoolND, _PointwiseAvgPoolAntialiasedND -import torch.nn.functional as F import torch -from typing import List, Tuple, Any, Union - -import math +from typing import Tuple, Union, Optional __all__ = [ "PointwiseAvgPool", @@ -21,19 +16,19 @@ ] -class PointwiseAvgPool2D(EquivariantModule): +class PointwiseAvgPool2D(_PointwisePoolND): def __init__(self, in_type: FieldType, kernel_size: Union[int, Tuple[int, int]], - stride: Union[int, Tuple[int, int]] = None, + stride: Optional[Union[int, Tuple[int, int]]] = None, padding: Union[int, Tuple[int, int]] = 0, ceil_mode: bool = False ): r""" Channel-wise average-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.AvgPool2D`, wrapping it in the + This module works exactly as :class:`torch.nn.AvgPool2d`, wrapping it in the :class:`~escnn.nn.EquivariantModule` interface. .. warning :: @@ -47,97 +42,122 @@ def __init__(self, ceil_mode: when ``True``, will use ceil instead of floor to compute the output shape """ + super().__init__( + in_type, 2, + pool=torch.nn.AvgPool2d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + ceil_mode=ceil_mode, + ), + ) + + +class PointwiseAvgPool3D(_PointwisePoolND): + + def __init__(self, + in_type: FieldType, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Optional[Union[int, Tuple[int, int, int]]] = None, + padding: Union[int, Tuple[int, int, int]] = 0, + ceil_mode: bool = False + ): + r""" - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == 2 - - # for r in in_type.representations: - # assert 'pointwise' in r.supported_nonlinearities, \ - # """Error! Representation "{}" does not support pointwise non-linearities - # so it is not possible to pool each channel independently""" + Channel-wise average-pooling: each channel is treated independently. + This module works exactly as :class:`torch.nn.AvgPool3d`, wrapping it in the + :class:`~escnn.nn.EquivariantModule` interface. - super(PointwiseAvgPool2D, self).__init__() - - self.space = in_type.gspace - self.in_type = in_type - self.out_type = in_type - - if isinstance(kernel_size, int): - self.kernel_size = (kernel_size, kernel_size) - else: - self.kernel_size = kernel_size - - if isinstance(stride, int): - self.stride = (stride, stride) - elif stride is None: - self.stride = self.kernel_size - else: - self.stride = stride - - if isinstance(padding, int): - self.padding = (padding, padding) - else: - self.padding = padding - - self.ceil_mode = ceil_mode + .. warning :: + Even if the input tensor has a `coords` attribute, the output of this module will not have one. - def forward(self, input: GeometricTensor) -> GeometricTensor: - r""" - Args: - input (GeometricTensor): the input feature map + in_type (FieldType): the input field type + kernel_size: the size of the window to take a average over + stride: the stride of the window. Default value is :attr:`kernel_size` + padding: implicit zero padding to be added on both sides + ceil_mode: when ``True``, will use ceil instead of floor to compute the output shape - Returns: - the resulting feature map - """ - - assert input.type == self.in_type - - # run the common max-pooling - output = F.avg_pool2d(input.tensor, - kernel_size=self.kernel_size, - stride=self.stride, - padding=self.padding, - ceil_mode=self.ceil_mode) - - # wrap the result in a GeometricTensor - return GeometricTensor(output, self.out_type, coords=None) - - def evaluate_output_shape(self, input_shape: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]: - assert len(input_shape) == 4 - assert input_shape[1] == self.in_type.size + super().__init__( + in_type, 3, + pool=torch.nn.AvgPool3d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + ceil_mode=ceil_mode, + ), + ) - b, c, hi, wi = input_shape + +class PointwiseAdaptiveAvgPool2D(_PointwisePoolND): - # compute the output shape (see 'torch.nn.MaxPool2D') - ho = (hi + 2 * self.padding[0] - (self.kernel_size[0] - 1) - 1) / self.stride[0] + 1 - wo = (wi + 2 * self.padding[1] - (self.kernel_size[1] - 1) - 1) / self.stride[1] + 1 - - if self.ceil_mode: - ho = math.ceil(ho) - wo = math.ceil(wo) - else: - ho = math.floor(ho) - wo = math.floor(wo) + def __init__(self, + in_type: FieldType, + output_size: Union[int, Tuple[int, int]] + ): + r""" - return b, self.out_type.size, ho, wo + Adaptive channel-wise average-pooling: each channel is treated independently. + This module works exactly as :class:`torch.nn.AdaptiveAvgPool2d`, wrapping it in + the :class:`~escnn.nn.EquivariantModule` interface. - def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: - - # this kind of pooling is not really equivariant so we can not test equivariance - pass - - -class PointwiseAvgPoolAntialiased2D(EquivariantModule): + Notice that not all representations support this kind of pooling. In general, only representations which support + pointwise non-linearities do. + + .. warning :: + Even if the input tensor has a `coords` attribute, the output of this module will not have one. + + Args: + in_type (FieldType): the input field type + output_size: the target output size of the image of the form H x W + + """ + super().__init__( + in_type, 2, + pool=torch.nn.AdaptiveAvgPool2d( + output_size=output_size, + ), + ) + + +class PointwiseAdaptiveAvgPool3D(_PointwisePoolND): + + def __init__(self, + in_type: FieldType, + output_size: Union[int, Tuple[int, int, int]] + ): + r""" + + Adaptive channel-wise average-pooling: each channel is treated independently. + This module works exactly as :class:`torch.nn.AdaptiveAvgPool3d`, wrapping it in + the :class:`~escnn.nn.EquivariantModule` interface. + + Notice that not all representations support this kind of pooling. In general, only representations which support + pointwise non-linearities do. + + .. warning :: + Even if the input tensor has a `coords` attribute, the output of this module will not have one. + + Args: + in_type (FieldType): the input field type + output_size: the target output size of the volume of the form H x W x D + + """ + super().__init__( + in_type, 3, + pool=torch.nn.AdaptiveAvgPool3d( + output_size=output_size, + ), + ) + +class PointwiseAvgPoolAntialiased2D(_PointwiseAvgPoolAntialiasedND): def __init__(self, in_type: FieldType, sigma: float, stride: Union[int, Tuple[int, int]], - # kernel_size: Union[int, Tuple[int, int]] = None, - padding: Union[int, Tuple[int, int]] = None, - #dilation: Union[int, Tuple[int, int]] = 1, + padding: Optional[Union[int, Tuple[int, int]]] = None, ): r""" @@ -159,98 +179,50 @@ def __init__(self, padding: additional zero padding to be added on both sides """ - - super(PointwiseAvgPoolAntialiased2D, self).__init__() + super().__init__( + in_type, 2, + sigma=sigma, + stride=stride, + padding=padding, + ) + +class PointwiseAvgPoolAntialiased3D(_PointwiseAvgPoolAntialiasedND): + + def __init__(self, + in_type: FieldType, + sigma: float, + stride: Union[int, Tuple[int, int, int]], + padding: Optional[Union[int, Tuple[int, int, int]]] = None, + ): + r""" - self.space = in_type.gspace - self.in_type = in_type - self.out_type = in_type - - assert sigma > 0. + Antialiased channel-wise average-pooling: each channel is treated independently. + It performs strided convolution with a Gaussian blur filter. - filter_size = 2*int(round(3*sigma))+1 + The size of the filter is computed as 3 standard deviations of the Gaussian curve. + By default, padding is added such that input size is preserved if stride is 1. - self.kernel_size = (filter_size, filter_size) + Inspired by `Making Convolutional Networks Shift-Invariant Again `_. - if isinstance(stride, int): - self.stride = (stride, stride) - elif stride is None: - self.stride = self.kernel_size - else: - self.stride = stride - - if padding is None: - padding = int((filter_size-1)//2) + .. warning :: + Even if the input tensor has a `coords` attribute, the output of this module will not have one. - if isinstance(padding, int): - self.padding = (padding, padding) - else: - self.padding = padding - - # Build the Gaussian smoothing filter - - grid_x = torch.arange(filter_size).repeat(filter_size).view(filter_size, filter_size) - grid_y = grid_x.t() - grid = torch.stack([grid_x, grid_y], dim=-1) - - mean = (filter_size - 1) / 2. - variance = sigma ** 2. - - # setting the dtype is needed, otherwise it becomes an integer tensor - r = -torch.sum((grid - mean) ** 2., dim=-1, dtype=torch.get_default_dtype()) - - # Build the gaussian kernel - _filter = torch.exp(r / (2 * variance)) - - # Normalize - _filter /= torch.sum(_filter) - - # The filter needs to be reshaped to be used in 2d depthwise convolution - _filter = _filter.view(1, 1, filter_size, filter_size).repeat((in_type.size, 1, 1, 1)) - - self.register_buffer('filter', _filter) - - ################################################################################################################ - - def forward(self, input: GeometricTensor) -> GeometricTensor: - r""" - Args: - input (GeometricTensor): the input feature map - - Returns: - the resulting feature map + in_type (FieldType): the input field type + sigma (float): standard deviation for the Gaussian blur filter + stride: the stride of the window. + padding: additional zero padding to be added on both sides """ + super().__init__( + in_type, 3, + sigma=sigma, + stride=stride, + padding=padding, + ) - assert input.type == self.in_type - - output = F.conv2d(input.tensor, self.filter, stride=self.stride, padding=self.padding, groups=input.shape[1]) - - # wrap the result in a GeometricTensor - return GeometricTensor(output, self.out_type, coords=None) - - def evaluate_output_shape(self, input_shape: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]: - assert len(input_shape) == 4 - assert input_shape[1] == self.in_type.size - - b, c, hi, wi = input_shape - - # compute the output shape - ho = (hi + 2 * self.padding[0] - (self.kernel_size[0] - 1) - 1) / self.stride[0] + 1 - wo = (wi + 2 * self.padding[1] - (self.kernel_size[1] - 1) - 1) / self.stride[1] + 1 - - ho = math.floor(ho) - wo = math.floor(wo) - - return b, self.out_type.size, ho, wo - - def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: - - # this kind of pooling is not really equivariant so we can't test equivariance - pass - # for backward compatibility PointwiseAvgPool = PointwiseAvgPool2D -PointwiseAvgPoolAntialiased = PointwiseAvgPoolAntialiased2D \ No newline at end of file +PointwiseAdaptiveAvgPool = PointwiseAdaptiveAvgPool2D +PointwiseAvgPoolAntialiased = PointwiseAvgPoolAntialiased2D diff --git a/escnn/nn/modules/pooling/pointwise_avg_3d.py b/escnn/nn/modules/pooling/pointwise_avg_3d.py deleted file mode 100644 index 045006fc..00000000 --- a/escnn/nn/modules/pooling/pointwise_avg_3d.py +++ /dev/null @@ -1,254 +0,0 @@ - - -from escnn.gspaces import * -from escnn.nn import FieldType -from escnn.nn import GeometricTensor - -from ..equivariant_module import EquivariantModule - -import torch.nn.functional as F -import torch - -from typing import List, Tuple, Any, Union - -import math - -__all__ = ["PointwiseAvgPool3D", "PointwiseAvgPoolAntialiased3D"] - - -class PointwiseAvgPool3D(EquivariantModule): - - def __init__(self, - in_type: FieldType, - kernel_size: Union[int, Tuple[int, int]], - stride: Union[int, Tuple[int, int]] = None, - padding: Union[int, Tuple[int, int]] = 0, - ceil_mode: bool = False - ): - r""" - - Channel-wise average-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.AvgPool3D`, wrapping it in the - :class:`~escnn.nn.EquivariantModule` interface. - - .. warning :: - Even if the input tensor has a `coords` attribute, the output of this module will not have one. - - Args: - in_type (FieldType): the input field type - kernel_size: the size of the window to take a average over - stride: the stride of the window. Default value is :attr:`kernel_size` - padding: implicit zero padding to be added on both sides - ceil_mode: when ``True``, will use ceil instead of floor to compute the output shape - - """ - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == 3 - - # for r in in_type.representations: - # assert 'pointwise' in r.supported_nonlinearities, \ - # """Error! Representation "{}" does not support pointwise non-linearities - # so it is not possible to pool each channel independently""" - - super(PointwiseAvgPool3D, self).__init__() - - self.space = in_type.gspace - self.in_type = in_type - self.out_type = in_type - - if isinstance(kernel_size, int): - self.kernel_size = (kernel_size, kernel_size, kernel_size) - else: - self.kernel_size = kernel_size - - if isinstance(stride, int): - self.stride = (stride, stride, stride) - elif stride is None: - self.stride = self.kernel_size - else: - self.stride = stride - - if isinstance(padding, int): - self.padding = (padding, padding, padding) - else: - self.padding = padding - - self.ceil_mode = ceil_mode - - def forward(self, input: GeometricTensor) -> GeometricTensor: - r""" - - Args: - input (GeometricTensor): the input feature map - - Returns: - the resulting feature map - - """ - - assert input.type == self.in_type - - # run the common max-pooling - output = F.avg_pool3d(input.tensor, - kernel_size=self.kernel_size, - stride=self.stride, - padding=self.padding, - ceil_mode=self.ceil_mode) - - # wrap the result in a GeometricTensor - return GeometricTensor(output, self.out_type, coords=None) - - def evaluate_output_shape(self, input_shape: Tuple[int, int, int, int, int]) -> Tuple[int, int, int, int, int]: - assert len(input_shape) == 5 - assert input_shape[1] == self.in_type.size - - b, c, hi, wi, di = input_shape - - # compute the output shape (see 'torch.nn.MaxPool2D') - ho = (hi + 2 * self.padding[0] - (self.kernel_size[0] - 1) - 1) / self.stride[0] + 1 - wo = (wi + 2 * self.padding[1] - (self.kernel_size[1] - 1) - 1) / self.stride[1] + 1 - do = (di + 2 * self.padding[2] - (self.kernel_size[2] - 1) - 1) / self.stride[2] + 1 - - if self.ceil_mode: - ho = math.ceil(ho) - wo = math.ceil(wo) - do = math.ceil(do) - else: - ho = math.floor(ho) - wo = math.floor(wo) - do = math.floor(do) - - return b, self.out_type.size, ho, wo, do - - def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: - - # this kind of pooling is not really equivariant so we can not test equivariance - pass - - -class PointwiseAvgPoolAntialiased3D(EquivariantModule): - - def __init__(self, - in_type: FieldType, - sigma: float, - stride: Union[int, Tuple[int, int]], - padding: Union[int, Tuple[int, int]] = None, - ): - r""" - - Antialiased channel-wise average-pooling: each channel is treated independently. - It performs strided convolution with a Gaussian blur filter. - - The size of the filter is computed as 3 standard deviations of the Gaussian curve. - By default, padding is added such that input size is preserved if stride is 1. - - Inspired by `Making Convolutional Networks Shift-Invariant Again `_. - - .. warning :: - Even if the input tensor has a `coords` attribute, the output of this module will not have one. - - Args: - in_type (FieldType): the input field type - sigma (float): standard deviation for the Gaussian blur filter - stride: the stride of the window. - padding: additional zero padding to be added on both sides - - """ - - super(PointwiseAvgPoolAntialiased3D, self).__init__() - - self.space = in_type.gspace - self.in_type = in_type - self.out_type = in_type - - assert sigma > 0. - - filter_size = 2*int(round(3*sigma))+1 - - self.kernel_size = (filter_size, filter_size, filter_size) - - if isinstance(stride, int): - self.stride = (stride, stride, stride) - elif stride is None: - self.stride = self.kernel_size - else: - self.stride = stride - - if padding is None: - padding = int((filter_size-1)//2) - - if isinstance(padding, int): - self.padding = (padding, padding, padding) - else: - self.padding = padding - - # Build the Gaussian smoothing filter - - # grid_x = torch.arange(filter_size).repeat(filter_size).view(filter_size, filter_size) - # grid_y = grid_x.t() - # grid = torch.stack([grid_x, grid_y], dim=-1) - grid_x, grid_y, grid_z = torch.meshgrid( - [torch.arange(filter_size)] * 3, - indexing='ij', - ) - grid = torch.stack([grid_x, grid_y, grid_z], dim=-1) - - mean = (filter_size - 1) / 2. - variance = sigma ** 2. - - # setting the dtype is needed, otherwise it becomes an integer tensor - r = -torch.sum((grid - mean) ** 2., dim=-1, dtype=torch.get_default_dtype()) - - # Build the gaussian kernel - filter = torch.exp(r / (2 * variance)) - - # Normalize - filter /= torch.sum(filter) - - # The filter needs to be reshaped to be used in 2d depthwise convolution - filter = filter.view(1, 1, filter_size, filter_size, filter_size).repeat((in_type.size, 1, 1, 1, 1)) - - self.register_buffer('filter', filter) - - ################################################################################################################ - - def forward(self, input: GeometricTensor) -> GeometricTensor: - r""" - - Args: - input (GeometricTensor): the input feature map - - Returns: - the resulting feature map - - """ - - assert input.type == self.in_type - - output = F.conv3d(input.tensor, self.filter, stride=self.stride, padding=self.padding, groups=input.shape[1]) - - # wrap the result in a GeometricTensor - return GeometricTensor(output, self.out_type, coords=None) - - def evaluate_output_shape(self, input_shape: Tuple[int, int, int, int, int]) -> Tuple[int, int, int, int, int]: - assert len(input_shape) == 5 - assert input_shape[1] == self.in_type.size - - b, c, hi, wi, di = input_shape - - # compute the output shape - ho = (hi + 2 * self.padding[0] - (self.kernel_size[0] - 1) - 1) / self.stride[0] + 1 - wo = (wi + 2 * self.padding[1] - (self.kernel_size[1] - 1) - 1) / self.stride[1] + 1 - do = (di + 2 * self.padding[2] - (self.kernel_size[2] - 1) - 1) / self.stride[2] + 1 - - ho = math.floor(ho) - wo = math.floor(wo) - do = math.floor(do) - - return b, self.out_type.size, ho, wo, do - - def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: - - # this kind of pooling is not really equivariant so we can't test equivariance - pass diff --git a/escnn/nn/modules/pooling/pointwise_max.py b/escnn/nn/modules/pooling/pointwise_max.py index 7fe1503e..9375460e 100644 --- a/escnn/nn/modules/pooling/pointwise_max.py +++ b/escnn/nn/modules/pooling/pointwise_max.py @@ -1,16 +1,11 @@ -from escnn.gspaces import * from escnn.nn import FieldType -from escnn.nn import GeometricTensor -from ..equivariant_module import EquivariantModule +from .pointwise import _PointwisePoolND, _PointwiseMaxPoolAntialiasedND -import torch.nn.functional as F import torch -from typing import List, Tuple, Any, Union - -import math +from typing import Tuple, Union, Optional __all__ = [ "PointwiseMaxPool2D", "PointwiseMaxPool", @@ -19,33 +14,30 @@ "PointwiseMaxPoolAntialiased3D", ] - -class _PointwiseMaxPoolND(EquivariantModule): +class PointwiseMaxPool2D(_PointwisePoolND): def __init__(self, in_type: FieldType, - d: int, - kernel_size: Union[int, Tuple], - stride: Union[int, Tuple] = None, - padding: Union[int, Tuple] = 0, - dilation: Union[int, Tuple] = 1, + kernel_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]] = None, + padding: Union[int, Tuple[int, int]] = 0, + dilation: Union[int, Tuple[int, int]] = 1, ceil_mode: bool = False ): r""" Channel-wise max-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.MaxPool2D` or :class:`torch.nn.MaxPool3D`, wrapping it in the + This module works exactly as :class:`torch.nn.MaxPool2d`, wrapping it in the :class:`~escnn.nn.EquivariantModule` interface. - + Notice that not all representations support this kind of pooling. In general, only representations which support pointwise non-linearities do. .. warning :: Even if the input tensor has a `coords` attribute, the output of this module will not have one. - + Args: in_type (FieldType): the input field type - d (int): dimensionality of the base space (2 for images, 3 for volumes) kernel_size: the size of the window to take a max over stride: the stride of the window. Default value is :attr:`kernel_size` padding: implicit zero padding to be added on both sides @@ -53,150 +45,33 @@ def __init__(self, ceil_mode: when True, will use ceil instead of floor to compute the output shape """ - - assert d in [2, 3], f"Only dimensionality 2 or 3 are currently suported by 'd={d}' found" - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == d, (in_type.gspace.dimensionality, d) - - for r in in_type.representations: - assert 'pointwise' in r.supported_nonlinearities, \ - f"""Error! Representation "{r.name}" does not support pointwise non-linearities - so it is not possible to pool each channel independently""" - - super(_PointwiseMaxPoolND, self).__init__() - - self.d = d - self.space = in_type.gspace - self.in_type = in_type - self.out_type = in_type - - if isinstance(kernel_size, int): - self.kernel_size = (kernel_size,) * self.d - else: - self.kernel_size = kernel_size - - if isinstance(stride, int): - self.stride = (stride,) * self.d - elif stride is None: - self.stride = self.kernel_size - else: - self.stride = stride - - if isinstance(padding, int): - self.padding = (padding,) * self.d - else: - self.padding = padding - - if isinstance(dilation, int): - self.dilation = (dilation,) * self.d - else: - self.dilation = dilation - - self.ceil_mode = ceil_mode - - def forward(self, input: GeometricTensor) -> GeometricTensor: - r""" - - Args: - input (GeometricTensor): the input feature map - - Returns: - the resulting feature map - - """ - - assert input.type == self.in_type - - # run the common max-pooling - if self.d == 2: - output = F.max_pool2d(input.tensor, - self.kernel_size, - self.stride, - self.padding, - self.dilation, - self.ceil_mode) - elif self.d == 3: - output = F.max_pool3d(input.tensor, - self.kernel_size, - self.stride, - self.padding, - self.dilation, - self.ceil_mode) - else: - raise NotImplementedError - - # wrap the result in a GeometricTensor - return GeometricTensor(output, self.out_type, coords=None) - - def evaluate_output_shape(self, input_shape: Tuple) -> Tuple: - assert len(input_shape) == 2 + self.d, (input_shape, self.d) - assert input_shape[1] == self.in_type.size - - b, c = input_shape[:2] - shape_i = input_shape[2:] - - # compute the output shape (see 'torch.nn.MaxPool2D') - shape_o = tuple( - (shape_i[i] + 2 * self.padding[i] - self.dilation[i] * (self.kernel_size[i] - 1) - 1) / self.stride[i] + 1 - for i in range(self.d) + super().__init__( + in_type, 2, + pool=torch.nn.MaxPool2d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + ceil_mode=ceil_mode, + ), ) - if self.ceil_mode: - shape_o = tuple( - math.ceil(shape_o[i]) - for i in range(self.d) - ) - else: - shape_o = tuple( - math.floor(shape_o[i]) - for i in range(self.d) - ) - - return (b, self.out_type.size, *shape_o) - - def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]: - # this kind of pooling is not really equivariant so we can not test equivariance - pass - - def export(self): - r""" - Export this module to a normal PyTorch :class:`torch.nn.MaxPool2d` module and set to "eval" mode. - - """ - - self.eval() - - if self.d == 2: - return torch.nn.MaxPool2d(self.kernel_size, self.stride, self.padding, self.dilation).eval() - elif self.d == 3: - return torch.nn.MaxPool3d(self.kernel_size, self.stride, self.padding, self.dilation).eval() - else: - raise NotImplementedError - - -class _PointwiseMaxPoolAntialiasedND(_PointwiseMaxPoolND): +class PointwiseMaxPool3D(_PointwisePoolND): def __init__(self, in_type: FieldType, - d: int, - kernel_size: Union[int, Tuple[int, int]], - stride: Union[int, Tuple[int, int]] = None, - padding: Union[int, Tuple[int, int]] = 0, - dilation: Union[int, Tuple[int, int]] = 1, - ceil_mode: bool = False, - sigma: float = 0.6, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]] = None, + padding: Union[int, Tuple[int, int, int]] = 0, + dilation: Union[int, Tuple[int, int, int]] = 1, + ceil_mode: bool = False ): r""" - Anti-aliased version of channel-wise max-pooling (each channel is treated independently). - - The max over a neighborhood is performed pointwise withot downsampling. - Then, convolution with a gaussian blurring filter is performed before downsampling the feature map. - - Based on `Making Convolutional Networks Shift-Invariant Again `_. - + Channel-wise max-pooling: each channel is treated independently. + This module works exactly as :class:`torch.nn.MaxPool3d`, wrapping it in the + :class:`~escnn.nn.EquivariantModule` interface. Notice that not all representations support this kind of pooling. In general, only representations which support pointwise non-linearities do. @@ -206,142 +81,66 @@ def __init__(self, Args: in_type (FieldType): the input field type - d (int): dimensionality of the base space (2 for images, 3 for volumes) kernel_size: the size of the window to take a max over stride: the stride of the window. Default value is :attr:`kernel_size` padding: implicit zero padding to be added on both sides dilation: a parameter that controls the stride of elements in the window - ceil_mode: when ``True``, will use ceil instead of floor to compute the output shape - sigma (float): standard deviation for the Gaussian blur filter + ceil_mode: when True, will use ceil instead of floor to compute the output shape """ + super().__init__( + in_type, 3, + pool=torch.nn.MaxPool3d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + ceil_mode=ceil_mode, + ), + ) - if dilation != 1: - raise NotImplementedError("Dilation larger than 1 is not supported yet") - - super(_PointwiseMaxPoolAntialiasedND, self).__init__(in_type, d, kernel_size, stride, padding, dilation, ceil_mode) - - assert sigma > 0. - - filter_size = 2 * int(round(4 * sigma)) + 1 - - # Build the Gaussian smoothing filter - - grid = torch.stack(torch.meshgrid(*[torch.arange(filter_size)] * self.d, indexing='ij'), dim=-1) - - mean = (filter_size - 1) / 2. - variance = sigma ** 2. - - # setting the dtype is needed, otherwise it becomes an integer tensor - r = -torch.sum((grid - mean) ** 2., dim=-1, dtype=torch.get_default_dtype()) - - # Build the gaussian kernel - _filter = torch.exp(r / (2 * variance)) - - # Normalize - _filter /= torch.sum(_filter) - - # The filter needs to be reshaped to be used in 2d depthwise convolution - _filter = _filter.view(1, 1, *[filter_size]*self.d).repeat((in_type.size, 1, *[1]*self.d)) - - self.register_buffer('filter', _filter) - self._pad = tuple(p + int((filter_size - 1) // 2) for p in self.padding) - - def forward(self, input: GeometricTensor) -> GeometricTensor: - r""" - - Args: - input (GeometricTensor): the input feature map - - Returns: - the resulting feature map - - """ - assert input.type == self.in_type - - if self.d == 2: - # evaluate the max operation densely (stride = 1) - output = F.max_pool2d(input.tensor, - self.kernel_size, - 1, - self.padding, - self.dilation, - self.ceil_mode) - output = F.conv2d(output, self.filter, stride=self.stride, padding=self._pad, groups=output.shape[1]) - elif self.d == 3: - # evaluate the max operation densely (stride = 1) - output = F.max_pool3d(input.tensor, - self.kernel_size, - 1, - self.padding, - self.dilation, - self.ceil_mode) - output = F.conv3d(output, self.filter, stride=self.stride, padding=self._pad, groups=output.shape[1]) - else: - raise NotImplementedError - - # wrap the result in a GeometricTensor - return GeometricTensor(output, self.out_type, coords=None) - - def export(self): - raise NotImplementedError - - -class PointwiseMaxPool2D(_PointwiseMaxPoolND): +class PointwiseAdaptiveMaxPool2D(_PointwisePoolND): def __init__(self, in_type: FieldType, - kernel_size: Union[int, Tuple[int, int]], - stride: Union[int, Tuple[int, int]] = None, - padding: Union[int, Tuple[int, int]] = 0, - dilation: Union[int, Tuple[int, int]] = 1, - ceil_mode: bool = False + output_size: Union[int, Tuple[int, int]] ): r""" - Channel-wise max-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.MaxPool2D`, wrapping it in the + Module that implements adaptive channel-wise max-pooling: each channel is treated independently. + This module works exactly as :class:`torch.nn.AdaptiveMaxPool2d`, wrapping it in the :class:`~escnn.nn.EquivariantModule` interface. - + Notice that not all representations support this kind of pooling. In general, only representations which support pointwise non-linearities do. .. warning :: Even if the input tensor has a `coords` attribute, the output of this module will not have one. - + Args: in_type (FieldType): the input field type - kernel_size: the size of the window to take a max over - stride: the stride of the window. Default value is :attr:`kernel_size` - padding: implicit zero padding to be added on both sides - dilation: a parameter that controls the stride of elements in the window - ceil_mode: when True, will use ceil instead of floor to compute the output shape + output_size: the target output size of the image of the form H x W """ - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == 2 - - super(PointwiseMaxPool2D, self).__init__( - in_type, 2, kernel_size, stride, padding, dilation, ceil_mode + super().__init__( + in_type, 2, + pool=torch.nn.AdaptiveMaxPool2d( + output_size=output_size, + ), ) -class PointwiseMaxPool3D(_PointwiseMaxPoolND): +class PointwiseAdaptiveMaxPool3D(_PointwisePoolND): def __init__(self, in_type: FieldType, - kernel_size: Union[int, Tuple[int, int]], - stride: Union[int, Tuple[int, int]] = None, - padding: Union[int, Tuple[int, int]] = 0, - dilation: Union[int, Tuple[int, int]] = 1, - ceil_mode: bool = False + output_size: Union[int, Tuple[int, int, int]] ): r""" - Channel-wise max-pooling: each channel is treated independently. - This module works exactly as :class:`torch.nn.MaxPool3D`, wrapping it in the + Module that implements adaptive channel-wise max-pooling: each channel is treated independently. + This module works exactly as :class:`torch.nn.AdaptiveMaxPool3d`, wrapping it in the :class:`~escnn.nn.EquivariantModule` interface. Notice that not all representations support this kind of pooling. In general, only representations which support @@ -352,30 +151,23 @@ def __init__(self, Args: in_type (FieldType): the input field type - kernel_size: the size of the window to take a max over - stride: the stride of the window. Default value is :attr:`kernel_size` - padding: implicit zero padding to be added on both sides - dilation: a parameter that controls the stride of elements in the window - ceil_mode: when True, will use ceil instead of floor to compute the output shape + output_size: the target output size of the image of the form H x W """ - - assert isinstance(in_type.gspace, GSpace) - assert in_type.gspace.dimensionality == 3 - - super(PointwiseMaxPool3D, self).__init__( - in_type, 3, kernel_size, stride, padding, dilation, ceil_mode + super().__init__( + in_type, 3, + pool=torch.nn.AdaptiveMaxPool3d( + output_size=output_size, + ), ) - class PointwiseMaxPoolAntialiased2D(_PointwiseMaxPoolAntialiasedND): def __init__(self, in_type: FieldType, kernel_size: Union[int, Tuple[int, int]], - stride: Union[int, Tuple[int, int]] = None, + stride: Optional[Union[int, Tuple[int, int]]] = None, padding: Union[int, Tuple[int, int]] = 0, - dilation: Union[int, Tuple[int, int]] = 1, ceil_mode: bool = False, sigma: float = 0.6, ): @@ -383,8 +175,8 @@ def __init__(self, Anti-aliased version of channel-wise max-pooling (each channel is treated independently). - The max over a neighborhood is performed pointwise withot downsampling. - Then, convolution with a gaussian blurring filter is performed before downsampling the feature map. + The max over a neighborhood is performed pointwise without downsampling. + Then, convolution with a Gaussian blurring filter is performed before downsampling the feature map. Based on `Making Convolutional Networks Shift-Invariant Again `_. @@ -405,9 +197,13 @@ def __init__(self, sigma (float): standard deviation for the Gaussian blur filter """ - - super(PointwiseMaxPoolAntialiased2D, self).__init__( - in_type, 2, kernel_size, stride, padding, dilation, ceil_mode, sigma + super().__init__( + in_type, 2, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ceil_mode=ceil_mode, + sigma=sigma, ) @@ -415,10 +211,9 @@ class PointwiseMaxPoolAntialiased3D(_PointwiseMaxPoolAntialiasedND): def __init__(self, in_type: FieldType, - kernel_size: Union[int, Tuple[int, int]], - stride: Union[int, Tuple[int, int]] = None, - padding: Union[int, Tuple[int, int]] = 0, - dilation: Union[int, Tuple[int, int]] = 1, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Optional[Union[int, Tuple[int, int, int]]] = None, + padding: Union[int, Tuple[int, int, int]] = 0, ceil_mode: bool = False, sigma: float = 0.6, ): @@ -448,14 +243,19 @@ def __init__(self, sigma (float): standard deviation for the Gaussian blur filter """ - - super(PointwiseMaxPoolAntialiased3D, self).__init__( - in_type, 3, kernel_size, stride, padding, dilation, ceil_mode, sigma + super().__init__( + in_type, 3, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ceil_mode=ceil_mode, + sigma=sigma, ) # for backward compatibility PointwiseMaxPool = PointwiseMaxPool2D +PointwiseAdaptiveMaxPool = PointwiseAdaptiveMaxPool2D PointwiseMaxPoolAntialiased = PointwiseMaxPoolAntialiased2D diff --git a/escnn/nn/modules/pooling/utils.py b/escnn/nn/modules/pooling/utils.py new file mode 100644 index 00000000..70ed2a9a --- /dev/null +++ b/escnn/nn/modules/pooling/utils.py @@ -0,0 +1,8 @@ +from numbers import Number + +def get_nd_tuple(scalar_or_tuple, d): + if isinstance(scalar_or_tuple, Number): + return (scalar_or_tuple,) * d + else: + return scalar_or_tuple + diff --git a/test/nn/test_pooling.py b/test/nn/test_pooling.py index 69d65e43..02881d4e 100644 --- a/test/nn/test_pooling.py +++ b/test/nn/test_pooling.py @@ -176,7 +176,7 @@ def test_induced_norm_pooling(self): .format(el, errs.max(), errs.mean(), errs.var()) # Not all modules are properly equivariant so it doesn't make sense to check their equivariance error - # At least, let's check they forward without raising any issues + # At least, let's check that they pool correctly. def test_PointwiseMaxPool2D(self): gs = rot2dOnR2() @@ -184,12 +184,27 @@ def test_PointwiseMaxPool2D(self): m = PointwiseMaxPool2D(ft, 3, 2, 0) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [0, 4, 0, 1, 0, 0], + [1, 0, 1, 0, 1, 1], + [1, 2, 0, 3, 0, 2], + [0, 0, 1, 0, 1, 1], + [1, 0, 1, 3, 2, 0], + ]) + x = x.view(1, 1, 5, 6) x = ft(x) m.eval() y = m(x) + # Manually verified. + y_expected = torch.Tensor([ + [4, 3], + [2, 3], + ]) + + torch.testing.assert_close(y.tensor, y_expected.view(1, 1, 2, 2)) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseMaxPoolAntialiased2D(self): @@ -198,12 +213,32 @@ def test_PointwiseMaxPoolAntialiased2D(self): m = PointwiseMaxPoolAntialiased2D(ft, 3, 2, 0, sigma=0.4) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [0, 4, 0, 1, 0, 0], + [1, 0, 1, 0, 1, 1], + [1, 2, 0, 3, 0, 2], + [0, 0, 1, 0, 1, 1], + [1, 0, 1, 3, 2, 0], + ]) + x = x.view(1, 1, 5, 6) x = ft(x) m.eval() y = m(x) + # Manually verified. + y_expected = torch.Tensor([ + [3.6075, 2.9159], + [1.8805, 2.8788], + ]) + + torch.testing.assert_close( + y.tensor, + y_expected.view(1, 1, 2, 2), + atol=1e-4, + rtol=0, + ) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseMaxPool3D(self): @@ -213,11 +248,59 @@ def test_PointwiseMaxPool3D(self): m = PointwiseMaxPool3D(ft, 3, 2, 0) x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + + x = torch.Tensor([ + [[0, 0, 0, 0, 1, 0, 0], + [2, 2, 1, 0, 2, 1, 1], + [0, 0, 1, 0, 1, 3, 1], + [1, 0, 2, 1, 2, 0, 1], + [1, 0, 0, 0, 1, 2, 2], + [3, 2, 2, 0, 1, 0, 1]], + + [[0, 0, 0, 2, 1, 0, 2], + [2, 1, 0, 1, 0, 0, 1], + [1, 0, 1, 3, 0, 1, 2], + [2, 1, 0, 1, 1, 0, 1], + [1, 1, 0, 1, 1, 1, 0], + [1, 0, 1, 1, 4, 1, 1]], + + [[1, 1, 2, 1, 0, 1, 1], + [1, 1, 1, 1, 1, 0, 1], + [0, 1, 0, 2, 0, 1, 2], + [2, 1, 2, 0, 0, 2, 0], + [2, 1, 0, 2, 2, 2, 3], + [2, 1, 1, 0, 0, 1, 0]], + + [[0, 2, 1, 2, 5, 1, 0], + [2, 1, 1, 5, 0, 1, 0], + [0, 2, 1, 0, 3, 1, 0], + [1, 1, 2, 0, 2, 1, 0], + [1, 0, 1, 2, 1, 0, 0], + [0, 1, 2, 0, 2, 2, 2]], + + [[1, 1, 0, 1, 1, 0, 2], + [1, 0, 1, 2, 1, 0, 1], + [3, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 2, 0, 0], + [3, 1, 0, 1, 1, 1, 1], + [1, 3, 1, 1, 2, 2, 0]], + ]) + x = x.view(1, 1, 5, 6, 7) x = ft(x) m.eval() y = m(x) + # Manually verified. + y_expected = torch.Tensor([ + [[2, 3, 3], + [2, 3, 3]], + [[3, 5, 5], + [3, 3, 3]], + ]) + + torch.testing.assert_close(y.tensor, y_expected.view(1, 1, 2, 2, 3)) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseMaxPoolAntialiased3D(self): @@ -226,12 +309,63 @@ def test_PointwiseMaxPoolAntialiased3D(self): m = PointwiseMaxPoolAntialiased3D(ft, 3, 2, 0, sigma=0.4) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [[0, 0, 0, 0, 1, 0, 0], + [2, 2, 1, 0, 2, 1, 1], + [0, 0, 1, 0, 1, 3, 1], + [1, 0, 2, 1, 2, 0, 1], + [1, 0, 0, 0, 1, 2, 2], + [3, 2, 2, 0, 1, 0, 1]], + + [[0, 0, 0, 2, 1, 0, 2], + [2, 1, 0, 1, 0, 0, 1], + [1, 0, 1, 3, 0, 1, 2], + [2, 1, 0, 1, 1, 0, 1], + [1, 1, 0, 1, 1, 1, 0], + [1, 0, 1, 1, 4, 1, 1]], + + [[1, 1, 2, 1, 0, 1, 1], + [1, 1, 1, 1, 1, 0, 1], + [0, 1, 0, 2, 0, 1, 2], + [2, 1, 2, 0, 0, 2, 0], + [2, 1, 0, 2, 2, 2, 3], + [2, 1, 1, 0, 0, 1, 0]], + + [[0, 2, 1, 2, 5, 1, 0], + [2, 1, 1, 5, 0, 1, 0], + [0, 2, 1, 0, 3, 1, 0], + [1, 1, 2, 0, 2, 1, 0], + [1, 0, 1, 2, 1, 0, 0], + [0, 1, 2, 0, 2, 2, 2]], + + [[1, 1, 0, 1, 1, 0, 2], + [1, 0, 1, 2, 1, 0, 1], + [3, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 2, 0, 0], + [3, 1, 0, 1, 1, 1, 1], + [1, 3, 1, 1, 2, 2, 0]], + ]) + x = x.view(1, 1, 5, 6, 7) x = ft(x) m.eval() y = m(x) + # Not manually verified. + y_expected = torch.Tensor([ + [[1.8076, 2.8401, 2.7224], + [1.9131, 2.9177, 2.7999]], + [[2.6897, 4.6042, 4.3470], + [2.6943, 2.8881, 2.7657]], + ]) + + torch.testing.assert_close( + y.tensor, + y_expected.view(1, 1, 2, 2, 3), + atol=1e-4, + rtol=0, + ) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseAvgPool2D(self): @@ -240,12 +374,27 @@ def test_PointwiseAvgPool2D(self): m = PointwiseAvgPool2D(ft, 3, 2, 0) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [0, 4, 0, 1, 0, 0], + [1, 0, 1, 0, 1, 1], + [1, 2, 0, 3, 0, 2], + [0, 0, 1, 0, 1, 1], + [1, 0, 1, 3, 2, 0], + ]) + x = x.view(1, 1, 5, 6) x = ft(x) m.eval() y = m(x) + # Manually verified. + y_expected = torch.Tensor([ + [ 9/9, 6/9], + [ 6/9, 11/9], + ]) + + torch.testing.assert_close(y.tensor, y_expected.view(1, 1, 2, 2)) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseAvgPoolAntialiased2D(self): @@ -254,12 +403,31 @@ def test_PointwiseAvgPoolAntialiased2D(self): m = PointwiseAvgPoolAntialiased2D(ft, 0.6, 2, 0) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [0, 4, 0, 1, 0, 0], + [1, 0, 1, 0, 1, 1], + [1, 2, 0, 3, 0, 2], + [0, 0, 1, 0, 1, 1], + [1, 0, 1, 3, 2, 0], + ]) + x = x.view(1, 1, 5, 6) x = ft(x) m.eval() y = m(x) + # Manually verified. + y_expected = torch.Tensor([ + [0.7773], + ]) + + torch.testing.assert_close( + y.tensor, + y_expected.view(1, 1, 1, 1), + atol=1e-4, + rtol=0, + ) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseAvgPool3D(self): @@ -268,12 +436,58 @@ def test_PointwiseAvgPool3D(self): m = PointwiseAvgPool3D(ft, 3, 2, 0) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [[0, 0, 0, 0, 1, 0, 0], + [2, 2, 1, 0, 2, 1, 1], + [0, 0, 1, 0, 1, 3, 1], + [1, 0, 2, 1, 2, 0, 1], + [1, 0, 0, 0, 1, 2, 2], + [3, 2, 2, 0, 1, 0, 1]], + + [[0, 0, 0, 2, 1, 0, 2], + [2, 1, 0, 1, 0, 0, 1], + [1, 0, 1, 3, 0, 1, 2], + [2, 1, 0, 1, 1, 0, 1], + [1, 1, 0, 1, 1, 1, 0], + [1, 0, 1, 1, 4, 1, 1]], + + [[1, 1, 2, 1, 0, 1, 1], + [1, 1, 1, 1, 1, 0, 1], + [0, 1, 0, 2, 0, 1, 2], + [2, 1, 2, 0, 0, 2, 0], + [2, 1, 0, 2, 2, 2, 3], + [2, 1, 1, 0, 0, 1, 0]], + + [[0, 2, 1, 2, 5, 1, 0], + [2, 1, 1, 5, 0, 1, 0], + [0, 2, 1, 0, 3, 1, 0], + [1, 1, 2, 0, 2, 1, 0], + [1, 0, 1, 2, 1, 0, 0], + [0, 1, 2, 0, 2, 2, 2]], + + [[1, 1, 0, 1, 1, 0, 2], + [1, 0, 1, 2, 1, 0, 1], + [3, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 2, 0, 0], + [3, 1, 0, 1, 1, 1, 1], + [1, 3, 1, 1, 2, 2, 0]], + ]) + x = x.view(1, 1, 5, 6, 7) x = ft(x) m.eval() y = m(x) + # Manually verified. + y_expected = torch.Tensor([ + [[19/27, 22/27, 24/27], + [21/27, 24/27, 32/27]], + [[26/27, 32/27, 23/27], + [27/27, 25/27, 25/27]], + ]) + + torch.testing.assert_close(y.tensor, y_expected.view(1, 1, 2, 2, 3)) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseAvgPoolAntialiased3D(self): @@ -282,70 +496,418 @@ def test_PointwiseAvgPoolAntialiased3D(self): m = PointwiseAvgPoolAntialiased3D(ft, 0.6, 2, 0) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [[0, 0, 0, 0, 1, 0, 0], + [2, 2, 1, 0, 2, 1, 1], + [0, 0, 1, 0, 1, 3, 1], + [1, 0, 2, 1, 2, 0, 1], + [1, 0, 0, 0, 1, 2, 2], + [3, 2, 2, 0, 1, 0, 1]], + + [[0, 0, 0, 2, 1, 0, 2], + [2, 1, 0, 1, 0, 0, 1], + [1, 0, 1, 3, 0, 1, 2], + [2, 1, 0, 1, 1, 0, 1], + [1, 1, 0, 1, 1, 1, 0], + [1, 0, 1, 1, 4, 1, 1]], + + [[1, 1, 2, 1, 0, 1, 1], + [1, 1, 1, 1, 1, 0, 1], + [0, 1, 0, 2, 0, 1, 2], + [2, 1, 2, 0, 0, 2, 0], + [2, 1, 0, 2, 2, 2, 3], + [2, 1, 1, 0, 0, 1, 0]], + + [[0, 2, 1, 2, 5, 1, 0], + [2, 1, 1, 5, 0, 1, 0], + [0, 2, 1, 0, 3, 1, 0], + [1, 1, 2, 0, 2, 1, 0], + [1, 0, 1, 2, 1, 0, 0], + [0, 1, 2, 0, 2, 2, 2]], + + [[1, 1, 0, 1, 1, 0, 2], + [1, 0, 1, 2, 1, 0, 1], + [3, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 2, 0, 0], + [3, 1, 0, 1, 1, 1, 1], + [1, 3, 1, 1, 2, 2, 0]], + ]) + x = x.view(1, 1, 5, 6, 7) x = ft(x) m.eval() y = m(x) + # Not manually verified. + y_expected = torch.Tensor([ + [0.8444, 0.7675], + ]) + + torch.testing.assert_close( + y.tensor, + y_expected.view(1, 1, 1, 1, 2), + atol=1e-4, + rtol=0, + ) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseAdaptiveAvgPool2D(self): gs = rot2dOnR2() ft = gs.type(gs.trivial_repr) - m = PointwiseAdaptiveAvgPool2D(ft, 3) + m = PointwiseAdaptiveAvgPool2D(ft, 2) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [0, 4, 0, 1, 0, 0], + [1, 0, 1, 0, 1, 1], + [1, 2, 0, 3, 0, 2], + [0, 0, 1, 0, 1, 1], + [1, 0, 1, 3, 2, 0], + ]) + x = x.view(1, 1, 5, 6) x = ft(x) m.eval() y = m(x) + # Manually verified, assuming kernel_size=3 and stride=(2, 3) + y_expected = torch.Tensor([ + [ 9/9, 8/9], + [ 6/9, 12/9], + ]) + + torch.testing.assert_close(y.tensor, y_expected.view(1, 1, 2, 2)) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseAdaptiveAvgPool3D(self): gs = rot2dOnR3() ft = gs.type(gs.trivial_repr) - m = PointwiseAdaptiveAvgPool3D(ft, 3) - - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + m = PointwiseAdaptiveAvgPool3D(ft, (2, 2, 3)) + + x = torch.Tensor([ + [[0, 0, 0, 0, 1, 0, 0], + [2, 2, 1, 0, 2, 1, 1], + [0, 0, 1, 0, 1, 3, 1], + [1, 0, 2, 1, 2, 0, 1], + [1, 0, 0, 0, 1, 2, 2], + [3, 2, 2, 0, 1, 0, 1]], + + [[0, 0, 0, 2, 1, 0, 2], + [2, 1, 0, 1, 0, 0, 1], + [1, 0, 1, 3, 0, 1, 2], + [2, 1, 0, 1, 1, 0, 1], + [1, 1, 0, 1, 1, 1, 0], + [1, 0, 1, 1, 4, 1, 1]], + + [[1, 1, 2, 1, 0, 1, 1], + [1, 1, 1, 1, 1, 0, 1], + [0, 1, 0, 2, 0, 1, 2], + [2, 1, 2, 0, 0, 2, 0], + [2, 1, 0, 2, 2, 2, 3], + [2, 1, 1, 0, 0, 1, 0]], + + [[0, 2, 1, 2, 5, 1, 0], + [2, 1, 1, 5, 0, 1, 0], + [0, 2, 1, 0, 3, 1, 0], + [1, 1, 2, 0, 2, 1, 0], + [1, 0, 1, 2, 1, 0, 0], + [0, 1, 2, 0, 2, 2, 2]], + + [[1, 1, 0, 1, 1, 0, 2], + [1, 0, 1, 2, 1, 0, 1], + [3, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 2, 0, 0], + [3, 1, 0, 1, 1, 1, 1], + [1, 3, 1, 1, 2, 2, 0]], + ]) + x = x.view(1, 1, 5, 6, 7) x = ft(x) m.eval() y = m(x) + # Manually verified, assuming kernel_size=3 and stride=(2,3,2). + y_expected = torch.Tensor([ + [[19/27, 22/27, 24/27], + [30/27, 26/27, 30/27]], + [[26/27, 32/27, 23/27], + [31/27, 28/27, 29/27]], + ]) + + torch.testing.assert_close(y.tensor, y_expected.view(1, 1, 2, 2, 3)) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseAdaptiveMaxPool2D(self): gs = rot2dOnR2() ft = gs.type(gs.trivial_repr) - m = PointwiseAdaptiveMaxPool2D(ft, 3) + m = PointwiseAdaptiveMaxPool2D(ft, 2) - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + x = torch.Tensor([ + [0, 4, 0, 1, 0, 0], + [1, 0, 1, 0, 1, 1], + [1, 2, 0, 3, 0, 2], + [0, 0, 1, 0, 1, 1], + [1, 0, 1, 3, 2, 0], + ]) + x = x.view(1, 1, 5, 6) x = ft(x) m.eval() y = m(x) + # Manually verified, assuming kernel_size=3 and stride=(2, 3) + y_expected = torch.Tensor([ + [4, 3], + [2, 3], + ]) + + torch.testing.assert_close(y.tensor, y_expected.view(1, 1, 2, 2)) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) def test_PointwiseAdaptiveMaxPool3D(self): gs = rot2dOnR3() ft = gs.type(gs.trivial_repr) - m = PointwiseAdaptiveMaxPool3D(ft, 3) - - x = torch.randn(5, ft.size, *[9+i for i in range(gs.dimensionality)]) + m = PointwiseAdaptiveMaxPool3D(ft, (2, 2, 3)) + + x = torch.Tensor([ + [[0, 0, 0, 0, 1, 0, 0], + [2, 2, 1, 0, 2, 1, 1], + [0, 0, 1, 0, 1, 3, 1], + [1, 0, 2, 1, 2, 0, 1], + [1, 0, 0, 0, 1, 2, 2], + [3, 2, 2, 0, 1, 0, 1]], + + [[0, 0, 0, 2, 1, 0, 2], + [2, 1, 0, 1, 0, 0, 1], + [1, 0, 1, 3, 0, 1, 2], + [2, 1, 0, 1, 1, 0, 1], + [1, 1, 0, 1, 1, 1, 0], + [1, 0, 1, 1, 4, 1, 1]], + + [[1, 1, 2, 1, 0, 1, 1], + [1, 1, 1, 1, 1, 0, 1], + [0, 1, 0, 2, 0, 1, 2], + [2, 1, 2, 0, 0, 2, 0], + [2, 1, 0, 2, 2, 2, 3], + [2, 1, 1, 0, 0, 1, 0]], + + [[0, 2, 1, 2, 5, 1, 0], + [2, 1, 1, 5, 0, 1, 0], + [0, 2, 1, 0, 3, 1, 0], + [1, 1, 2, 0, 2, 1, 0], + [1, 0, 1, 2, 1, 0, 0], + [0, 1, 2, 0, 2, 2, 2]], + + [[1, 1, 0, 1, 1, 0, 2], + [1, 0, 1, 2, 1, 0, 1], + [3, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 2, 0, 0], + [3, 1, 0, 1, 1, 1, 1], + [1, 3, 1, 1, 2, 2, 0]], + ]) + x = x.view(1, 1, 5, 6, 7) x = ft(x) m.eval() y = m(x) + # Manually verified, assuming kernel_size=3 and stride=(2,3,2). + y_expected = torch.Tensor([ + [[2, 3, 3], + [3, 4, 4]], + [[3, 5, 5], + [3, 2, 3]], + ]) + + torch.testing.assert_close(y.tensor, y_expected.view(1, 1, 2, 2, 3)) + self.assertEqual(y.shape, m.evaluate_output_shape(x.shape)) + def test_output_shape_2d(self): + # The main purpose of this function is to make sure that all of the + # arguments to the pooling module have the expected effect, i.e. they + # aren't ignored or misinterpreted. + + gs = rot2dOnR2() + ft = gs.type(gs.trivial_repr) + + cases = [ + # Pointwise pooling: + # Use `PointwiseMaxPool2D` as a representative of all the + # `_PointwisePoolND` subclasses. + dict( + module=PointwiseMaxPool2D( + ft, + kernel_size=3, + stride=2, + ), + in_shape=(2, 1, 5, 5), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPool2D( + ft, + kernel_size=5, + stride=2, + ), + in_shape=(2, 1, 7, 7), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPool2D( + ft, + kernel_size=3, + stride=1, + ), + in_shape=(2, 1, 4, 4), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPool2D( + ft, + kernel_size=3, + stride=None, # Means same as kernel size + ), + in_shape=(2, 1, 6, 6), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPool2D( + ft, + kernel_size=3, + stride=2, + padding=1, + ), + in_shape=(2, 1, 3, 3), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPool2D( + ft, + kernel_size=3, + stride=2, + dilation=2, + ), + in_shape=(2, 1, 7, 7), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPool2D( + ft, + kernel_size=3, + stride=2, + ceil_mode=True, + ), + in_shape=(2, 1, 4, 4), + out_shape=(2, 1, 2, 2), + ), + + # Antialiased average pooling: + dict( + module=PointwiseAvgPoolAntialiased2D( + ft, + sigma=0.6, # kernel size: 5 + stride=2, + ), + in_shape=(2, 1, 3, 3), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseAvgPoolAntialiased2D( + ft, + sigma=0.4, # kernel size: 3 + stride=2, + ), + in_shape=(2, 1, 3, 3), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseAvgPoolAntialiased2D( + ft, + sigma=0.6, # kernel size: 5 + stride=1, + ), + in_shape=(2, 1, 2, 2), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseAvgPoolAntialiased2D( + ft, + sigma=0.6, # kernel size: 5 + stride=2, + padding=0, + ), + in_shape=(2, 1, 7, 7), + out_shape=(2, 1, 2, 2), + ), + + # Antialiased max pooling: + dict( + module=PointwiseMaxPoolAntialiased2D( + ft, + kernel_size=3, + ), + in_shape=(2, 1, 6, 6), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPoolAntialiased2D( + ft, + kernel_size=5, + ), + in_shape=(2, 1, 10, 10), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPoolAntialiased2D( + ft, + kernel_size=3, + # This value of sigma reduces the size of the blur + # filter from 5 (the default) to 3. But because the + # blur filter is hard-coded to use "relative padding", + # the size of the blur filter has no effect on the size + # of the output. + sigma=0.3, + ), + in_shape=(2, 1, 6, 6), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPoolAntialiased2D( + ft, + kernel_size=3, + stride=2, + ), + in_shape=(2, 1, 5, 5), + out_shape=(2, 1, 2, 2), + ), + dict( + module=PointwiseMaxPoolAntialiased2D( + ft, + kernel_size=3, + padding=1 + ), + in_shape=(2, 1, 4, 4), + out_shape=(2, 1, 2, 2), + ), + ] + + for case in cases: + f = case['module'] + + with self.subTest(f): + x = ft(torch.zeros(*case['in_shape'])) + y = f(x) + + self.assertEqual(y.shape, case['out_shape']) + self.assertEqual(y.shape, f.evaluate_output_shape(x.shape)) if __name__ == '__main__': unittest.main() From cee4e58911aa0fdfe93bd98774fca343acd7be35 Mon Sep 17 00:00:00 2001 From: Kale Kundert Date: Wed, 1 Nov 2023 17:09:46 -0400 Subject: [PATCH 2/6] Make it easier to customize the kernel size of the Gaussian blur Signed-off-by: Kale Kundert --- escnn/nn/modules/pooling/gaussian_blur.py | 73 ++++++++++++----------- escnn/nn/modules/pooling/pointwise.py | 12 ++-- 2 files changed, 44 insertions(+), 41 deletions(-) diff --git a/escnn/nn/modules/pooling/gaussian_blur.py b/escnn/nn/modules/pooling/gaussian_blur.py index aec75c8e..deb16ad7 100644 --- a/escnn/nn/modules/pooling/gaussian_blur.py +++ b/escnn/nn/modules/pooling/gaussian_blur.py @@ -15,16 +15,14 @@ class GaussianBlurND(torch.nn.Module): def __init__( self, - d: int, *, - sigma: float = 0.6, + sigma: float, + kernel_size: int, stride: Union[int, Tuple[int, ...]] = 1, + padding: Optional[Union[int, Tuple[int, ...]]] = None, rel_padding: Optional[Union[int, Tuple[int, ...]]] = None, - abs_padding: Optional[Union[int, Tuple[int, ...]]] = None, + d: int, channels: int, - - # 4 for max pooling, 3 for average pooling. - _kernel_size_factor: float, ): """ Apply a Gaussian blur to the input. @@ -32,53 +30,56 @@ def __init__( This is equivalent to a depth-wise convolution with a Gaussian filter. Args: - d (int): Dimensionality of the base space (2 for images, 3 for volumes) + sigma (float): Standard deviation of the Gaussian making up the + blur filter. - sigma (float): Standard deviation for the Gaussian blur filter. + kernel_size (int): Size of the convolutional filter used to apply + the blur. Note that this should be related to the value of + *sigma*; larger standard deviations require larger kernels to + overlap the same density. You can think of the Gaussian blur + as being truncated to zero beyond the bounds of the filter. - stride (int): The stride of the blur filter. + stride (int): Stride of the convolutional filter used to apply the + blur. - abs_padding: Implicit zero padding to be added on all sides of the - input, without regard to the size of the filter. Note that the - size of the filter depends on *sigma* and *_kernel_size_factor*, - and in this padding mode, the shape of the output tensor - depends on the size of the filter. It is an error to specify - *abs_padding* and *rel_padding*. + padding: Implicit zero padding to be added on all sides of the + input, without regard to the size of the filter. It is an + error to specify *padding* and *rel_padding*. rel_padding: Implicit zero padding to be added on all sides of the - input, treating the filter as if it were 1x1 (or 1x1x1), no - matter what size it really is. This means that the shape of - the output tensor is independent of the filter size. It is an - error to specify *abs_padding* and *rel_padding*. + input, treating the filter as if it were 1x1 (or 1x1x1, etc.), + no matter what size it really is. This means that the shape of + the output tensor is independent of the filter size. This is + helpful when, for example, the *kernel_size* argument is + dynamically calculated as a function of *sigma*. It is an + error to specify *padding* and *rel_padding*. - channels (int): The channel dimension of the input. + d (int): Dimensionality of the base space (2 for images, 3 for + volumes) - _kernel_size_factor (float): How big the Gaussian blur filter - should be, in terms of *sigma*. See the code for the exact - expression, but the basic idea is that larger filters are - needed for larger standard deviations. + channels (int): Channel dimension of the input. """ super().__init__() assert sigma > 0. - if rel_padding is not None and abs_padding is not None: - raise ValueError("can't specify `rel_padding` and `max_padding`") + if padding is not None and rel_padding is not None: + raise ValueError("can't specify `padding` and `rel_padding`") - self.d = d + self.kernel_size = kernel_size self.sigma = sigma self.stride = stride - self.kernel_size = 2 * int(round(_kernel_size_factor * sigma)) + 1 + self.d = d # Build the Gaussian smoothing filter grid = torch.meshgrid( - *[torch.arange(self.kernel_size)] * d, + *[torch.arange(kernel_size)] * d, indexing='ij', ) grid = torch.stack(grid, dim=-1) - mean = (self.kernel_size - 1) / 2. + mean = (kernel_size - 1) / 2. variance = sigma ** 2. # setting the dtype is needed, otherwise it becomes an integer tensor @@ -92,22 +93,22 @@ def __init__( # The filter needs to be reshaped to be used in depthwise convolution _filter = _filter\ - .view(1, 1, *[self.kernel_size]*d)\ + .view(1, 1, *[kernel_size]*d)\ .repeat((channels, 1, *[1]*d)) self.register_buffer('filter', _filter) - if abs_padding is not None: - self.padding = abs_padding + if padding is not None: + self.padding = padding else: - half_kernel_size = (self.kernel_size - 1) // 2 + half_kernel_size = (kernel_size - 1) // 2 self.padding = tuple( p + half_kernel_size for p in get_nd_tuple(rel_padding or 0, d) ) def __repr__(self): - return f'{self.__class__.__name__}(d={self.d}, sigma={self.sigma}, stride={self.stride}, padding={self.padding}, channels={self.filter.shape[1]})' + return f'{self.__class__.__name__}(sigma={self.sigma}, kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}, d={self.d}, channels={self.filter.shape[1]})' def forward(self, x): return _CONV[self.d]( @@ -118,3 +119,5 @@ def forward(self, x): groups=x.shape[1], ) +def kernel_size_from_radius(radius): + return 2 * int(round(radius)) + 1 diff --git a/escnn/nn/modules/pooling/pointwise.py b/escnn/nn/modules/pooling/pointwise.py index a979b758..bed3f671 100644 --- a/escnn/nn/modules/pooling/pointwise.py +++ b/escnn/nn/modules/pooling/pointwise.py @@ -4,7 +4,7 @@ from escnn.gspaces import GSpace from ..equivariant_module import EquivariantModule -from .gaussian_blur import GaussianBlurND +from .gaussian_blur import GaussianBlurND, kernel_size_from_radius from .utils import get_nd_tuple import torch @@ -137,12 +137,12 @@ def __init__( self.out_type = in_type self.blur = GaussianBlurND( - d=d, sigma=sigma, + kernel_size=kernel_size_from_radius(sigma * 3), stride=stride, - abs_padding=padding, + padding=padding, + d=d, channels=in_type.size, - _kernel_size_factor=3, ) def forward(self, input: GeometricTensor) -> GeometricTensor: @@ -234,12 +234,12 @@ def __init__( padding=padding, ) blur = GaussianBlurND( - d=d, sigma=sigma, + kernel_size=kernel_size_from_radius(sigma * 4), stride=stride if stride is not None else kernel_size, rel_padding=padding, channels=in_type.size, - _kernel_size_factor=4, + d=d, ) self.layers = torch.nn.Sequential( OrderedDict([('pool', pool), ('blur', blur)]), From ce99035d87b5456ffc6f40c3e7f1212cba6c5bd5 Mon Sep 17 00:00:00 2001 From: Kale Kundert Date: Fri, 3 Nov 2023 13:11:59 -0400 Subject: [PATCH 3/6] Add an option to account for edge effects when blurring Signed-off-by: Kale Kundert --- escnn/nn/modules/pooling/gaussian_blur.py | 110 +++++++++++++++------- test/nn/test_pooling.py | 110 ++++++++++++++++++++++ 2 files changed, 188 insertions(+), 32 deletions(-) diff --git a/escnn/nn/modules/pooling/gaussian_blur.py b/escnn/nn/modules/pooling/gaussian_blur.py index deb16ad7..87a89eea 100644 --- a/escnn/nn/modules/pooling/gaussian_blur.py +++ b/escnn/nn/modules/pooling/gaussian_blur.py @@ -4,6 +4,10 @@ import torch import torch.nn.functional as F +from torch.nn.modules import Module +from torch.nn.modules.lazy import LazyModuleMixin +from torch.nn.parameter import UninitializedBuffer, is_lazy + from typing import Optional, Union, Tuple _CONV = { @@ -11,18 +15,19 @@ 3: F.conv3d, } -class GaussianBlurND(torch.nn.Module): +class GaussianBlurND(LazyModuleMixin, Module): def __init__( self, *, + d: int, sigma: float, kernel_size: int, stride: Union[int, Tuple[int, ...]] = 1, padding: Optional[Union[int, Tuple[int, ...]]] = None, rel_padding: Optional[Union[int, Tuple[int, ...]]] = None, - d: int, - channels: int, + edge_correction: bool = False, + channels: Optional[int] = None, ): """ Apply a Gaussian blur to the input. @@ -66,38 +71,12 @@ def __init__( if padding is not None and rel_padding is not None: raise ValueError("can't specify `padding` and `rel_padding`") - self.kernel_size = kernel_size self.sigma = sigma + self.kernel_size = kernel_size self.stride = stride + self.edge_correction = edge_correction self.d = d - # Build the Gaussian smoothing filter - - grid = torch.meshgrid( - *[torch.arange(kernel_size)] * d, - indexing='ij', - ) - grid = torch.stack(grid, dim=-1) - - mean = (kernel_size - 1) / 2. - variance = sigma ** 2. - - # setting the dtype is needed, otherwise it becomes an integer tensor - r = torch.sum((grid - mean) ** 2., dim=-1, dtype=torch.get_default_dtype()) - - # Build the gaussian kernel - _filter = torch.exp(-r / (2 * variance)) - - # Normalize - _filter /= torch.sum(_filter) - - # The filter needs to be reshaped to be used in depthwise convolution - _filter = _filter\ - .view(1, 1, *[kernel_size]*d)\ - .repeat((channels, 1, *[1]*d)) - - self.register_buffer('filter', _filter) - if padding is not None: self.padding = padding else: @@ -107,10 +86,50 @@ def __init__( for p in get_nd_tuple(rel_padding or 0, d) ) + if channels is not None: + filter_ = make_gaussian_filter(sigma, kernel_size, d, channels) + else: + filter_ = UninitializedBuffer() + + self.register_buffer('filter', filter_) + + if edge_correction: + self.register_buffer('weights', UninitializedBuffer()) + def __repr__(self): - return f'{self.__class__.__name__}(sigma={self.sigma}, kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}, d={self.d}, channels={self.filter.shape[1]})' + return f'{self.__class__.__name__}(sigma={self.sigma}, kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}, edge_correction={self.edge_correction}, d={self.d})' + + def initialize_parameters(self, x): + if is_lazy(self.filter): + assert len(x.shape) == 2 + self.d + channels = x.shape[1] + filter_ = make_gaussian_filter( + self.sigma, + self.kernel_size, + self.d, + channels, + ) + + self.filter.materialize(shape=filter_.shape, dtype=filter_.dtype) + self.filter.copy_(filter_) + + if self.edge_correction: + ones = torch.ones(x.shape, dtype=x.dtype) + weights = self.blur(ones) + + self.weights.materialize(shape=weights.shape, dtype=weights.dtype) + self.weights.copy_(weights) def forward(self, x): + y = self.blur(x) + + if self.edge_correction: + assert y.shape == self.weights.shape + y /= self.weights + + return y + + def blur(self, x): return _CONV[self.d]( x, self.filter, @@ -119,5 +138,32 @@ def forward(self, x): groups=x.shape[1], ) +def make_gaussian_filter(sigma, kernel_size, d, channels): + grid = torch.meshgrid( + *[torch.arange(kernel_size)] * d, + indexing='ij', + ) + grid = torch.stack(grid, dim=-1) + + mean = (kernel_size - 1) / 2. + variance = sigma ** 2. + + # setting the dtype is needed, otherwise it becomes an integer tensor + r = torch.sum((grid - mean) ** 2., dim=-1, dtype=torch.get_default_dtype()) + + # Build the gaussian kernel + _filter = torch.exp(-r / (2 * variance)) + + # Normalize + _filter /= torch.sum(_filter) + + # The filter needs to be reshaped to be used in depthwise convolution + _filter = _filter\ + .view(1, 1, *[kernel_size]*d)\ + .repeat((channels, 1, *[1]*d)) + + return _filter + + def kernel_size_from_radius(radius): return 2 * int(round(radius)) + 1 diff --git a/test/nn/test_pooling.py b/test/nn/test_pooling.py index 02881d4e..6d3edde6 100644 --- a/test/nn/test_pooling.py +++ b/test/nn/test_pooling.py @@ -3,6 +3,7 @@ from escnn.nn import * from escnn.gspaces import * +from escnn.nn.modules.pooling.gaussian_blur import GaussianBlurND import torch @@ -909,5 +910,114 @@ def test_output_shape_2d(self): self.assertEqual(y.shape, case['out_shape']) self.assertEqual(y.shape, f.evaluate_output_shape(x.shape)) + + def test_gaussian_blur(self): + # Checking to make sure that layers and training examples are all + # completely independent from each other. + x = torch.Tensor([ + [[[1, 1, 1], + [1, 1, 1], + [1, 1, 1]], + + [[2, 2, 2], + [2, 2, 2], + [2, 2, 2]]], + + [[[3, 3, 3], + [3, 3, 3], + [3, 3, 3]], + + [[4, 4, 4], + [4, 4, 4], + [4, 4, 4]]], + ]) + + gb = GaussianBlurND(d=2, sigma=0.6, kernel_size=3) + + y = gb(x) + + y_expected = torch.Tensor([ + [[[0.6949, 0.8336, 0.6949], + [0.8336, 1.0000, 0.8336], + [0.6949, 0.8336, 0.6949]], + + [[1.3898, 1.6672, 1.3898], + [1.6672, 2.0000, 1.6672], + [1.3898, 1.6672, 1.3898]]], + + [[[2.0847, 2.5008, 2.0847], + [2.5008, 3.0000, 2.5008], + [2.0847, 2.5008, 2.0847]], + + [[2.7796, 3.3344, 2.7796], + [3.3344, 4.0000, 3.3344], + [2.7796, 3.3344, 2.7796]]], + ]) + + torch.testing.assert_close(y, y_expected, atol=1e-4, rtol=0) + + def test_gaussian_blur_edge_correction(self): + x = torch.Tensor([ + [[[1, 1, 1], + [1, 1, 1], + [1, 1, 1]], + + [[2, 2, 2], + [2, 2, 2], + [2, 2, 2]]], + + [[[3, 3, 3], + [3, 3, 3], + [3, 3, 3]], + + [[4, 4, 4], + [4, 4, 4], + [4, 4, 4]]], + ]) + + gb = GaussianBlurND(d=2, sigma=0.6, kernel_size=3, edge_correction=True) + + y = gb(x) + + y_expected = torch.Tensor([ + [[[1, 1, 1], + [1, 1, 1], + [1, 1, 1]], + + [[2, 2, 2], + [2, 2, 2], + [2, 2, 2]]], + + [[[3, 3, 3], + [3, 3, 3], + [3, 3, 3]], + + [[4, 4, 4], + [4, 4, 4], + [4, 4, 4]]], + ]) + + torch.testing.assert_close(y, y_expected) + + def test_gaussian_blur_repr(self): + # Non-default values for all constructor arguments: + blur = GaussianBlurND( + sigma=0.6, + kernel_size=5, + stride=2, + padding=1, + edge_correction=True, + d=3, + ) + blur_copy = eval(repr(blur), globals()) + + self.assertEqual(blur_copy.sigma, 0.6) + self.assertEqual(blur_copy.kernel_size, 5) + self.assertEqual(blur_copy.stride, 2) + self.assertEqual(blur_copy.padding, 1) + self.assertEqual(blur_copy.edge_correction, True) + self.assertEqual(blur_copy.d, 3) + + if __name__ == '__main__': unittest.main() From 948307d0fc278b20a709b6271c2fb1f69240b7e6 Mon Sep 17 00:00:00 2001 From: Kale Kundert Date: Fri, 17 Nov 2023 17:42:24 -0500 Subject: [PATCH 4/6] fix: use correct device Signed-off-by: Kale Kundert --- escnn/nn/modules/pooling/gaussian_blur.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/escnn/nn/modules/pooling/gaussian_blur.py b/escnn/nn/modules/pooling/gaussian_blur.py index 87a89eea..266b8e77 100644 --- a/escnn/nn/modules/pooling/gaussian_blur.py +++ b/escnn/nn/modules/pooling/gaussian_blur.py @@ -113,8 +113,8 @@ def initialize_parameters(self, x): self.filter.materialize(shape=filter_.shape, dtype=filter_.dtype) self.filter.copy_(filter_) - if self.edge_correction: - ones = torch.ones(x.shape, dtype=x.dtype) + if self.edge_correction and is_lazy(self.weights): + ones = torch.ones(x.shape, dtype=x.dtype, device=x.device) weights = self.blur(ones) self.weights.materialize(shape=weights.shape, dtype=weights.dtype) From 0d230db154a5dbbb0d5236b5d805dbdf72d0f5b1 Mon Sep 17 00:00:00 2001 From: Kale Kundert Date: Tue, 21 Nov 2023 17:53:31 -0500 Subject: [PATCH 5/6] Accommodate johnmarktaylor91/torchlens#18 Signed-off-by: Kale Kundert --- escnn/nn/modules/pooling/gaussian_blur.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/escnn/nn/modules/pooling/gaussian_blur.py b/escnn/nn/modules/pooling/gaussian_blur.py index 266b8e77..a48ec938 100644 --- a/escnn/nn/modules/pooling/gaussian_blur.py +++ b/escnn/nn/modules/pooling/gaussian_blur.py @@ -10,11 +10,6 @@ from typing import Optional, Union, Tuple -_CONV = { - 2: F.conv2d, - 3: F.conv3d, -} - class GaussianBlurND(LazyModuleMixin, Module): def __init__( @@ -76,6 +71,7 @@ def __init__( self.stride = stride self.edge_correction = edge_correction self.d = d + self.conv = getattr(F, f'conv{d}d') if padding is not None: self.padding = padding @@ -130,7 +126,7 @@ def forward(self, x): return y def blur(self, x): - return _CONV[self.d]( + return self.conv( x, self.filter, stride=self.stride, From 476725c889817df5743a1cd8a6bd13b73bd5299c Mon Sep 17 00:00:00 2001 From: Kale Kundert Date: Mon, 18 Dec 2023 12:36:22 -0500 Subject: [PATCH 6/6] Allow inputs with different batch sizes Signed-off-by: Kale Kundert --- escnn/nn/modules/pooling/gaussian_blur.py | 19 ++++++++++++------- test/nn/test_pooling.py | 16 ++++++++++++++++ 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/escnn/nn/modules/pooling/gaussian_blur.py b/escnn/nn/modules/pooling/gaussian_blur.py index a48ec938..37e2eec4 100644 --- a/escnn/nn/modules/pooling/gaussian_blur.py +++ b/escnn/nn/modules/pooling/gaussian_blur.py @@ -55,9 +55,14 @@ def __init__( error to specify *padding* and *rel_padding*. d (int): Dimensionality of the base space (2 for images, 3 for - volumes) - - channels (int): Channel dimension of the input. + volumes). + + channels (int): Channel dimension of the input. If specified, the + convolutional filter can be constructed immediately. + Otherwise, it will be constructed during the first forward + pass. Really, the only reason to specify this parameter is to + double check that the input has the expected number of + channels. """ super().__init__() @@ -87,10 +92,10 @@ def __init__( else: filter_ = UninitializedBuffer() - self.register_buffer('filter', filter_) + self.register_buffer('filter', filter_, persistent=False) if edge_correction: - self.register_buffer('weights', UninitializedBuffer()) + self.register_buffer('weights', UninitializedBuffer(), persistent=False) def __repr__(self): return f'{self.__class__.__name__}(sigma={self.sigma}, kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}, edge_correction={self.edge_correction}, d={self.d})' @@ -110,7 +115,8 @@ def initialize_parameters(self, x): self.filter.copy_(filter_) if self.edge_correction and is_lazy(self.weights): - ones = torch.ones(x.shape, dtype=x.dtype, device=x.device) + shape = 1, 1, *x.shape[2:] + ones = torch.ones(shape, dtype=x.dtype, device=x.device) weights = self.blur(ones) self.weights.materialize(shape=weights.shape, dtype=weights.dtype) @@ -120,7 +126,6 @@ def forward(self, x): y = self.blur(x) if self.edge_correction: - assert y.shape == self.weights.shape y /= self.weights return y diff --git a/test/nn/test_pooling.py b/test/nn/test_pooling.py index 6d3edde6..e4604a6b 100644 --- a/test/nn/test_pooling.py +++ b/test/nn/test_pooling.py @@ -956,6 +956,14 @@ def test_gaussian_blur(self): torch.testing.assert_close(y, y_expected, atol=1e-4, rtol=0) + # Make sure that the same module can process inputs with different + # batch dimensions: + x2 = x[0:2] + y2 = gb(x2) + y2_expected = y_expected[0:2] + + torch.testing.assert_close(y2, y2_expected, atol=1e-4, rtol=0) + def test_gaussian_blur_edge_correction(self): x = torch.Tensor([ [[[1, 1, 1], @@ -999,6 +1007,14 @@ def test_gaussian_blur_edge_correction(self): torch.testing.assert_close(y, y_expected) + # Make sure that the same module can process inputs with different + # batch dimensions: + x2 = x[0:2] + y2 = gb(x2) + y2_expected = y_expected[0:2] + + torch.testing.assert_close(y2, y2_expected, atol=1e-4, rtol=0) + def test_gaussian_blur_repr(self): # Non-default values for all constructor arguments: blur = GaussianBlurND(