def __init__(self,
in_chs,
out_chs,
stride,
exp_ratio=1.0,
use_se=True,
se_rd=12,
ch_div=1):
super(LinearBottleneck, self).__init__()
self.use_shortcut = stride == 1 and in_chs <= out_chs
self.in_channels = in_chs
self.out_channels = out_chs
if exp_ratio != 1.:
dw_chs = make_divisible(round(in_chs * exp_ratio), divisor=ch_div)
self.conv_exp = ConvBnAct(in_chs, dw_chs, act_layer="swish")
else:
dw_chs = in_chs
self.conv_exp = None
self.conv_dw = ConvBnAct(dw_chs,
dw_chs,
3,
stride=stride,
groups=dw_chs,
apply_act=False)
self.se = SEWithNorm(dw_chs, reduction=se_rd,
divisor=ch_div) if use_se else None
self.act_dw = nn.ReLU6()
self.conv_pwl = ConvBnAct(dw_chs, out_chs, 1, apply_act=False)
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, 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 = ConvBnAct(mid_chs,
out_chs,
kernel_size=3,
dilation=dilation,
groups=groups,
**ckwargs)
self.attn = create_attn(attn_layer, channels=out_chs)
self.drop_path = drop_path
def __init__(self,
block_cfg,
num_classes=1000,
in_chans=3,
output_stride=32,
act_layer=nn.ReLU,
norm_layer=nn.BatchNorm2d,
norm_kwargs=None,
drop_rate=0.,
global_pool='avg'):
super(XceptionAligned, self).__init__()
self.num_classes = num_classes
self.drop_rate = drop_rate
assert output_stride in (8, 16, 32)
norm_kwargs = norm_kwargs if norm_kwargs is not None else {}
layer_args = dict(act_layer=act_layer,
norm_layer=norm_layer,
norm_kwargs=norm_kwargs)
self.stem = nn.Sequential(*[
ConvBnAct(in_chans, 32, kernel_size=3, stride=2, **layer_args),
ConvBnAct(32, 64, kernel_size=3, stride=1, **layer_args)
])
curr_dilation = 1
curr_stride = 2
self.feature_info = []
self.blocks = nn.Sequential()
for i, b in enumerate(block_cfg):
b['dilation'] = curr_dilation
if b['stride'] > 1:
self.feature_info += [
dict(num_chs=tup_triple(b['out_chs'])[-2],
reduction=curr_stride,
module=f'blocks.{i}.stack.act3')
]
next_stride = curr_stride * b['stride']
if next_stride > output_stride:
curr_dilation *= b['stride']
b['stride'] = 1
else:
curr_stride = next_stride
self.blocks.add_module(str(i), XceptionModule(**b, **layer_args))
self.num_features = self.blocks[-1].out_channels
self.feature_info += [
dict(num_chs=self.num_features,
reduction=curr_stride,
module='blocks.' + str(len(self.blocks) - 1))
]
self.head = ClassifierHead(in_chs=self.num_features,
num_classes=num_classes,
pool_type=global_pool,
drop_rate=drop_rate)
def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
cardinality=1,
base_width=64,
sk_kwargs=None,
reduce_first=1,
dilation=1,
first_dilation=None,
act_layer=nn.ReLU,
norm_layer=nn.BatchNorm2d,
attn_layer=None,
aa_layer=None,
drop_block=None,
drop_path=None):
super(SelectiveKernelBottleneck, self).__init__()
sk_kwargs = sk_kwargs or {}
conv_kwargs = dict(drop_block=drop_block,
act_layer=act_layer,
norm_layer=norm_layer,
aa_layer=aa_layer)
width = int(math.floor(planes * (base_width / 64)) * cardinality)
first_planes = width // reduce_first
outplanes = planes * self.expansion
first_dilation = first_dilation or dilation
self.conv1 = ConvBnAct(inplanes,
first_planes,
kernel_size=1,
**conv_kwargs)
self.conv2 = SelectiveKernelConv(first_planes,
width,
stride=stride,
dilation=first_dilation,
groups=cardinality,
**conv_kwargs,
**sk_kwargs)
conv_kwargs['act_layer'] = None
self.conv3 = ConvBnAct(width, outplanes, kernel_size=1, **conv_kwargs)
self.se = create_attn(attn_layer, outplanes)
self.act = act_layer(inplace=True)
self.downsample = downsample
self.stride = stride
self.dilation = dilation
self.drop_block = drop_block
self.drop_path = drop_path
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, num_classes=1001, in_chans=3, output_stride=32, drop_rate=0., global_pool='avg', pad_type=''):
super(PNASNet5Large, self).__init__()
self.num_classes = num_classes
self.drop_rate = drop_rate
self.num_features = 4320
assert output_stride == 32
self.conv_0 = ConvBnAct(
in_chans, 96, kernel_size=3, stride=2, padding=0,
norm_kwargs=dict(eps=0.001, momentum=0.1), act_layer=None)
self.cell_stem_0 = CellStem0(
in_chs_left=96, out_chs_left=54, in_chs_right=96, out_chs_right=54, pad_type=pad_type)
self.cell_stem_1 = Cell(
in_chs_left=96, out_chs_left=108, in_chs_right=270, out_chs_right=108, pad_type=pad_type,
match_prev_layer_dims=True, is_reduction=True)
self.cell_0 = Cell(
in_chs_left=270, out_chs_left=216, in_chs_right=540, out_chs_right=216, pad_type=pad_type,
match_prev_layer_dims=True)
self.cell_1 = Cell(
in_chs_left=540, out_chs_left=216, in_chs_right=1080, out_chs_right=216, pad_type=pad_type)
self.cell_2 = Cell(
in_chs_left=1080, out_chs_left=216, in_chs_right=1080, out_chs_right=216, pad_type=pad_type)
self.cell_3 = Cell(
in_chs_left=1080, out_chs_left=216, in_chs_right=1080, out_chs_right=216, pad_type=pad_type)
self.cell_4 = Cell(
in_chs_left=1080, out_chs_left=432, in_chs_right=1080, out_chs_right=432, pad_type=pad_type,
is_reduction=True)
self.cell_5 = Cell(
in_chs_left=1080, out_chs_left=432, in_chs_right=2160, out_chs_right=432, pad_type=pad_type,
match_prev_layer_dims=True)
self.cell_6 = Cell(
in_chs_left=2160, out_chs_left=432, in_chs_right=2160, out_chs_right=432, pad_type=pad_type)
self.cell_7 = Cell(
in_chs_left=2160, out_chs_left=432, in_chs_right=2160, out_chs_right=432, pad_type=pad_type)
self.cell_8 = Cell(
in_chs_left=2160, out_chs_left=864, in_chs_right=2160, out_chs_right=864, pad_type=pad_type,
is_reduction=True)
self.cell_9 = Cell(
in_chs_left=2160, out_chs_left=864, in_chs_right=4320, out_chs_right=864, pad_type=pad_type,
match_prev_layer_dims=True)
self.cell_10 = Cell(
in_chs_left=4320, out_chs_left=864, in_chs_right=4320, out_chs_right=864, pad_type=pad_type)
self.cell_11 = Cell(
in_chs_left=4320, out_chs_left=864, in_chs_right=4320, out_chs_right=864, pad_type=pad_type)
self.act = nn.ReLU()
self.feature_info = [
dict(num_chs=96, reduction=2, module='conv_0'),
dict(num_chs=270, reduction=4, module='cell_stem_1.conv_1x1.act'),
dict(num_chs=1080, reduction=8, module='cell_4.conv_1x1.act'),
dict(num_chs=2160, reduction=16, module='cell_8.conv_1x1.act'),
dict(num_chs=4320, reduction=32, module='act'),
]
self.global_pool, self.last_linear = create_classifier(
self.num_features, self.num_classes, pool_type=global_pool)
def __init__(self,
in_chs,
mid_chs,
out_chs,
layer_per_block,
residual=False,
depthwise=False,
attn='',
norm_layer=BatchNormAct2d,
act_layer=nn.ReLU):
super(OsaBlock, self).__init__()
self.residual = residual
self.depthwise = depthwise
conv_kwargs = dict(norm_layer=norm_layer, act_layer=act_layer)
next_in_chs = in_chs
if self.depthwise and next_in_chs != mid_chs:
assert not residual
self.conv_reduction = ConvBnAct(next_in_chs, mid_chs, 1,
**conv_kwargs)
else:
self.conv_reduction = None
mid_convs = []
for i in range(layer_per_block):
if self.depthwise:
conv = SeparableConvBnAct(mid_chs, mid_chs, **conv_kwargs)
else:
conv = ConvBnAct(next_in_chs, mid_chs, 3, **conv_kwargs)
next_in_chs = mid_chs
mid_convs.append(conv)
self.conv_mid = SequentialAppendList(*mid_convs)
# feature aggregation
next_in_chs = in_chs + layer_per_block * mid_chs
self.conv_concat = ConvBnAct(next_in_chs, out_chs, **conv_kwargs)
if attn:
self.attn = create_attn(attn, out_chs)
else:
self.attn = 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,
in_chs,
out_chs,
dilation=1,
bottle_ratio=0.25,
groups=1,
act_layer=nn.ReLU,
norm_layer=nn.BatchNorm2d,
attn_last=False,
attn_layer=None,
aa_layer=None,
drop_block=None,
drop_path=None):
super(ResBottleneck, 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 = ConvBnAct(mid_chs,
mid_chs,
kernel_size=3,
dilation=dilation,
groups=groups,
**ckwargs)
self.attn2 = create_attn(attn_layer,
channels=mid_chs) if not attn_last else None
self.conv3 = ConvBnAct(mid_chs,
out_chs,
kernel_size=1,
apply_act=False,
**ckwargs)
self.attn3 = create_attn(attn_layer,
channels=out_chs) if attn_last else None
self.drop_path = drop_path
self.act3 = act_layer(inplace=True)
def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
cardinality=1,
base_width=64,
sk_kwargs=None,
reduce_first=1,
dilation=1,
first_dilation=None,
act_layer=nn.ReLU,
norm_layer=nn.BatchNorm2d,
attn_layer=None,
aa_layer=None,
drop_block=None,
drop_path=None):
super(SelectiveKernelBasic, self).__init__()
sk_kwargs = sk_kwargs or {}
conv_kwargs = dict(drop_block=drop_block,
act_layer=act_layer,
norm_layer=norm_layer,
aa_layer=aa_layer)
assert cardinality == 1, 'BasicBlock only supports cardinality of 1'
assert base_width == 64, 'BasicBlock doest not support changing base width'
first_planes = planes // reduce_first
outplanes = planes * self.expansion
first_dilation = first_dilation or dilation
self.conv1 = SelectiveKernelConv(inplanes,
first_planes,
stride=stride,
dilation=first_dilation,
**conv_kwargs,
**sk_kwargs)
conv_kwargs['act_layer'] = None
self.conv2 = ConvBnAct(first_planes,
outplanes,
kernel_size=3,
dilation=dilation,
**conv_kwargs)
self.se = create_attn(attn_layer, outplanes)
self.act = act_layer(inplace=True)
self.downsample = downsample
self.stride = stride
self.dilation = dilation
self.drop_block = drop_block
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_chs,
out_chs,
stride=1,
dilation=1,
pad_type='',
start_with_relu=True,
no_skip=False,
act_layer=nn.ReLU,
norm_layer=None,
norm_kwargs=None):
super(XceptionModule, self).__init__()
norm_kwargs = norm_kwargs if norm_kwargs is not None else {}
out_chs = tup_triple(out_chs)
self.in_channels = in_chs
self.out_channels = out_chs[-1]
self.no_skip = no_skip
if not no_skip and (self.out_channels != self.in_channels
or stride != 1):
self.shortcut = ConvBnAct(in_chs,
self.out_channels,
1,
stride=stride,
norm_layer=norm_layer,
norm_kwargs=norm_kwargs,
act_layer=None)
else:
self.shortcut = None
separable_act_layer = None if start_with_relu else act_layer
self.stack = nn.Sequential()
for i in range(3):
if start_with_relu:
self.stack.add_module(f'act{i + 1}', nn.ReLU(inplace=i > 0))
self.stack.add_module(
f'conv{i + 1}',
SeparableConv2d(in_chs,
out_chs[i],
3,
stride=stride if i == 2 else 1,
dilation=dilation,
padding=pad_type,
act_layer=separable_act_layer,
norm_layer=norm_layer,
norm_kwargs=norm_kwargs))
in_chs = out_chs[i]
def __init__(self,
in_chs,
out_chs,
stride,
dilation,
depth,
block_ratio=1.,
bottle_ratio=1.,
groups=1,
first_dilation=None,
block_fn=ResBottleneck,
block_dpr=None,
**block_kwargs):
super(DarkStage, self).__init__()
first_dilation = first_dilation or dilation
self.conv_down = ConvBnAct(in_chs,
out_chs,
kernel_size=3,
stride=stride,
dilation=first_dilation,
groups=groups,
act_layer=block_kwargs.get('act_layer'),
norm_layer=block_kwargs.get('norm_layer'),
aa_layer=block_kwargs.get('aa_layer', None))
prev_chs = out_chs
block_out_chs = int(round(out_chs * block_ratio))
self.blocks = nn.Sequential()
for i in range(depth):
drop_path = DropPath(
block_dpr[i]) if block_dpr and block_dpr[i] else None
self.blocks.add_module(
str(i),
block_fn(prev_chs,
block_out_chs,
dilation,
bottle_ratio,
groups,
drop_path=drop_path,
**block_kwargs))
prev_chs = block_out_chs
def _build_blocks(block_cfg,
prev_chs,
width_mult,
se_rd=12,
ch_div=1,
feature_location='bottleneck'):
feat_exp = feature_location == 'expansion'
feat_chs = [prev_chs]
feature_info = []
curr_stride = 2
features = []
for block_idx, (chs, exp_ratio, stride, se) in enumerate(block_cfg):
if stride > 1:
fname = 'stem' if block_idx == 0 else f'features.{block_idx - 1}'
if block_idx > 0 and feat_exp:
fname += '.act_dw'
feature_info += [
dict(num_chs=feat_chs[-1], reduction=curr_stride, module=fname)
]
curr_stride *= stride
features.append(
LinearBottleneck(in_chs=prev_chs,
out_chs=chs,
exp_ratio=exp_ratio,
stride=stride,
use_se=se,
se_rd=se_rd,
ch_div=ch_div))
prev_chs = chs
feat_chs += [features[-1].feat_channels(feat_exp)]
pen_chs = make_divisible(1280 * width_mult, divisor=ch_div)
feature_info += [
dict(num_chs=pen_chs if feat_exp else feat_chs[-1],
reduction=curr_stride,
module=f'features.{len(features) - int(not feat_exp)}')
]
features.append(ConvBnAct(prev_chs, pen_chs, act_layer="swish"))
return features, feature_info
def __init__(self,
in_chans=3,
num_classes=1000,
global_pool='avg',
output_stride=32,
initial_chs=16,
final_chs=180,
width_mult=1.0,
depth_mult=1.0,
use_se=True,
se_rd=12,
ch_div=1,
drop_rate=0.2,
feature_location='bottleneck'):
super(ReXNetV1, self).__init__()
self.drop_rate = drop_rate
assert output_stride == 32 # FIXME support dilation
stem_base_chs = 32 / width_mult if width_mult < 1.0 else 32
stem_chs = make_divisible(round(stem_base_chs * width_mult),
divisor=ch_div)
self.stem = ConvBnAct(in_chans,
stem_chs,
3,
stride=2,
act_layer='swish')
block_cfg = _block_cfg(width_mult, depth_mult, initial_chs, final_chs,
use_se, ch_div)
features, self.feature_info = _build_blocks(block_cfg, stem_chs,
width_mult, se_rd, ch_div,
feature_location)
self.num_features = features[-1].out_channels
self.features = nn.Sequential(*features)
self.head = ClassifierHead(self.num_features, num_classes, global_pool,
drop_rate)
def create_stem(in_chans=3,
out_chs=32,
kernel_size=3,
stride=2,
pool='',
act_layer=None,
norm_layer=None,
aa_layer=None):
stem = nn.Sequential()
if not isinstance(out_chs, (tuple, list)):
out_chs = [out_chs]
assert len(out_chs)
in_c = in_chans
for i, out_c in enumerate(out_chs):
conv_name = f'conv{i + 1}'
stem.add_module(
conv_name,
ConvBnAct(in_c,
out_c,
kernel_size,
stride=stride if i == 0 else 1,
act_layer=act_layer,
norm_layer=norm_layer))
in_c = out_c
last_conv = conv_name
if pool:
if aa_layer is not None:
stem.add_module('pool',
nn.MaxPool2d(kernel_size=3, stride=1, padding=1))
stem.add_module('aa', aa_layer(channels=in_c, stride=2))
else:
stem.add_module('pool',
nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
return stem, dict(num_chs=in_c,
reduction=stride,
module='.'.join(['stem', last_conv]))
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):
super().__init__()
# TODO add drop block, drop path, anti-aliasing, custom bn/act args
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, **stage_args, se_ratio=se_ratio))
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,
small=False,
num_init_features=64,
k_r=96,
groups=32,
b=False,
k_sec=(3, 4, 20, 3),
inc_sec=(16, 32, 24, 128),
output_stride=32,
num_classes=1000,
in_chans=3,
drop_rate=0.,
global_pool='avg',
fc_act=nn.ELU):
super(DPN, self).__init__()
self.num_classes = num_classes
self.drop_rate = drop_rate
self.b = b
assert output_stride == 32 # FIXME look into dilation support
bw_factor = 1 if small else 4
blocks = OrderedDict()
# conv1
blocks['conv1_1'] = ConvBnAct(in_chans,
num_init_features,
kernel_size=3 if small else 7,
stride=2,
norm_kwargs=dict(eps=.001))
blocks['conv1_pool'] = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.feature_info = [
dict(num_chs=num_init_features,
reduction=2,
module='features.conv1_1')
]
# conv2
bw = 64 * bw_factor
inc = inc_sec[0]
r = (k_r * bw) // (64 * bw_factor)
blocks['conv2_1'] = DualPathBlock(num_init_features, r, r, bw, inc,
groups, 'proj', b)
in_chs = bw + 3 * inc
for i in range(2, k_sec[0] + 1):
blocks['conv2_' + str(i)] = DualPathBlock(in_chs, r, r, bw, inc,
groups, 'normal', b)
in_chs += inc
self.feature_info += [
dict(num_chs=in_chs,
reduction=4,
module=f'features.conv2_{k_sec[0]}')
]
# conv3
bw = 128 * bw_factor
inc = inc_sec[1]
r = (k_r * bw) // (64 * bw_factor)
blocks['conv3_1'] = DualPathBlock(in_chs, r, r, bw, inc, groups,
'down', b)
in_chs = bw + 3 * inc
for i in range(2, k_sec[1] + 1):
blocks['conv3_' + str(i)] = DualPathBlock(in_chs, r, r, bw, inc,
groups, 'normal', b)
in_chs += inc
self.feature_info += [
dict(num_chs=in_chs,
reduction=8,
module=f'features.conv3_{k_sec[1]}')
]
# conv4
bw = 256 * bw_factor
inc = inc_sec[2]
r = (k_r * bw) // (64 * bw_factor)
blocks['conv4_1'] = DualPathBlock(in_chs, r, r, bw, inc, groups,
'down', b)
in_chs = bw + 3 * inc
for i in range(2, k_sec[2] + 1):
blocks['conv4_' + str(i)] = DualPathBlock(in_chs, r, r, bw, inc,
groups, 'normal', b)
in_chs += inc
self.feature_info += [
dict(num_chs=in_chs,
reduction=16,
module=f'features.conv4_{k_sec[2]}')
]
# conv5
bw = 512 * bw_factor
inc = inc_sec[3]
r = (k_r * bw) // (64 * bw_factor)
blocks['conv5_1'] = DualPathBlock(in_chs, r, r, bw, inc, groups,
'down', b)
in_chs = bw + 3 * inc
for i in range(2, k_sec[3] + 1):
blocks['conv5_' + str(i)] = DualPathBlock(in_chs, r, r, bw, inc,
groups, 'normal', b)
in_chs += inc
self.feature_info += [
dict(num_chs=in_chs,
reduction=32,
module=f'features.conv5_{k_sec[3]}')
]
def _fc_norm(f, eps):
return BatchNormAct2d(f, eps=eps, act_layer=fc_act, inplace=False)
blocks['conv5_bn_ac'] = CatBnAct(in_chs, norm_layer=_fc_norm)
self.num_features = in_chs
self.features = nn.Sequential(blocks)
# Using 1x1 conv for the FC layer to allow the extra pooling scheme
self.global_pool, self.classifier = create_classifier(
self.num_features,
self.num_classes,
pool_type=global_pool,
use_conv=True)
def __init__(self,
cfg,
in_chans=3,
num_classes=1000,
global_pool='avg',
drop_rate=0.,
stem_stride=4,
output_stride=32,
norm_layer=BatchNormAct2d,
act_layer=nn.ReLU):
""" VovNet (v2)
"""
super(VovNet, self).__init__()
self.num_classes = num_classes
self.drop_rate = drop_rate
assert stem_stride in (4, 2)
assert output_stride == 32 # FIXME support dilation
stem_chs = cfg["stem_chs"]
stage_conv_chs = cfg["stage_conv_chs"]
stage_out_chs = cfg["stage_out_chs"]
block_per_stage = cfg["block_per_stage"]
layer_per_block = cfg["layer_per_block"]
conv_kwargs = dict(norm_layer=norm_layer, act_layer=act_layer)
# Stem module
last_stem_stride = stem_stride // 2
conv_type = SeparableConvBnAct if cfg["depthwise"] else ConvBnAct
self.stem = nn.Sequential(*[
ConvBnAct(in_chans, stem_chs[0], 3, stride=2, **conv_kwargs),
conv_type(stem_chs[0], stem_chs[1], 3, stride=1, **conv_kwargs),
conv_type(stem_chs[1],
stem_chs[2],
3,
stride=last_stem_stride,
**conv_kwargs),
])
self.feature_info = [
dict(num_chs=stem_chs[1],
reduction=2,
module=f'stem.{1 if stem_stride == 4 else 2}')
]
current_stride = stem_stride
# OSA stages
in_ch_list = stem_chs[-1:] + stage_out_chs[:-1]
stage_args = dict(residual=cfg["residual"],
depthwise=cfg["depthwise"],
attn=cfg["attn"],
**conv_kwargs)
stages = []
for i in range(4): # num_stages
downsample = stem_stride == 2 or i > 0 # first stage has no stride/downsample if stem_stride is 4
stages += [
OsaStage(in_ch_list[i],
stage_conv_chs[i],
stage_out_chs[i],
block_per_stage[i],
layer_per_block,
downsample=downsample,
**stage_args)
]
self.num_features = stage_out_chs[i]
current_stride *= 2 if downsample else 1
self.feature_info += [
dict(num_chs=self.num_features,
reduction=current_stride,
module=f'stages.{i}')
]
self.stages = nn.Sequential(*stages)
self.head = ClassifierHead(self.num_features,
num_classes,
pool_type=global_pool,
drop_rate=drop_rate)
for n, m in self.named_modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1.)
nn.init.constant_(m.bias, 0.)
elif isinstance(m, nn.Linear):
nn.init.zeros_(m.bias)
def __init__(self,
num_classes=1000,
in_chans=1,
stem_size=96,
channel_multiplier=2,
num_features=4032,
output_stride=32,
drop_rate=0.,
global_pool='avg',
pad_type='same'):
super(NASNetALarge, self).__init__()
self.num_classes = num_classes
self.stem_size = stem_size
self.num_features = num_features
self.channel_multiplier = channel_multiplier
self.drop_rate = drop_rate
assert output_stride == 32
channels = self.num_features // 24
# 24 is default value for the architecture
self.conv0 = ConvBnAct(in_channels=in_chans,
out_channels=self.stem_size,
kernel_size=3,
padding=0,
stride=2,
norm_kwargs=dict(eps=0.001, momentum=0.1),
act_layer=None)
self.cell_stem_0 = CellStem0(self.stem_size,
num_channels=channels //
(channel_multiplier**2),
pad_type=pad_type)
self.cell_stem_1 = CellStem1(self.stem_size,
num_channels=channels //
channel_multiplier,
pad_type=pad_type)
self.cell_0 = FirstCell(in_chs_left=channels,
out_chs_left=channels // 2,
in_chs_right=2 * channels,
out_chs_right=channels,
pad_type=pad_type)
self.cell_1 = NormalCell(in_chs_left=2 * channels,
out_chs_left=channels,
in_chs_right=6 * channels,
out_chs_right=channels,
pad_type=pad_type)
self.cell_2 = NormalCell(in_chs_left=6 * channels,
out_chs_left=channels,
in_chs_right=6 * channels,
out_chs_right=channels,
pad_type=pad_type)
self.cell_3 = NormalCell(in_chs_left=6 * channels,
out_chs_left=channels,
in_chs_right=6 * channels,
out_chs_right=channels,
pad_type=pad_type)
self.cell_4 = NormalCell(in_chs_left=6 * channels,
out_chs_left=channels,
in_chs_right=6 * channels,
out_chs_right=channels,
pad_type=pad_type)
self.cell_5 = NormalCell(in_chs_left=6 * channels,
out_chs_left=channels,
in_chs_right=6 * channels,
out_chs_right=channels,
pad_type=pad_type)
self.reduction_cell_0 = ReductionCell0(in_chs_left=6 * channels,
out_chs_left=2 * channels,
in_chs_right=6 * channels,
out_chs_right=2 * channels,
pad_type=pad_type)
self.cell_6 = FirstCell(in_chs_left=6 * channels,
out_chs_left=channels,
in_chs_right=8 * channels,
out_chs_right=2 * channels,
pad_type=pad_type)
self.cell_7 = NormalCell(in_chs_left=8 * channels,
out_chs_left=2 * channels,
in_chs_right=12 * channels,
out_chs_right=2 * channels,
pad_type=pad_type)
self.cell_8 = NormalCell(in_chs_left=12 * channels,
out_chs_left=2 * channels,
in_chs_right=12 * channels,
out_chs_right=2 * channels,
pad_type=pad_type)
self.cell_9 = NormalCell(in_chs_left=12 * channels,
out_chs_left=2 * channels,
in_chs_right=12 * channels,
out_chs_right=2 * channels,
pad_type=pad_type)
self.cell_10 = NormalCell(in_chs_left=12 * channels,
out_chs_left=2 * channels,
in_chs_right=12 * channels,
out_chs_right=2 * channels,
pad_type=pad_type)
self.cell_11 = NormalCell(in_chs_left=12 * channels,
out_chs_left=2 * channels,
in_chs_right=12 * channels,
out_chs_right=2 * channels,
pad_type=pad_type)
self.reduction_cell_1 = ReductionCell1(in_chs_left=12 * channels,
out_chs_left=4 * channels,
in_chs_right=12 * channels,
out_chs_right=4 * channels,
pad_type=pad_type)
self.cell_12 = FirstCell(in_chs_left=12 * channels,
out_chs_left=2 * channels,
in_chs_right=16 * channels,
out_chs_right=4 * channels,
pad_type=pad_type)
self.cell_13 = NormalCell(in_chs_left=16 * channels,
out_chs_left=4 * channels,
in_chs_right=24 * channels,
out_chs_right=4 * channels,
pad_type=pad_type)
self.cell_14 = NormalCell(in_chs_left=24 * channels,
out_chs_left=4 * channels,
in_chs_right=24 * channels,
out_chs_right=4 * channels,
pad_type=pad_type)
self.cell_15 = NormalCell(in_chs_left=24 * channels,
out_chs_left=4 * channels,
in_chs_right=24 * channels,
out_chs_right=4 * channels,
pad_type=pad_type)
self.cell_16 = NormalCell(in_chs_left=24 * channels,
out_chs_left=4 * channels,
in_chs_right=24 * channels,
out_chs_right=4 * channels,
pad_type=pad_type)
self.cell_17 = NormalCell(in_chs_left=24 * channels,
out_chs_left=4 * channels,
in_chs_right=24 * channels,
out_chs_right=4 * channels,
pad_type=pad_type)
self.act = nn.ReLU(inplace=True)
self.feature_info = [
dict(num_chs=96, reduction=2, module='conv0'),
dict(num_chs=168, reduction=4, module='cell_stem_1.conv_1x1.act'),
dict(num_chs=1008,
reduction=8,
module='reduction_cell_0.conv_1x1.act'),
dict(num_chs=2016,
reduction=16,
module='reduction_cell_1.conv_1x1.act'),
dict(num_chs=4032, reduction=32, module='act'),
]
self.global_pool, self.last_linear = create_classifier(
self.num_features, self.num_classes, pool_type=global_pool)
def __init__(self,
in_chs,
out_chs,
stride,
dilation,
depth,
block_ratio=1.,
bottle_ratio=1.,
exp_ratio=1.,
groups=1,
first_dilation=None,
down_growth=False,
cross_linear=False,
block_dpr=None,
block_fn=ResBottleneck,
**block_kwargs):
super(CrossStage, self).__init__()
first_dilation = first_dilation or dilation
down_chs = out_chs if down_growth else in_chs # grow downsample channels to output channels
exp_chs = int(round(out_chs * exp_ratio))
block_out_chs = int(round(out_chs * block_ratio))
conv_kwargs = dict(act_layer=block_kwargs.get('act_layer'),
norm_layer=block_kwargs.get('norm_layer'))
if stride != 1 or first_dilation != dilation:
self.conv_down = ConvBnAct(in_chs,
down_chs,
kernel_size=3,
stride=stride,
dilation=first_dilation,
groups=groups,
aa_layer=block_kwargs.get(
'aa_layer', None),
**conv_kwargs)
prev_chs = down_chs
else:
self.conv_down = None
prev_chs = in_chs
# FIXME this 1x1 expansion is pushed down into the cross and block paths in the darknet cfgs. Also,
# there is also special case for the first stage for some of the model that results in uneven split
# across the two paths. I did it this way for simplicity for now.
self.conv_exp = ConvBnAct(prev_chs,
exp_chs,
kernel_size=1,
apply_act=not cross_linear,
**conv_kwargs)
prev_chs = exp_chs // 2 # output of conv_exp is always split in two
self.blocks = nn.Sequential()
for i in range(depth):
drop_path = DropPath(
block_dpr[i]) if block_dpr and block_dpr[i] else None
self.blocks.add_module(
str(i),
block_fn(prev_chs,
block_out_chs,
dilation,
bottle_ratio,
groups,
drop_path=drop_path,
**block_kwargs))
prev_chs = block_out_chs
# transition convs
self.conv_transition_b = ConvBnAct(prev_chs,
exp_chs // 2,
kernel_size=1,
**conv_kwargs)
self.conv_transition = ConvBnAct(exp_chs,
out_chs,
kernel_size=1,
**conv_kwargs)