From 3a287a6e764d54a984263f97923c60e8e101cc8b Mon Sep 17 00:00:00 2001 From: Chris Ha <15088501+VRandme@users.noreply.github.com> Date: Tue, 10 Mar 2020 23:27:24 +0900 Subject: [PATCH 1/6] Create blurpool.py Initial implementation of blur layer. currently tests as correct against Downsample of original github --- timm/models/layers/blurpool.py | 68 ++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 timm/models/layers/blurpool.py diff --git a/timm/models/layers/blurpool.py b/timm/models/layers/blurpool.py new file mode 100644 index 00000000..0ce4263e --- /dev/null +++ b/timm/models/layers/blurpool.py @@ -0,0 +1,68 @@ +'''independent attempt to implement + +MaxBlurPool2d in a more general fashion(separate maxpooling from BlurPool) +which was again inspired by +Making Convolutional Networks Shift-Invariant Again :cite:`zhang2019shiftinvar` + +''' + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +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 + blur_filter_size (int): filter size for blurring. currently supports either 3 or 5 (most common) + defaults to 3. + stride (int): downsampling filter stride + Shape: + Returns: + torch.Tensor: the transformed tensor. + Examples: + """ + + def __init__(self, channels=None, blur_filter_size=3, stride=2) -> None: + super(BlurPool2d, self).__init__() + assert blur_filter_size in [3, 5] + self.channels = channels + self.blur_filter_size = blur_filter_size + self.stride = stride + + if blur_filter_size == 3: + pad_size = [1] * 4 + blur_matrix = torch.Tensor([[1., 2., 1]]) / 4 # binomial kernel b2 + else: + pad_size = [2] * 4 + blur_matrix = torch.Tensor([[1., 4., 6., 4., 1.]]) / 16 # binomial filter kernel 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 + if not torch.is_tensor(input_tensor): + raise TypeError("Input input type is not a torch.Tensor. Got {}" + .format(type(input_tensor))) + if not len(input_tensor.shape) == 4: + 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 From ce3d82b58b2eb3e526de5df608cfc7d181b8c916 Mon Sep 17 00:00:00 2001 From: Chris Ha <15088501+VRandme@users.noreply.github.com> Date: Wed, 11 Mar 2020 22:19:10 +0900 Subject: [PATCH 2/6] Update blurpool.py clean up code for PR --- timm/models/layers/blurpool.py | 32 +++++++++----------------------- 1 file changed, 9 insertions(+), 23 deletions(-) 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 From acd1b6cccd9a333bd790b2d00d646a966e710151 Mon Sep 17 00:00:00 2001 From: Chris Ha <15088501+VRandme@users.noreply.github.com> Date: Sun, 22 Mar 2020 22:42:55 +0900 Subject: [PATCH 3/6] Implement Functional Blur on resnet.py 1. add ResNet argument blur='' 2. implement blur for maxpool and strided convs in downsampling blocks --- timm/models/layers/__init__.py | 1 + timm/models/layers/blurpool.py | 2 +- timm/models/resnet.py | 79 +++++++++++++++++++++++++++++----- 3 files changed, 70 insertions(+), 12 deletions(-) diff --git a/timm/models/layers/__init__.py b/timm/models/layers/__init__.py index f012c3cf..33450483 100644 --- a/timm/models/layers/__init__.py +++ b/timm/models/layers/__init__.py @@ -15,3 +15,4 @@ from .adaptive_avgmax_pool import \ from .drop import DropBlock2d, DropPath, drop_block_2d, drop_path from .test_time_pool import TestTimePoolHead, apply_test_time_pool from .split_batchnorm import SplitBatchNorm2d, convert_splitbn_model +from .blurpool import BlurPool2d diff --git a/timm/models/layers/blurpool.py b/timm/models/layers/blurpool.py index 96937114..0b37a90c 100644 --- a/timm/models/layers/blurpool.py +++ b/timm/models/layers/blurpool.py @@ -17,7 +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): binomial filter size for blurring. currently supports 3(default) and 5. + blur_filter_size (int): binomial filter size for blurring. currently supports 3 (default) and 5. stride (int): downsampling filter stride Shape: Returns: diff --git a/timm/models/resnet.py b/timm/models/resnet.py index 0013cbe0..057eca6c 100644 --- a/timm/models/resnet.py +++ b/timm/models/resnet.py @@ -12,7 +12,7 @@ import torch.nn.functional as F from .registry import register_model from .helpers import load_pretrained -from .layers import SelectAdaptivePool2d, DropBlock2d, DropPath, AvgPool2dSame, create_attn +from .layers import SelectAdaptivePool2d, DropBlock2d, DropPath, AvgPool2dSame, create_attn, BlurPool2d from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD @@ -104,6 +104,8 @@ default_cfgs = { interpolation='bicubic'), 'ecaresnet18': _cfg(), 'ecaresnet50': _cfg(), + 'resnetblur18': _cfg(), + 'resnetblur50': _cfg() } @@ -117,7 +119,7 @@ class BasicBlock(nn.Module): def __init__(self, inplanes, planes, stride=1, downsample=None, cardinality=1, base_width=64, reduce_first=1, dilation=1, first_dilation=None, act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, - attn_layer=None, drop_block=None, drop_path=None): + attn_layer=None, drop_block=None, drop_path=None, blur=False): super(BasicBlock, self).__init__() assert cardinality == 1, 'BasicBlock only supports cardinality of 1' @@ -125,10 +127,19 @@ class BasicBlock(nn.Module): first_planes = planes // reduce_first outplanes = planes * self.expansion first_dilation = first_dilation or dilation + self.blur = blur - self.conv1 = nn.Conv2d( + if blur and stride==2: + self.conv1 = nn.Conv2d( + inplanes, first_planes, kernel_size=3, stride=1, padding=first_dilation, + dilation=first_dilation, bias=False) + self.blurpool=BlurPool2d(channels=first_planes) + else: + self.conv1 = nn.Conv2d( inplanes, first_planes, kernel_size=3, stride=stride, padding=first_dilation, dilation=first_dilation, bias=False) + self.blurpool = None + self.bn1 = norm_layer(first_planes) self.act1 = act_layer(inplace=True) self.conv2 = nn.Conv2d( @@ -154,7 +165,11 @@ class BasicBlock(nn.Module): x = self.bn1(x) if self.drop_block is not None: x = self.drop_block(x) - x = self.act1(x) + if self.blurpool is not None: + x = self.act1(x) + x = self.blurpool(x) + else: + x = self.act1(x) x = self.conv2(x) x = self.bn2(x) @@ -181,20 +196,30 @@ class Bottleneck(nn.Module): def __init__(self, inplanes, planes, stride=1, downsample=None, cardinality=1, base_width=64, reduce_first=1, dilation=1, first_dilation=None, act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, - attn_layer=None, drop_block=None, drop_path=None): + attn_layer=None, drop_block=None, drop_path=None, blur=False): super(Bottleneck, self).__init__() width = int(math.floor(planes * (base_width / 64)) * cardinality) first_planes = width // reduce_first outplanes = planes * self.expansion first_dilation = first_dilation or dilation + self.blur = blur self.conv1 = nn.Conv2d(inplanes, first_planes, kernel_size=1, bias=False) self.bn1 = norm_layer(first_planes) self.act1 = act_layer(inplace=True) - self.conv2 = nn.Conv2d( - first_planes, width, kernel_size=3, stride=stride, - padding=first_dilation, dilation=first_dilation, groups=cardinality, bias=False) + + if blur and stride==2: + self.conv2 = nn.Conv2d( + first_planes, width, kernel_size=3, stride=1, + padding=first_dilation, dilation=first_dilation, groups=cardinality, bias=False) + self.blurpool = BlurPool2d(channels=width) + else: + self.conv2 = nn.Conv2d( + first_planes, width, kernel_size=3, stride=stride, + padding=first_dilation, dilation=first_dilation, groups=cardinality, bias=False) + self.blurpool = None + self.bn2 = norm_layer(width) self.act2 = act_layer(inplace=True) self.conv3 = nn.Conv2d(width, outplanes, kernel_size=1, bias=False) @@ -345,12 +370,19 @@ class ResNet(nn.Module): Dropout probability before classifier, for training global_pool : str, default 'avg' Global pooling type. One of 'avg', 'max', 'avgmax', 'catavgmax' + blur : str, default '' + Location of Blurring: + * '', default - Not applied + * 'max' - only stem layer MaxPool will be blurred + * 'strided' - only strided convolutions in the downsampling blocks (assembled-cnn style) + * 'max_strided' - on both stem MaxPool and strided convolutions (zhang2019shiftinvar style for ResNets) + """ def __init__(self, block, layers, num_classes=1000, in_chans=3, cardinality=1, base_width=64, stem_width=64, stem_type='', block_reduce_first=1, down_kernel_size=1, avg_down=False, output_stride=32, act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, drop_rate=0.0, drop_path_rate=0., - drop_block_rate=0., global_pool='avg', zero_init_last_bn=True, block_args=None): + drop_block_rate=0., global_pool='avg', blur='', zero_init_last_bn=True, block_args=None): block_args = block_args or dict() self.num_classes = num_classes deep_stem = 'deep' in stem_type @@ -359,6 +391,7 @@ class ResNet(nn.Module): self.base_width = base_width self.drop_rate = drop_rate self.expansion = block.expansion + self.blur = 'strided' in blur super(ResNet, self).__init__() # Stem @@ -379,7 +412,13 @@ class ResNet(nn.Module): self.conv1 = nn.Conv2d(in_chans, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = norm_layer(self.inplanes) self.act1 = act_layer(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + # Stem Blur + if 'max' in blur : + self.maxpool = nn.Sequential(*[ + nn.MaxPool2d(kernel_size=3, stride=1, padding=1), + BlurPool2d(channels=self.inplanes)]) + else : + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) # Feature Blocks dp = DropPath(drop_path_rate) if drop_path_rate else None @@ -432,7 +471,7 @@ class ResNet(nn.Module): block_kwargs = dict( cardinality=self.cardinality, base_width=self.base_width, reduce_first=reduce_first, dilation=dilation, **kwargs) - layers = [block(self.inplanes, planes, stride, downsample, first_dilation=first_dilation, **block_kwargs)] + layers = [block(self.inplanes, planes, stride, downsample, first_dilation=first_dilation, blur=self.blur, **block_kwargs)] self.inplanes = planes * block.expansion layers += [block(self.inplanes, planes, **block_kwargs) for _ in range(1, blocks)] @@ -1022,3 +1061,21 @@ def ecaresnet50(pretrained=False, num_classes=1000, in_chans=3, **kwargs): if pretrained: load_pretrained(model, default_cfg, num_classes, in_chans) return model + +@register_model +def resnetblur18(pretrained=False, num_classes=1000, in_chans=3, **kwargs): + """Constructs a ResNet-18 model. With original style blur + """ + default_cfg = default_cfgs['resnetblur18'] + model = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, in_chans=in_chans, blur='max_strided',**kwargs) + model.default_cfg = default_cfg + return model + +@register_model +def resnetblur50(pretrained=False, num_classes=1000, in_chans=3, **kwargs): + """Constructs a ResNet-50 model. With assembled-cnn style blur + """ + default_cfg = default_cfgs['resnetblur18'] + model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, in_chans=in_chans, blur='strided', **kwargs) + model.default_cfg = default_cfg + return model \ No newline at end of file From 6cdeca24a389cce7693e043b89db04e230b3a405 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Wed, 1 Apr 2020 12:01:59 -0700 Subject: [PATCH 4/6] Some cleanup and fixes for initial BlurPool impl. Still some testing and tweaks to go... --- timm/models/layers/blurpool.py | 43 +++++++++++++++------------- timm/models/resnet.py | 52 +++++++++++++--------------------- 2 files changed, 43 insertions(+), 52 deletions(-) diff --git a/timm/models/layers/blurpool.py b/timm/models/layers/blurpool.py index 0b37a90c..a12274d8 100644 --- a/timm/models/layers/blurpool.py +++ b/timm/models/layers/blurpool.py @@ -1,14 +1,17 @@ -''' +""" BlurPool layer inspired by -Kornia's Max_BlurPool2d -and -Making Convolutional Networks Shift-Invariant Again :cite:`zhang2019shiftinvar` + - Kornia's Max_BlurPool2d + - Making Convolutional Networks Shift-Invariant Again :cite:`zhang2019shiftinvar` + +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 .padding import get_padding class BlurPool2d(nn.Module): @@ -25,30 +28,30 @@ class BlurPool2d(nn.Module): Examples: """ - def __init__(self, channels=None, blur_filter_size=3, stride=2) -> None: + def __init__(self, channels, blur_filter_size=3, stride=2) -> None: super(BlurPool2d, self).__init__() - assert blur_filter_size in [3, 5] + assert blur_filter_size > 1 self.channels = channels self.blur_filter_size = blur_filter_size self.stride = stride - if blur_filter_size == 3: - pad_size = [1] * 4 - 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 b4 - + pad_size = [get_padding(blur_filter_size, stride, dilation=1)] * 4 self.padding = nn.ReflectionPad2d(pad_size) - blur_filter = blur_matrix * blur_matrix.T + + blur_matrix = (np.poly1d((0.5, 0.5)) ** (blur_filter_size - 1)).coeffs + blur_filter = torch.Tensor(blur_matrix[:, None] * blur_matrix[None, :]) + # FIXME figure a clean hack to prevent the filter from getting saved in weights, but still + # plays nice with recursive module apply for fn like .cuda(), .type(), etc -RW 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 if not torch.is_tensor(input_tensor): - raise TypeError("Input input type is not a torch.Tensor. Got {}" - .format(type(input_tensor))) + raise TypeError("Input input type is not a torch.Tensor. Got {}".format(type(input_tensor))) if not len(input_tensor.shape) == 4: - raise ValueError("Invalid input shape, we expect BxCxHxW. Got: {}" - .format(input_tensor.shape)) + 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]) \ No newline at end of file + return F.conv2d( + self.padding(input_tensor), + self.blur_filter.type(input_tensor.dtype), + stride=self.stride, + groups=input_tensor.shape[1]) diff --git a/timm/models/resnet.py b/timm/models/resnet.py index 057eca6c..fdb8097c 100644 --- a/timm/models/resnet.py +++ b/timm/models/resnet.py @@ -127,21 +127,14 @@ class BasicBlock(nn.Module): first_planes = planes // reduce_first outplanes = planes * self.expansion first_dilation = first_dilation or dilation - self.blur = blur - if blur and stride==2: - self.conv1 = nn.Conv2d( - inplanes, first_planes, kernel_size=3, stride=1, padding=first_dilation, - dilation=first_dilation, bias=False) - self.blurpool=BlurPool2d(channels=first_planes) - else: - self.conv1 = nn.Conv2d( - inplanes, first_planes, kernel_size=3, stride=stride, padding=first_dilation, + self.conv1 = nn.Conv2d( + inplanes, first_planes, kernel_size=3, stride=1 if blur else stride, padding=first_dilation, dilation=first_dilation, bias=False) - self.blurpool = None - self.bn1 = norm_layer(first_planes) self.act1 = act_layer(inplace=True) + self.blurpool = BlurPool2d(channels=first_planes) if stride == 2 and blur else None + self.conv2 = nn.Conv2d( first_planes, outplanes, kernel_size=3, padding=dilation, dilation=dilation, bias=False) self.bn2 = norm_layer(outplanes) @@ -165,11 +158,9 @@ class BasicBlock(nn.Module): x = self.bn1(x) if self.drop_block is not None: x = self.drop_block(x) + x = self.act1(x) if self.blurpool is not None: - x = self.act1(x) x = self.blurpool(x) - else: - x = self.act1(x) x = self.conv2(x) x = self.bn2(x) @@ -209,19 +200,13 @@ class Bottleneck(nn.Module): self.bn1 = norm_layer(first_planes) self.act1 = act_layer(inplace=True) - if blur and stride==2: - self.conv2 = nn.Conv2d( - first_planes, width, kernel_size=3, stride=1, - padding=first_dilation, dilation=first_dilation, groups=cardinality, bias=False) - self.blurpool = BlurPool2d(channels=width) - else: - self.conv2 = nn.Conv2d( - first_planes, width, kernel_size=3, stride=stride, - padding=first_dilation, dilation=first_dilation, groups=cardinality, bias=False) - self.blurpool = None - + self.conv2 = nn.Conv2d( + first_planes, width, kernel_size=3, stride=1 if blur else stride, + padding=first_dilation, dilation=first_dilation, groups=cardinality, bias=False) self.bn2 = norm_layer(width) self.act2 = act_layer(inplace=True) + self.blurpool = BlurPool2d(channels=width) if stride == 2 and blur else None + self.conv3 = nn.Conv2d(width, outplanes, kernel_size=1, bias=False) self.bn3 = norm_layer(outplanes) @@ -251,6 +236,8 @@ class Bottleneck(nn.Module): if self.drop_block is not None: x = self.drop_block(x) x = self.act2(x) + if self.blurpool is not None: + x = self.blurpool(x) x = self.conv3(x) x = self.bn3(x) @@ -412,11 +399,12 @@ class ResNet(nn.Module): self.conv1 = nn.Conv2d(in_chans, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = norm_layer(self.inplanes) self.act1 = act_layer(inplace=True) - # Stem Blur + # Stem Pooling if 'max' in blur : self.maxpool = nn.Sequential(*[ - nn.MaxPool2d(kernel_size=3, stride=1, padding=1), - BlurPool2d(channels=self.inplanes)]) + nn.MaxPool2d(kernel_size=3, stride=1, padding=1), + BlurPool2d(channels=self.inplanes, stride=2) + ]) else : self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) @@ -470,8 +458,8 @@ class ResNet(nn.Module): block_kwargs = dict( cardinality=self.cardinality, base_width=self.base_width, reduce_first=reduce_first, - dilation=dilation, **kwargs) - layers = [block(self.inplanes, planes, stride, downsample, first_dilation=first_dilation, blur=self.blur, **block_kwargs)] + dilation=dilation, blur=self.blur, **kwargs) + layers = [block(self.inplanes, planes, stride, downsample, first_dilation=first_dilation, **block_kwargs)] self.inplanes = planes * block.expansion layers += [block(self.inplanes, planes, **block_kwargs) for _ in range(1, blocks)] @@ -1075,7 +1063,7 @@ def resnetblur18(pretrained=False, num_classes=1000, in_chans=3, **kwargs): def resnetblur50(pretrained=False, num_classes=1000, in_chans=3, **kwargs): """Constructs a ResNet-50 model. With assembled-cnn style blur """ - default_cfg = default_cfgs['resnetblur18'] - model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, in_chans=in_chans, blur='strided', **kwargs) + default_cfg = default_cfgs['resnetblur50'] + model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, in_chans=in_chans, blur='max_strided', **kwargs) model.default_cfg = default_cfg return model \ No newline at end of file From f17b42bc33eca08308192f4079f8f609dada81db Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Wed, 1 Apr 2020 13:15:06 -0700 Subject: [PATCH 5/6] Blur filter no longer a buffer --- timm/models/layers/blurpool.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/timm/models/layers/blurpool.py b/timm/models/layers/blurpool.py index a12274d8..6df2b748 100644 --- a/timm/models/layers/blurpool.py +++ b/timm/models/layers/blurpool.py @@ -40,18 +40,16 @@ class BlurPool2d(nn.Module): blur_matrix = (np.poly1d((0.5, 0.5)) ** (blur_filter_size - 1)).coeffs blur_filter = torch.Tensor(blur_matrix[:, None] * blur_matrix[None, :]) - # FIXME figure a clean hack to prevent the filter from getting saved in weights, but still - # plays nice with recursive module apply for fn like .cuda(), .type(), etc -RW - self.register_buffer('blur_filter', blur_filter[None, None, :, :].repeat((self.channels, 1, 1, 1))) + self.blur_filter = blur_filter[None, None, :, :] + + def _apply(self, fn): + # override nn.Module _apply to prevent need for blur_filter to be registered as a buffer, + # this keeps it out of state dict, but allows .cuda(), .type(), etc to work as expected + super(BlurPool2d, self)._apply(fn) + self.blur_filter = fn(self.blur_filter) 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))) - if not len(input_tensor.shape) == 4: - 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.type(input_tensor.dtype), - stride=self.stride, - groups=input_tensor.shape[1]) + self.blur_filter.type(input_tensor.dtype).expand(C, -1, -1, -1), + stride=self.stride, groups=input_tensor.shape[1]) From 1a9ab07307a5dccaec08566bc599b3c36c89dae8 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Wed, 1 Apr 2020 13:19:08 -0700 Subject: [PATCH 6/6] One too many changes at a time, fix missing C --- timm/models/layers/blurpool.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/timm/models/layers/blurpool.py b/timm/models/layers/blurpool.py index 6df2b748..57af3e0e 100644 --- a/timm/models/layers/blurpool.py +++ b/timm/models/layers/blurpool.py @@ -49,7 +49,7 @@ class BlurPool2d(nn.Module): self.blur_filter = fn(self.blur_filter) def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: # type: ignore + C = input_tensor.shape[1] return F.conv2d( self.padding(input_tensor), - self.blur_filter.type(input_tensor.dtype).expand(C, -1, -1, -1), - stride=self.stride, groups=input_tensor.shape[1]) + self.blur_filter.type(input_tensor.dtype).expand(C, -1, -1, -1), stride=self.stride, groups=C)