| """ PyTorch Involution Layer | |
| Official impl: https://github.com/d-li14/involution/blob/main/cls/mmcls/models/utils/involution_naive.py | |
| Paper: `Involution: Inverting the Inherence of Convolution for Visual Recognition` - https://arxiv.org/abs/2103.06255 | |
| """ | |
| import torch.nn as nn | |
| from .conv_bn_act import ConvBnAct | |
| from .create_conv2d import create_conv2d | |
| class Involution(nn.Module): | |
| def __init__( | |
| self, | |
| channels, | |
| kernel_size=3, | |
| stride=1, | |
| group_size=16, | |
| rd_ratio=4, | |
| norm_layer=nn.BatchNorm2d, | |
| act_layer=nn.ReLU, | |
| ): | |
| super(Involution, self).__init__() | |
| self.kernel_size = kernel_size | |
| self.stride = stride | |
| self.channels = channels | |
| self.group_size = group_size | |
| self.groups = self.channels // self.group_size | |
| self.conv1 = ConvBnAct( | |
| in_channels=channels, | |
| out_channels=channels // rd_ratio, | |
| kernel_size=1, | |
| norm_layer=norm_layer, | |
| act_layer=act_layer) | |
| self.conv2 = self.conv = create_conv2d( | |
| in_channels=channels // rd_ratio, | |
| out_channels=kernel_size**2 * self.groups, | |
| kernel_size=1, | |
| stride=1) | |
| self.avgpool = nn.AvgPool2d(stride, stride) if stride == 2 else nn.Identity() | |
| self.unfold = nn.Unfold(kernel_size, 1, (kernel_size-1)//2, stride) | |
| def forward(self, x): | |
| weight = self.conv2(self.conv1(self.avgpool(x))) | |
| B, C, H, W = weight.shape | |
| KK = int(self.kernel_size ** 2) | |
| weight = weight.view(B, self.groups, KK, H, W).unsqueeze(2) | |
| out = self.unfold(x).view(B, self.groups, self.group_size, KK, H, W) | |
| out = (weight * out).sum(dim=3).view(B, self.channels, H, W) | |
| return out | |