fixed multilabel configs (#67)

openvinotoolkit · Oct 10, 2022 · bc04c98 · bc04c98
1 parent 7efc90a
commit bc04c98
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 8 deletions.
diff --git a/models/classification/ote_efficientnet_b0_multilabel_al.yaml b/models/classification/ote_efficientnet_b0_multilabel_al.yaml
@@ -2,11 +2,9 @@ _base_: ./ote_efficientnet_b0_multilabel.yaml
 
 model:
     head:
-        type: CustomMultiLabelNonLinearClsHead
+        type: CustomMultiLabelLinearClsHead
         normalized: True
         scale: 7.0
-        act_cfg:
-            type: PReLU
         loss:
             type: AsymmetricAngularLossWithIgnore
             gamma_pos: 0.0

diff --git a/models/classification/ote_efficientnet_v2_s_multilabel_al.yaml b/models/classification/ote_efficientnet_v2_s_multilabel_al.yaml
@@ -2,11 +2,9 @@ _base_: ./ote_efficientnet_v2_s_multilabel.yaml
 
 model:
     head:
-        type: CustomMultiLabelNonLinearClsHead
+        type: CustomMultiLabelLinearClsHead
         normalized: True
         scale: 7.0
-        act_cfg:
-            type: PReLU
         loss:
             type: AsymmetricAngularLossWithIgnore
             gamma_pos: 0.0

diff --git a/mpa/modules/models/classifiers/sam_classifier.py b/mpa/modules/models/classifiers/sam_classifier.py
@@ -171,7 +171,7 @@ def load_state_dict_pre_hook(module, state_dict, *args, **kwargs):
                         k = k.replace('asl', 'fc')
                         v = v.t()
                 state_dict[k] = v
-        
+
         elif backbone_type == 'OTEEfficientNetV2':
             for k in list(state_dict.keys()):
                 v = state_dict.pop(k)
@@ -222,7 +222,7 @@ def load_state_dict_mixing_hook(model, model_classes, chkpt_classes, chkpt_dict,
         for model_name in param_names:
             model_param = model_dict[model_name].clone()
             if backbone_type == 'OTEMobileNetV3':
-                chkpt_name = 'head.'+model_name.replace('4', '3')
+                chkpt_name = 'head.' + model_name.replace('4', '3')
                 if model.multilabel:
                     model_param = model_param.t()
             elif backbone_type in 'OTEEfficientNet':