def downsample_avg(in_chs,
out_chs,
kernel_size,
stride=1,
dilation=1,
norm_layer=None):
""" AvgPool Downsampling as in 'D' ResNet variants. This is not in RegNet space but I might experiment."""
norm_layer = norm_layer or nn.BatchNorm2d
avg_stride = stride if dilation == 1 else 1
pool = nn.Identity()
if stride > 1 or dilation > 1:
avg_pool_fn = AvgPool2dSame if avg_stride == 1 and dilation > 1 else nn.AvgPool2d
pool = avg_pool_fn(2,
avg_stride,
ceil_mode=True,
count_include_pad=False)
return nn.Sequential(*[
pool,
ConvBnAct(in_chs,
out_chs,
1,
stride=1,
norm_layer=norm_layer,
act_layer=None)
])
def __init__(self,
in_chs,
out_chs,
dilation=1,
bottle_ratio=0.5,
groups=1,
act_layer=nn.ReLU,
norm_layer=nn.BatchNorm2d,
attn_layer=None,
aa_layer=None,
drop_block=None,
drop_path=None):
super(DarkBlock_ACON, self).__init__()
mid_chs = int(round(out_chs * bottle_ratio))
ckwargs = dict(act_layer=act_layer,
norm_layer=norm_layer,
aa_layer=aa_layer,
drop_block=drop_block)
self.conv1 = ConvBnAct(in_chs, mid_chs, kernel_size=1, **ckwargs)
self.conv2 = Conv3x3_ACON(mid_chs,
out_chs,
kernel_size=3,
dilation=dilation,
groups=groups)
self.attn = create_attn(attn_layer, channels=out_chs)
self.drop_path = drop_path
def downsample_conv(
in_chs, out_chs, kernel_size, stride=1, dilation=1, norm_layer=None):
norm_layer = norm_layer or nn.BatchNorm2d
kernel_size = 1 if stride == 1 and dilation == 1 else kernel_size
dilation = dilation if kernel_size > 1 else 1
return ConvBnAct(
in_chs, out_chs, kernel_size, stride=stride, dilation=dilation, norm_layer=norm_layer, act_layer=None)
def __init__(self,
in_channels,
out_channels,
stride=1,
bottleneck_factor=1,
group_width=1,
se_ratio=0.0):
super(Bottleneck, self).__init__()
self.use_se = (se_ratio > 0.0)
assert (bottleneck_factor == 1)
mid_channels = out_channels // bottleneck_factor
groups = mid_channels // group_width
self.conv1 = ConvBnAct(in_channels, mid_channels, kernel_size=1)
self.conv2 = ConvBnAct(mid_channels,
mid_channels,
kernel_size=3,
stride=stride,
dilation=1,
groups=groups)
if self.use_se:
self.se = SEModule(mid_channels,
reduction_channels=int(
round(in_channels * se_ratio)))
self.conv3 = ConvBnAct(mid_channels,
out_channels,
kernel_size=1,
act_layer=None)
self.act3 = nn.ReLU(inplace=True)
if (in_channels != out_channels) or (stride != 1):
self.downsample = ConvBnAct(in_channels=in_channels,
out_channels=out_channels,
kernel_size=1,
stride=stride,
dilation=1,
norm_layer=nn.BatchNorm2d,
act_layer=None)
else:
self.downsample = None
def __init__(self,
in_chs,
out_chs,
stride=1,
dilation=1,
bottleneck_ratio=1,
group_width=1,
se_ratio=0.25,
downsample=None,
act_layer=nn.ReLU,
norm_layer=nn.BatchNorm2d,
aa_layer=None,
drop_block=None,
drop_path=None):
super(Bottleneck, self).__init__()
bottleneck_chs = int(round(out_chs * bottleneck_ratio))
groups = bottleneck_chs // group_width
cargs = dict(act_layer=act_layer,
norm_layer=norm_layer,
aa_layer=aa_layer,
drop_block=drop_block)
self.conv1 = ConvBnAct(in_chs, bottleneck_chs, kernel_size=1, **cargs)
self.conv2 = ConvBnAct(bottleneck_chs,
bottleneck_chs,
kernel_size=3,
stride=stride,
dilation=dilation,
groups=groups,
**cargs)
if se_ratio:
se_channels = int(round(in_chs * se_ratio))
self.se = SEModule(bottleneck_chs, reduction_channels=se_channels)
else:
self.se = None
cargs['act_layer'] = None
self.conv3 = ConvBnAct(bottleneck_chs, out_chs, kernel_size=1, **cargs)
self.act3 = act_layer(inplace=True)
self.downsample = downsample
self.drop_path = drop_path
def __init__(self, cfg, in_chans=3, num_classes=1000, output_stride=32, global_pool='avg', drop_rate=0.,
drop_path_rate=0., zero_init_last_bn=True, in_channels=3, in_size=(224, 224)):
super().__init__()
# TODO add drop block, drop path, anti-aliasing, custom bn/act args
self.in_size = in_size
self.num_classes = num_classes
self.drop_rate = drop_rate
assert output_stride in (8, 16, 32)
# Construct the stem
stem_width = cfg['stem_width']
self.stem = ConvBnAct(in_chans, stem_width, 3, stride=2)
self.feature_info = [dict(num_chs=stem_width, reduction=2, module='stem')]
# Construct the stages
prev_width = stem_width
curr_stride = 2
stage_params = self._get_stage_params(cfg, output_stride=output_stride, drop_path_rate=drop_path_rate)
se_ratio = cfg['se_ratio']
for i, stage_args in enumerate(stage_params):
stage_name = "s{}".format(i + 1)
self.add_module(stage_name, RegStage(prev_width, se_ratio=se_ratio, **stage_args))
prev_width = stage_args['out_chs']
curr_stride *= stage_args['stride']
self.feature_info += [dict(num_chs=prev_width, reduction=curr_stride, module=stage_name)]
# Construct the head
self.num_features = prev_width
self.head = ClassifierHead(
in_chs=prev_width, num_classes=num_classes, pool_type=global_pool, drop_rate=drop_rate)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, mean=0.0, std=0.01)
nn.init.zeros_(m.bias)
if zero_init_last_bn:
for m in self.modules():
if hasattr(m, 'zero_init_last_bn'):
m.zero_init_last_bn()
def __init__(self,
cfg,
in_channels=3,
in_size=(224, 224),
num_classes=1000):
super().__init__()
self.in_size = in_size
self.num_classes = num_classes
self.features = nn.Sequential()
stem_width = 32
# Construct the stem
self.features.add_module(
"stem",
ConvBnAct(in_channels=in_channels,
out_channels=stem_width,
kernel_size=3,
stride=2))
# Construct the stages
in_channels = stem_width
stage_params = self._get_stage_params(cfg)
se_ratio = cfg['se_ratio']
for i, stage_args in enumerate(stage_params):
stage_name = "s{}".format(i + 1)
self.features.add_module(
stage_name,
RegStage(in_channels=in_channels,
se_ratio=se_ratio,
**stage_args))
in_channels = stage_args['out_channels']
self.features.add_module("global_pool",
nn.AdaptiveAvgPool2d(output_size=1))
self.fc = nn.Linear(in_features=in_channels,
out_features=num_classes,
bias=True)
pass