Exemplo n.º 1
0
    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()
Exemplo n.º 2
0
    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()
Exemplo n.º 3
0
 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()
Exemplo n.º 4
0
    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
Exemplo n.º 5
0
    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()
Exemplo n.º 6
0
    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()
Exemplo n.º 8
0
    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()