def __init__(self,
num_classes,
blocks,
layers,
channels,
feature_dim=256,
loss='softmax',
conv1_IN=False,
dropout_prob=0,
activation=nn.ReLU,
**kwargs):
super(OSNet, self).__init__()
num_blocks = len(blocks)
assert num_blocks == len(layers)
assert num_blocks == len(channels) - 1
self.loss = loss
self.feature_dim = feature_dim
# convolutional backbone
self.conv1 = ConvLayer(3,
channels[0],
7,
stride=2,
padding=3,
IN=conv1_IN)
self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
self.conv2 = self._make_layer(blocks[0], layers[0], channels[0],
channels[1])
self.pool2 = nn.Sequential(Conv1x1(channels[1], channels[1]),
nn.AvgPool2d(2, stride=2))
self.conv3 = self._make_layer(blocks[1], layers[1], channels[1],
channels[2])
self.pool3 = nn.Sequential(Conv1x1(channels[2], channels[2]),
nn.AvgPool2d(2, stride=2))
self.conv4 = self._make_layer(blocks[2], layers[2], channels[2],
channels[3])
self.conv5 = Conv1x1(channels[3], channels[3])
self.global_avgpool = nn.Conv2d(channels[3],
channels[3], (16, 8),
groups=channels[3])
# nn.AdaptiveAvgPool2d(1)
# fully connected layer
self.fc = self._construct_fc_layer(self.feature_dim,
channels[3],
dropout_p=dropout_prob)
# identity classification layer
if self.loss not in ['am_softmax', 'adacos', 'd_softmax']:
self.classifier = nn.Linear(self.feature_dim, num_classes)
else:
from torchreid.losses import AngleSimpleLinear
self.classifier = AngleSimpleLinear(self.feature_dim, num_classes)
self._init_params()
def __init__(self, mode, num_classes, width_mult=1.0, feature_dim=512, loss='softmax', input_instance_norm=False,
dropout_prob=None, **kwargs):
super(MobileNetV3, self).__init__()
self.feature_dim = feature_dim
self.loss = loss
# config definition
assert mode in ['large', 'small']
self.cfg = MobileNetV3.arch_settings[mode]
self.instance_norm_input = None
if input_instance_norm:
self.instance_norm_input = nn.InstanceNorm2d(3, affine=True)
# building first layer
input_channel = make_divisible(16 * width_mult, 8)
layers = [conv_3x3_bn(3, input_channel, 2)]
# building inverted residual blocks
block = InvertedResidual
for k, exp_size, c, use_se, use_hs, s in self.cfg:
output_channel = make_divisible(c * width_mult, 8)
layers.append(block(input_channel, exp_size, output_channel, k, s, use_se, use_hs, dropout_prob))
input_channel = output_channel
self.features = nn.Sequential(*layers)
# building last several layers
output_channel = make_divisible(exp_size * width_mult, 8)
self.conv = nn.Sequential(
conv_1x1_bn(input_channel, output_channel),
SELayer(output_channel) if mode == 'small' else nn.Sequential()
)
# building embedding layer
self.global_avgpool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(output_channel, feature_dim),
nn.BatchNorm1d(feature_dim)
)
# building classifier layer
if loss not in ['am_softmax']:
self.classifier = nn.Linear(feature_dim, num_classes)
else:
self.classifier = AngleSimpleLinear(feature_dim, num_classes)
self._init_weights()
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,
model_name,
num_classes,
loss='softmax',
IN_first=False,
pooling_type='avg',
**kwargs):
super().__init__(**kwargs)
assert self.is_classification(), f"{model_name} model is adapted for classification tasks only"
self.pooling_type = pooling_type
self.loss = loss
assert isinstance(num_classes, int)
model = get_model(model_name, num_classes=1000, pretrained=self.pretrained)
assert hasattr(model, 'features') and isinstance(model.features, nn.Sequential)
self.features = model.features
self.features = self.features[:-1] # remove pooling, since it can have a fixed size
if self.loss not in ['am_softmax']:
self.output_conv = nn.Conv2d(in_channels=model.output.in_channels, out_channels=num_classes, kernel_size=1, stride=1, bias=False)
else:
self.output_conv = AngleSimpleLinear(model.output.in_channels, num_classes)
self.num_features = self.num_head_features = model.output.in_channels
self.input_IN = nn.InstanceNorm2d(3, affine=True) if IN_first else None
def __init__(self,
num_classes,
blocks,
layers,
channels,
feature_dim=256,
loss='softmax',
instance_norm=False,
dropout_cfg=None,
fpn_cfg=None,
pooling_type='avg',
input_size=(256, 128),
IN_first=False,
extra_blocks=False,
lct_gate=False,
**kwargs):
self.dropout_cfg = dropout_cfg
self.extra_blocks = extra_blocks
self.channel_gate = LCTGate if lct_gate else ChannelGate
if self.extra_blocks:
for i, l in enumerate(layers):
layers[i] = l + 1
super(OSNetFPN, self).__init__(num_classes, blocks, layers, channels,
feature_dim, loss, instance_norm)
self.feature_scales = (4, 8, 16, 16)
if fpn_cfg is not None:
self.fpn_enable = fpn_cfg.enable
self.fpn_dim = fpn_cfg.dim
self.fpn_process = fpn_cfg.process
assert self.fpn_process in [
'concatenation', 'max_pooling', 'elementwise_sum'
]
else:
self.fpn_enable = False
self.feature_dim = feature_dim
self.use_IN_first = IN_first
if IN_first:
self.in_first = nn.InstanceNorm2d(3, affine=True)
self.conv1 = ConvLayer(3,
channels[0],
7,
stride=2,
padding=3,
IN=self.use_IN_first)
if self.fpn_enable:
self.fpn = FPN(channels, self.feature_scales, self.fpn_dim,
self.fpn_dim)
fpn_out_dim = self.fpn_dim if self.fpn_process in ['max_pooling', 'elementwise_sum'] \
else feature_dim
self.fc = self._construct_fc_layer(feature_dim, fpn_out_dim,
dropout_cfg)
else:
self.fpn = None
self.fc = self._construct_fc_layer(feature_dim, channels[3],
dropout_cfg)
if self.loss not in [
'am_softmax',
]:
self.classifier = nn.Linear(self.feature_dim, num_classes)
else:
self.classifier = AngleSimpleLinear(self.feature_dim, num_classes)
if 'conv' in pooling_type:
kernel_size = (input_size[0] // self.feature_scales[-1],
input_size[1] // self.feature_scales[-1])
if self.fpn_enable:
self.global_avgpool = nn.Conv2d(fpn_out_dim,
fpn_out_dim,
kernel_size,
groups=fpn_out_dim)
else:
self.global_avgpool = nn.Conv2d(channels[3],
channels[3],
kernel_size,
groups=channels[3])
elif 'avg' in pooling_type:
self.global_avgpool = nn.AdaptiveAvgPool2d(1)
elif 'gmp' in pooling_type:
self.global_avgpool = GeneralizedMeanPooling()
else:
raise ValueError('Incorrect pooling type')
if self.fpn_enable and self.fpn_process == 'concatenation':
self.fpn_extra_conv = ConvLayer(self.fpn_dim *
len(self.fpn.dims_out),
feature_dim,
3,
stride=1,
padding=1,
IN=False)
else:
self.fpn_extra_conv = None
self._init_params()
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()