composer.algorithms.blurpool.blurpool_layers#

composer.algorithms.blurpool.blurpool_layers

Functions

blur_2d

Applies a spatial low-pass filter.

blurmax_pool2d

Max-pooling with anti-aliasing.

Classes

BlurConv2d

This module is a drop-in replacement for torch.nn.Conv2d, but with an anti-aliasing filter.

BlurMaxPool2d

This module is a (nearly) drop-in replacement for torch.nn.MaxPool2d, but with an anti-aliasing filter.

BlurPool2d

This module just calls blur_2d() in forward using the provided arguments.

Attributes

  • Optional

class composer.algorithms.blurpool.blurpool_layers.BlurConv2d(in_channels, out_channels, kernel_size, stride=None, padding=0, dilation=1, groups=1, bias=True, blur_first=True)[source]#

Bases: torch.nn.modules.module.Module

This module is a drop-in replacement for torch.nn.Conv2d, but with an anti-aliasing filter.

The one new parameter is blur_first. When set to True, the anti-aliasing filter is applied before the underlying convolution and vice-versa when set to False. This mostly makes a difference when the stride is greater than one. In the former case, the only overhead is the cost of doing the anti-aliasing operation. In the latter case, the Conv2d is applied with a stride of one to the input, and then the anti-aliasing is applied with the provided stride to the result. Setting the stride of the convolution to 1 can greatly increase the computational cost. E.g., replacing a stride of (2, 2) with a stride of 1 increases the number of operations by a factor of (2/1) * (2/1) = 4. However, this approach most closely matches the behavior specified in the paper.

This module should only be used to replace strided convolutions.

See the associated paper for more details, experimental results, etc.

See also: blur_2d().

class composer.algorithms.blurpool.blurpool_layers.BlurMaxPool2d(kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False)[source]#

Bases: torch.nn.modules.module.Module

This module is a (nearly) drop-in replacement for torch.nn.MaxPool2d, but with an anti-aliasing filter.

The only API difference is that the parameter return_indices is not available, because it is ill-defined when using anti-aliasing.

See the associated paper for more details, experimental results, etc.

See blurmax_pool2d() for details.

class composer.algorithms.blurpool.blurpool_layers.BlurPool2d(stride=2, padding=1)[source]#

Bases: torch.nn.modules.module.Module

This module just calls blur_2d() in forward using the provided arguments.

composer.algorithms.blurpool.blurpool_layers.blur_2d(input, stride=1, filter=None)[source]#

Applies a spatial low-pass filter.

Parameters

  • input (Tensor) โ€“ A 4d tensor of shape NCHW

  • stride (int | tuple, optional) โ€“ Stride(s) along H and W axes. If a single value is passed, this value is used for both dimensions.

  • filter (Tensor, optional) โ€“ A 2d or 4d tensor to be cross-correlated with the input tensor at each spatial position, within each channel. If 4d, the structure is required to be (C, 1, kH, kW) where C is the number of channels in the input tensor and kH and kW are the spatial sizes of the filter.

By default, the filter used is:

[1 2 1]
[2 4 2] * 1/16
[1 2 1]
Returns

The blurred input

composer.algorithms.blurpool.blurpool_layers.blurmax_pool2d(input, kernel_size=None, stride=2, padding=0, dilation=1, ceil_mode=False, filter=None)[source]#

Max-pooling with anti-aliasing.

This is a nearly drop-in replacement for PyTorchโ€™s torch.nn.functional.max_pool2d(). The only API difference is that the parameter return_indices is not available, because it is ill-defined when using anti-aliasing.

See the associated paper for more details, experimental results, etc.

This function can be understood as decoupling the max from the pooling, and inserting a low-pass filtering step between the two. Concretely, this function computes the max within spatial neighborhoods of shape kernel_size, then applies an anti-aliasing filter to smooth the maxes, and only then pools according to stride.

See also: blur_2d().

Parameters

  • input (Tensor) โ€“ A 4d tensor of shape NCHW

  • kernel_size (int | tuple, optional) โ€“ Size(s) of the spatial neighborhoods over which to pool. This is mostly commonly 2x2. If only a scalar s is provided, the neighborhood is of size (s, s). Default: (2, 2).

  • stride (int | tuple, optional) โ€“ Stride(s) along H and W axes. If a single value is passed, this value is used for both dimensions. Default: 2.

  • padding (int | tuple, optional) โ€“ implicit zero-padding to use. For the default 3x3 low-pass filter, padding=1 (the default) returns output of the same size as the input. Default: 0.

  • dilation (int | tuple, optional) โ€“ Amount by which to โ€œstretchโ€ the pooling region for a given total size. See torch.nn.MaxPool2d for our favorite explanation of how this works. Default: 1.

  • ceil_mode (bool) โ€“ When True, will use ceil instead of floor to compute the output shape. Default: False.

  • filter (Tensor, optional) โ€“ A 2d or 4d tensor to be cross-correlated with the input tensor at each spatial position, within each channel. If 4d, the structure is required to be (C, 1, kH, kW) where C is the number of channels in the input tensor and kH and kW are the spatial sizes of the filter.

By default, the filter used is:

[1 2 1]
[2 4 2] * 1/16
[1 2 1]
Returns

The blurred and max-pooled input