diff --git a/timm/models/layers/blurpool.py b/timm/models/layers/blurpool.py index 0ce4263e..96937114 100644 --- a/timm/models/layers/blurpool.py +++ b/timm/models/layers/blurpool.py @@ -1,7 +1,7 @@ -'''independent attempt to implement - -MaxBlurPool2d in a more general fashion(separate maxpooling from BlurPool) -which was again inspired by +''' +BlurPool layer inspired by +Kornia's Max_BlurPool2d +and Making Convolutional Networks Shift-Invariant Again :cite:`zhang2019shiftinvar` ''' @@ -17,8 +17,7 @@ class BlurPool2d(nn.Module): Corresponds to the Downsample class, which does blurring and subsampling Args: channels = Number of input channels - blur_filter_size (int): filter size for blurring. currently supports either 3 or 5 (most common) - defaults to 3. + blur_filter_size (int): binomial filter size for blurring. currently supports 3(default) and 5. stride (int): downsampling filter stride Shape: Returns: @@ -35,16 +34,16 @@ class BlurPool2d(nn.Module): if blur_filter_size == 3: pad_size = [1] * 4 - blur_matrix = torch.Tensor([[1., 2., 1]]) / 4 # binomial kernel b2 + blur_matrix = torch.Tensor([[1., 2., 1]]) / 4 # binomial filter b2 else: pad_size = [2] * 4 - blur_matrix = torch.Tensor([[1., 4., 6., 4., 1.]]) / 16 # binomial filter kernel b4 + blur_matrix = torch.Tensor([[1., 4., 6., 4., 1.]]) / 16 # binomial filter b4 self.padding = nn.ReflectionPad2d(pad_size) blur_filter = blur_matrix * blur_matrix.T self.register_buffer('blur_filter', blur_filter[None, None, :, :].repeat((self.channels, 1, 1, 1))) - def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: # type: ignore + def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: # type: ignore if not torch.is_tensor(input_tensor): raise TypeError("Input input type is not a torch.Tensor. Got {}" .format(type(input_tensor))) @@ -52,17 +51,4 @@ class BlurPool2d(nn.Module): raise ValueError("Invalid input shape, we expect BxCxHxW. Got: {}" .format(input_tensor.shape)) # apply blur_filter on input - return F.conv2d(self.padding(input_tensor), self.blur_filter, stride=self.stride, groups=input_tensor.shape[1]) - - -###################### -# functional interface -###################### - - -'''def blur_pool2d() -> torch.Tensor: - r"""Creates a module that computes pools and blurs and downsample a given - feature map. - See :class:`~kornia.contrib.MaxBlurPool2d` for details. - """ - return BlurPool2d(kernel_size, ceil_mode)(input)''' \ No newline at end of file + return F.conv2d(self.padding(input_tensor), self.blur_filter, stride=self.stride, groups=input_tensor.shape[1]) \ No newline at end of file