def __init__(self, in_channels, bn_norm, norm_opt, reduc_ratio=2):
super(Non_local, self).__init__()
self.in_channels = in_channels
self.inter_channels = in_channels // reduc_ratio
self.g = meta_conv2d(in_channels=self.in_channels,
out_channels=self.inter_channels,
kernel_size=1,
stride=1,
padding=0)
self.W = nn.Sequential(
meta_conv2d(in_channels=self.inter_channels,
out_channels=self.in_channels,
kernel_size=1,
stride=1,
padding=0),
meta_norm(bn_norm, self.in_channels, norm_opt=norm_opt),
)
nn.init.constant_(self.W[1].weight, 0.0)
nn.init.constant_(self.W[1].bias, 0.0)
self.theta = meta_conv2d(in_channels=self.in_channels,
out_channels=self.inter_channels,
kernel_size=1,
stride=1,
padding=0)
self.phi = meta_conv2d(in_channels=self.in_channels,
out_channels=self.inter_channels,
kernel_size=1,
stride=1,
padding=0)
def __init__(self, inplanes, planes, bn_norm, norm_opt, num_splits, with_ibn=False, with_se=False,
stride=1, downsample=None, reduction=16):
super(BasicBlock, self).__init__()
self.conv1 = meta_conv2d(inplanes, planes, kernel_size = 3, stride=stride, padding=1, bias=False)
self.bn1 = meta_norm(bn_norm, planes, norm_opt=norm_opt)
self.conv2 = meta_conv2d(planes, planes, kernel_size = 3, stride=1, padding=1, bias=False)
self.bn2 = meta_norm(bn_norm, planes, norm_opt=norm_opt)
self.relu = nn.ReLU(inplace=True)
if with_se:
self.se = SELayer(planes, reduction)
else:
self.se = nn.Identity()
self.downsample = downsample
self.stride = stride
def __init__(self,
in_channels,
out_channels,
expansion_factor,
stride=1,
bn_norm='BN',
norm_opt=None):
super(Bottleneck, self).__init__()
mid_channels = in_channels * expansion_factor
self.use_residual = stride == 1 and in_channels == out_channels
self.conv1 = ConvBlock(in_channels,
mid_channels,
1,
bn_norm=bn_norm,
norm_opt=norm_opt)
self.dwconv2 = ConvBlock(mid_channels,
mid_channels,
3,
s=stride,
p=1,
g=mid_channels,
bn_norm=bn_norm,
norm_opt=norm_opt)
# self.conv3 = nn.Sequential(
# meta_conv2d(mid_channels, out_channels, 1, bias=False),
# meta_norm(bn_norm, out_channels, norm_opt = norm_opt),
# )
self.conv3 = meta_conv2d(mid_channels, out_channels, 1, bias=False)
self.bn = meta_norm(bn_norm, out_channels, norm_opt=norm_opt)
def __init__(self, last_stride, bn_norm, norm_opt, num_splits, with_ibn, with_se, with_nl, block, layers, non_layers):
self.inplanes = 64
super().__init__()
self.conv1 = meta_conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = meta_norm(bn_norm, 64, norm_opt = norm_opt)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0], 1, bn_norm, norm_opt, num_splits, with_ibn, with_se)
self.layer2 = self._make_layer(block, 128, layers[1], 2, bn_norm, norm_opt, num_splits, with_ibn, with_se)
self.layer3 = self._make_layer(block, 256, layers[2], 2, bn_norm, norm_opt, num_splits, with_ibn, with_se)
self.layer4 = self._make_layer(block, 512, layers[3], last_stride, bn_norm, norm_opt, num_splits, with_se=with_se)
self.random_init()
if with_nl:
self._build_nonlocal(layers, non_layers, bn_norm, norm_opt, num_splits)
else:
self.NL_1_idx = self.NL_2_idx = self.NL_3_idx = self.NL_4_idx = []
def __init__(self,
in_c,
out_c,
k,
s=1,
p=0,
g=1,
bn_norm='BN',
norm_opt=None):
super(ConvBlock, self).__init__()
# self.conv = nn.Conv2d(in_c, out_c, k, stride=s, padding=p, bias=False, groups=g)
self.conv = meta_conv2d(in_c,
out_c,
k,
stride=s,
padding=p,
bias=False,
groups=g)
self.bn = meta_norm(bn_norm, out_c, norm_opt=norm_opt)
def __init__(self, in_channels, out_channels, kernel_size, stride, bias, bn_norm, norm_opt, bn_out_channels, num_splits):
super(Downsample, self).__init__()
# self = nn.Sequential()
# self.downsample = nn.Module
self.conv = meta_conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias)
self.bn = meta_norm(bn_norm, bn_out_channels, norm_opt =norm_opt)