Selective Kernel Convolution Module As described in Selective Kernel Networks (https://arxiv.org/abs/1903.06586) with some modifications. Largest change is the input split, which divides the input channels across each convolution path, this can be viewed as a grouping of s
(self, in_channels, out_channels=None, kernel_size=None, stride=1, dilation=1, groups=1,
rd_ratio=1./16, rd_channels=None, rd_divisor=8, keep_3x3=True, split_input=True,
act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, aa_layer=None, drop_layer=None)
| 49 | class SelectiveKernel(nn.Module): |
| 50 | |
| 51 | def __init__(self, in_channels, out_channels=None, kernel_size=None, stride=1, dilation=1, groups=1, |
| 52 | rd_ratio=1./16, rd_channels=None, rd_divisor=8, keep_3x3=True, split_input=True, |
| 53 | act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, aa_layer=None, drop_layer=None): |
| 54 | """ Selective Kernel Convolution Module |
| 55 | |
| 56 | As described in Selective Kernel Networks (https://arxiv.org/abs/1903.06586) with some modifications. |
| 57 | |
| 58 | Largest change is the input split, which divides the input channels across each convolution path, this can |
| 59 | be viewed as a grouping of sorts, but the output channel counts expand to the module level value. This keeps |
| 60 | the parameter count from ballooning when the convolutions themselves don't have groups, but still provides |
| 61 | a noteworthy increase in performance over similar param count models without this attention layer. -Ross W |
| 62 | |
| 63 | Args: |
| 64 | in_channels (int): module input (feature) channel count |
| 65 | out_channels (int): module output (feature) channel count |
| 66 | kernel_size (int, list): kernel size for each convolution branch |
| 67 | stride (int): stride for convolutions |
| 68 | dilation (int): dilation for module as a whole, impacts dilation of each branch |
| 69 | groups (int): number of groups for each branch |
| 70 | rd_ratio (int, float): reduction factor for attention features |
| 71 | keep_3x3 (bool): keep all branch convolution kernels as 3x3, changing larger kernels for dilations |
| 72 | split_input (bool): split input channels evenly across each convolution branch, keeps param count lower, |
| 73 | can be viewed as grouping by path, output expands to module out_channels count |
| 74 | act_layer (nn.Module): activation layer to use |
| 75 | norm_layer (nn.Module): batchnorm/norm layer to use |
| 76 | aa_layer (nn.Module): anti-aliasing module |
| 77 | drop_layer (nn.Module): spatial drop module in convs (drop block, etc) |
| 78 | """ |
| 79 | super(SelectiveKernel, self).__init__() |
| 80 | out_channels = out_channels or in_channels |
| 81 | kernel_size = kernel_size or [3, 5] # default to one 3x3 and one 5x5 branch. 5x5 -> 3x3 + dilation |
| 82 | _kernel_valid(kernel_size) |
| 83 | if not isinstance(kernel_size, list): |
| 84 | kernel_size = [kernel_size] * 2 |
| 85 | if keep_3x3: |
| 86 | dilation = [dilation * (k - 1) // 2 for k in kernel_size] |
| 87 | kernel_size = [3] * len(kernel_size) |
| 88 | else: |
| 89 | dilation = [dilation] * len(kernel_size) |
| 90 | self.num_paths = len(kernel_size) |
| 91 | self.in_channels = in_channels |
| 92 | self.out_channels = out_channels |
| 93 | self.split_input = split_input |
| 94 | if self.split_input: |
| 95 | assert in_channels % self.num_paths == 0 |
| 96 | in_channels = in_channels // self.num_paths |
| 97 | groups = min(out_channels, groups) |
| 98 | |
| 99 | conv_kwargs = dict( |
| 100 | stride=stride, groups=groups, act_layer=act_layer, norm_layer=norm_layer, |
| 101 | aa_layer=aa_layer, drop_layer=drop_layer) |
| 102 | self.paths = nn.ModuleList([ |
| 103 | ConvNormActAa(in_channels, out_channels, kernel_size=k, dilation=d, **conv_kwargs) |
| 104 | for k, d in zip(kernel_size, dilation)]) |
| 105 | |
| 106 | attn_channels = rd_channels or make_divisible(out_channels * rd_ratio, divisor=rd_divisor) |
| 107 | self.attn = SelectiveKernelAttn(out_channels, self.num_paths, attn_channels) |
| 108 |
no test coverage detected