You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
59 lines
2.1 KiB
59 lines
2.1 KiB
5 years ago
|
"""
|
||
|
|
BlurPool layer inspired by
|
||
|
|
- Kornia's Max_BlurPool2d
|
||
|
|
- Making Convolutional Networks Shift-Invariant Again :cite:`zhang2019shiftinvar`
|
||
|
|
|
||
|
|
FIXME merge this impl with those in `anti_aliasing.py`
|
||
|
|
|
||
|
|
Hacked together by Chris Ha and Ross Wightman
|
||
|
|
"""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
import numpy as np
|
||
|
|
from typing import Dict
|
||
|
|
from .padding import get_padding
|
||
|
|
|
||
|
|
|
||
|
|
class BlurPool2d(nn.Module):
|
||
|
|
r"""Creates a module that computes blurs and downsample a given feature map.
|
||
|
|
See :cite:`zhang2019shiftinvar` for more details.
|
||
|
|
Corresponds to the Downsample class, which does blurring and subsampling
|
||
|
|
|
||
|
|
Args:
|
||
|
|
channels = Number of input channels
|
||
|
|
filt_size (int): binomial filter size for blurring. currently supports 3 (default) and 5.
|
||
|
|
stride (int): downsampling filter stride
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
torch.Tensor: the transformed tensor.
|
||
|
|
"""
|
||
|
|
filt: Dict[str, torch.Tensor]
|
||
|
|
|
||
|
|
def __init__(self, channels, filt_size=3, stride=2) -> None:
|
||
|
|
super(BlurPool2d, self).__init__()
|
||
|
|
assert filt_size > 1
|
||
|
|
self.channels = channels
|
||
|
|
self.filt_size = filt_size
|
||
|
|
self.stride = stride
|
||
|
|
pad_size = [get_padding(filt_size, stride, dilation=1)] * 4
|
||
|
|
self.padding = nn.ReflectionPad2d(pad_size)
|
||
|
|
self._coeffs = torch.tensor((np.poly1d((0.5, 0.5)) ** (self.filt_size - 1)).coeffs) # for torchscript compat
|
||
|
|
self.filt = {} # lazy init by device for DataParallel compat
|
||
|
|
|
||
|
|
def _create_filter(self, like: torch.Tensor):
|
||
|
|
blur_filter = (self._coeffs[:, None] * self._coeffs[None, :]).to(dtype=like.dtype, device=like.device)
|
||
|
|
return blur_filter[None, None, :, :].repeat(self.channels, 1, 1, 1)
|
||
|
|
|
||
|
|
def _apply(self, fn):
|
||
|
|
# override nn.Module _apply, reset filter cache if used
|
||
|
|
self.filt = {}
|
||
|
|
super(BlurPool2d, self)._apply(fn)
|
||
|
|
|
||
|
|
def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
|
||
|
|
C = input_tensor.shape[1]
|
||
|
|
blur_filt = self.filt.get(str(input_tensor.device), self._create_filter(input_tensor))
|
||
|
|
return F.conv2d(
|
||
|
|
self.padding(input_tensor), blur_filt, stride=self.stride, groups=C)
|