def __init__(self,
in_channels,
out_channels,
channel_gate,
reduction=4,
T=4,
dropout_cfg=None,
**kwargs):
super(OSBlockINin, self).__init__()
assert T >= 1
assert out_channels >= reduction and out_channels % reduction == 0
mid_channels = out_channels // reduction
self.conv1 = Conv1x1(in_channels, mid_channels)
self.conv2 = nn.ModuleList()
for t in range(1, T + 1):
self.conv2 += [LightConvStream(mid_channels, mid_channels, t)]
self.gate = channel_gate(mid_channels)
self.conv3 = Conv1x1Linear(mid_channels, out_channels, bn=False)
self.downsample = None
if in_channels != out_channels:
self.downsample = Conv1x1Linear(in_channels, out_channels)
self.IN = nn.InstanceNorm2d(out_channels, affine=True)
self.dropout = None
if dropout_cfg is not None:
self.dropout = Dropout(**dropout_cfg)
def _construct_fc_layer(input_dim, output_dim, dropout=None):
layers = []
if dropout:
layers.append(Dropout(**dropout))
layers.extend(
[nn.Linear(input_dim, output_dim),
nn.BatchNorm1d(output_dim)])
return nn.Sequential(*layers)
def _construct_fc_layer(input_dim, output_dim, dropout=False):
layers = []
if dropout:
layers.append(Dropout(p=0.2, dist='gaussian'))
layers.extend([
nn.Linear(input_dim, output_dim),
nn.BatchNorm1d(output_dim)
])
return nn.Sequential(*layers)
def __init__(self, inplaces, hidden_dim, outplaces, kernel_size, stride, use_se, use_hs, dropout_prob=None):
super(InvertedResidual, self).__init__()
assert stride in [1, 2]
self.identity = stride == 1 and inplaces == outplaces
if inplaces == hidden_dim:
self.conv = nn.Sequential(
# dw
nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, (kernel_size - 1) // 2,
groups=hidden_dim, bias=False),
nn.BatchNorm2d(hidden_dim),
HSwish() if use_hs else nn.ReLU(inplace=True),
# Squeeze-and-Excite
SELayer(hidden_dim) if use_se else nn.Sequential(),
# pw-linear
nn.Conv2d(hidden_dim, outplaces, 1, 1, 0, bias=False),
nn.BatchNorm2d(outplaces),
)
else:
self.conv = nn.Sequential(
# pw
nn.Conv2d(inplaces, hidden_dim, 1, 1, 0, bias=False),
nn.BatchNorm2d(hidden_dim),
HSwish() if use_hs else nn.ReLU(inplace=True),
# dw
nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, (kernel_size - 1) // 2,
groups=hidden_dim, bias=False),
nn.BatchNorm2d(hidden_dim),
# Squeeze-and-Excite
SELayer(hidden_dim) if use_se else nn.Sequential(),
HSwish() if use_hs else nn.ReLU(inplace=True),
# pw-linear
nn.Conv2d(hidden_dim, outplaces, 1, 1, 0, bias=False),
nn.BatchNorm2d(outplaces),
)
self.dropout = None
if dropout_prob is not None and dropout_prob > 0.0:
self.dropout = Dropout(p=dropout_prob)
def __init__(self, in_channels, out_channels, channel_gate, reduction=4, T=4, dropout_prob=None, **kwargs):
super(OSBlock, self).__init__()
assert T >= 1
assert out_channels >= reduction and out_channels % reduction == 0
mid_channels = out_channels // reduction
self.conv1 = Conv1x1(in_channels, mid_channels)
self.conv2 = nn.ModuleList()
for t in range(1, T + 1):
self.conv2 += [LightConvStream(mid_channels, mid_channels, t)]
self.gate = channel_gate(mid_channels)
self.conv3 = Conv1x1Linear(mid_channels, out_channels)
self.downsample = None
if in_channels != out_channels:
self.downsample = Conv1x1Linear(in_channels, out_channels)
self.dropout = None
if dropout_prob is not None and dropout_prob > 0.0:
self.dropout = Dropout(p=dropout_prob)
def __init__(self,
in_channels,
out_channels,
IN=False,
bottleneck_reduction=4,
dropout_cfg=None,
channel_gate=ChannelGate,
**kwargs):
super(OSBlock, self).__init__()
mid_channels = out_channels // bottleneck_reduction
self.conv1 = Conv1x1(in_channels, mid_channels)
self.conv2a = LightConv3x3(mid_channels, mid_channels)
self.conv2b = nn.Sequential(
LightConv3x3(mid_channels, mid_channels),
LightConv3x3(mid_channels, mid_channels),
)
self.conv2c = nn.Sequential(
LightConv3x3(mid_channels, mid_channels),
LightConv3x3(mid_channels, mid_channels),
LightConv3x3(mid_channels, mid_channels),
)
self.conv2d = nn.Sequential(
LightConv3x3(mid_channels, mid_channels),
LightConv3x3(mid_channels, mid_channels),
LightConv3x3(mid_channels, mid_channels),
LightConv3x3(mid_channels, mid_channels),
)
self.gate = channel_gate(mid_channels)
self.conv3 = Conv1x1Linear(mid_channels, out_channels)
self.downsample = None
if in_channels != out_channels:
self.downsample = Conv1x1Linear(in_channels, out_channels)
self.IN = None
if IN:
self.IN = nn.InstanceNorm2d(out_channels, affine=True)
self.dropout = None
if dropout_cfg is not None:
self.dropout = Dropout(**dropout_cfg)
def __init__(self,
model_name,
pretrained=False,
dropout_cls = None,
pooling_type='avg',
**kwargs):
super().__init__(**kwargs)
assert self.is_classification(), f"{model_name} model is adapted for classification tasks only"
self.pretrained = pretrained
self.is_mobilenet = True if model_name in ["mobilenetv3_large_100_miil_in21k", "mobilenetv3_large_100_miil"] else False
self.model = timm.create_model(model_name,
pretrained=pretrained,
num_classes=self.num_classes)
self.num_head_features = self.model.num_features
self.num_features = (self.model.conv_head.in_channels if self.is_mobilenet
else self.model.num_features)
self.dropout = Dropout(**dropout_cls)
self.pooling_type = pooling_type
if self.loss in ["am_softmax", "am_binary"]:
self.model.act2 = nn.PReLU()
self.classifier = AngleSimpleLinear(self.num_head_features, self.num_classes)
else:
assert self.loss in ["softmax", "asl", "bce"]
self.classifier = self.model.get_classifier()
def __init__(self,
channels,
init_block_channels,
final_block_channels,
kernel_sizes,
strides_per_stage,
expansion_factors,
tf_mode=False,
bn_eps=1e-5,
in_channels=3,
in_size=(224, 224),
dropout_cls = None,
pooling_type='avg',
bn_eval=False,
bn_frozen=False,
IN_first=False,
IN_conv1=False,
**kwargs):
super().__init__(**kwargs)
self.in_size = in_size
self.input_IN = nn.InstanceNorm2d(3, affine=True) if IN_first else None
self.bn_eval = bn_eval
self.bn_frozen = bn_frozen
self.pooling_type = pooling_type
self.num_features = self.num_head_features = final_block_channels
activation = "swish"
self.features = nn.Sequential()
self.features.add_module("init_block", EffiInitBlock(
in_channels=in_channels,
out_channels=init_block_channels,
bn_eps=bn_eps,
activation=activation,
tf_mode=tf_mode,
IN_conv1=IN_conv1))
in_channels = init_block_channels
for i, channels_per_stage in enumerate(channels):
kernel_sizes_per_stage = kernel_sizes[i]
expansion_factors_per_stage = expansion_factors[i]
stage = nn.Sequential()
for j, out_channels in enumerate(channels_per_stage):
kernel_size = kernel_sizes_per_stage[j]
expansion_factor = expansion_factors_per_stage[j]
stride = strides_per_stage[i] if (j == 0) else 1
if i == 0:
stage.add_module("unit{}".format(j + 1), EffiDwsConvUnit(
in_channels=in_channels,
out_channels=out_channels,
stride=stride,
bn_eps=bn_eps,
activation=activation,
tf_mode=tf_mode))
else:
stage.add_module("unit{}".format(j + 1), EffiInvResUnit(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
exp_factor=expansion_factor,
se_factor=4,
bn_eps=bn_eps,
activation=activation,
tf_mode=tf_mode))
in_channels = out_channels
self.features.add_module("stage{}".format(i + 1), stage)
activation = activation if self.loss == 'softmax' else lambda: nn.PReLU(init=0.25)
self.features.add_module("final_block", conv1x1_block(
in_channels=in_channels,
out_channels=final_block_channels,
bn_eps=bn_eps,
activation=activation))
self.output = nn.Sequential()
if dropout_cls:
self.output.add_module("dropout", Dropout(**dropout_cls))
if self.loss in ['softmax', 'asl']:
self.output.add_module("fc", nn.Linear(
in_features=final_block_channels,
out_features=self.num_classes))
else:
assert self.loss in ['am_softmax', 'am_binary']
self.output.add_module("asl", AngleSimpleLinear(
in_features=final_block_channels,
out_features=self.num_classes))
self._init_params()
def __init__(self,
in_channels=3,
in_size=(224, 224),
num_classes=1000,
dropout_cls=None,
pooling_type='avg',
bn_eval=False,
bn_frozen=False,
loss='softmax',
IN_first=False,
IN_conv1=False,
self_challenging_cfg=False,
**kwargs):
super().__init__(**kwargs)
self.in_size = in_size
self.num_classes = num_classes
self.input_IN = nn.InstanceNorm2d(3, affine=True) if IN_first else None
self.bn_eval = bn_eval
self.bn_frozen = bn_frozen
self.pooling_type = pooling_type
self.loss = loss
self.self_challenging_cfg = self_challenging_cfg
self.num_features = self.num_head_features = 1536
layers = [4, 8, 4]
normal_units = [InceptionAUnit, InceptionBUnit, InceptionCUnit]
reduction_units = [ReductionAUnit, ReductionBUnit]
self.features = nn.Sequential()
self.features.add_module(
"init_block",
InceptInitBlock(in_channels=in_channels,
IN_conv1=IN_conv1,
in_size=in_size[0]))
for i, layers_per_stage in enumerate(layers):
stage = nn.Sequential()
for j in range(layers_per_stage):
if (j == 0) and (i != 0):
unit = reduction_units[i - 1]
else:
unit = normal_units[i]
stage.add_module("unit{}".format(j + 1), unit())
self.features.add_module("stage{}".format(i + 1), stage)
self.output = nn.Sequential()
if dropout_cls:
self.output.add_module("dropout", Dropout(**dropout_cls))
if self.loss in ['softmax', 'asl']:
self.output.add_module(
"fc",
nn.Linear(in_features=self.num_features,
out_features=num_classes))
else:
assert self.loss in ['am_softmax', 'am_binary']
self.output.add_module(
"asl",
AngleSimpleLinear(in_features=self.num_features,
out_features=num_classes))
self._init_params()
def __init__(self,
cfgs,
mode,
IN_conv1=False,
width_mult=1.,
in_channels=3,
input_size=(224, 224),
dropout_cls=None,
pooling_type='avg',
IN_first=False,
self_challenging_cfg=False,
**kwargs):
super().__init__(**kwargs)
self.in_size = input_size
self.input_IN = nn.InstanceNorm2d(in_channels,
affine=True) if IN_first else None
self.pooling_type = pooling_type
self.self_challenging_cfg = self_challenging_cfg
self.width_mult = width_mult
self.dropout_cls = dropout_cls
# setting of inverted residual blocks
self.cfgs = cfgs
assert mode in ['large', 'small']
# building first layer
input_channel = make_divisible(16 * self.width_mult, 8)
stride = 1 if self.in_size[0] < 100 else 2
layers = [conv_3x3_bn(3, input_channel, stride, IN_conv1)]
# building inverted residual blocks
block = InvertedResidual
flag = True
for k, t, c, use_se, use_hs, s in self.cfgs:
if (self.in_size[0] < 100) and (s == 2) and flag:
s = 1
flag = False
output_channel = make_divisible(c * self.width_mult, 8)
exp_size = make_divisible(input_channel * t, 8)
layers.append(
block(input_channel, exp_size, output_channel, k, s, use_se,
use_hs))
input_channel = output_channel
self.features = nn.Sequential(*layers)
# building last several layers
self.conv = conv_1x1_bn(input_channel, exp_size, self.loss)
output_channel = {'large': 1280, 'small': 1024}
output_channel = make_divisible(
output_channel[mode] * self.width_mult,
8) if self.width_mult > 1.0 else output_channel[mode]
self.num_head_features = output_channel
self.num_features = exp_size
if self.loss == 'softmax' or self.loss == 'asl':
self.classifier = nn.Sequential(
nn.Linear(exp_size, output_channel),
nn.BatchNorm1d(output_channel),
HSwish(),
Dropout(**self.dropout_cls),
nn.Linear(output_channel, self.num_classes),
)
else:
assert self.loss in ['am_softmax', 'am_binary']
self.classifier = nn.Sequential(
nn.Linear(exp_size, output_channel),
nn.BatchNorm1d(output_channel),
nn.PReLU(),
Dropout(**self.dropout_cls),
AngleSimpleLinear(output_channel, self.num_classes),
)
self._initialize_weights()