mirror of https://github.com/JDAI-CV/fast-reid.git
80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
# encoding: utf-8
|
|
"""
|
|
@author: xingyu liao
|
|
@contact: liaoxingyu5@jd.com
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.nn import Conv2d, Module, ReLU
|
|
from torch.nn.modules.utils import _pair
|
|
|
|
|
|
class SplAtConv2d(Module):
|
|
"""Split-Attention Conv2d
|
|
"""
|
|
|
|
def __init__(self, in_channels, channels, kernel_size, stride=(1, 1), padding=(0, 0),
|
|
dilation=(1, 1), groups=1, bias=True,
|
|
radix=2, reduction_factor=4,
|
|
rectify=False, rectify_avg=False, norm_layer=None,
|
|
dropblock_prob=0.0, **kwargs):
|
|
super(SplAtConv2d, self).__init__()
|
|
padding = _pair(padding)
|
|
self.rectify = rectify and (padding[0] > 0 or padding[1] > 0)
|
|
self.rectify_avg = rectify_avg
|
|
inter_channels = max(in_channels * radix // reduction_factor, 32)
|
|
self.radix = radix
|
|
self.cardinality = groups
|
|
self.channels = channels
|
|
self.dropblock_prob = dropblock_prob
|
|
if self.rectify:
|
|
from rfconv import RFConv2d
|
|
self.conv = RFConv2d(in_channels, channels * radix, kernel_size, stride, padding, dilation,
|
|
groups=groups * radix, bias=bias, average_mode=rectify_avg, **kwargs)
|
|
else:
|
|
self.conv = Conv2d(in_channels, channels * radix, kernel_size, stride, padding, dilation,
|
|
groups=groups * radix, bias=bias, **kwargs)
|
|
self.use_bn = norm_layer is not None
|
|
self.bn0 = norm_layer(channels * radix)
|
|
self.relu = ReLU(inplace=True)
|
|
self.fc1 = Conv2d(channels, inter_channels, 1, groups=self.cardinality)
|
|
self.bn1 = norm_layer(inter_channels)
|
|
self.fc2 = Conv2d(inter_channels, channels * radix, 1, groups=self.cardinality)
|
|
if dropblock_prob > 0.0:
|
|
self.dropblock = DropBlock2D(dropblock_prob, 3)
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
if self.use_bn:
|
|
x = self.bn0(x)
|
|
if self.dropblock_prob > 0.0:
|
|
x = self.dropblock(x)
|
|
x = self.relu(x)
|
|
|
|
batch, channel = x.shape[:2]
|
|
if self.radix > 1:
|
|
splited = torch.split(x, channel // self.radix, dim=1)
|
|
gap = sum(splited)
|
|
else:
|
|
gap = x
|
|
gap = F.adaptive_avg_pool2d(gap, 1)
|
|
gap = self.fc1(gap)
|
|
|
|
if self.use_bn:
|
|
gap = self.bn1(gap)
|
|
gap = self.relu(gap)
|
|
|
|
atten = self.fc2(gap).view((batch, self.radix, self.channels))
|
|
if self.radix > 1:
|
|
atten = F.softmax(atten, dim=1).view(batch, -1, 1, 1)
|
|
else:
|
|
atten = F.sigmoid(atten, dim=1).view(batch, -1, 1, 1)
|
|
|
|
if self.radix > 1:
|
|
atten = torch.split(atten, channel // self.radix, dim=1)
|
|
out = sum([att * split for (att, split) in zip(atten, splited)])
|
|
else:
|
|
out = atten * x
|
|
return out.contiguous()
|