Support other features only modes for EfficientNet

reinhud · Jun 14, 2023 · a09c88e · a09c88e
1 parent 2d597b1
commit a09c88e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/timm/models/_efficientnet_builder.py b/timm/models/_efficientnet_builder.py
@@ -419,9 +419,11 @@ def __call__(self, in_chs, model_block_args):
                 if extract_features:
                     feature_info = dict(
                         stage=stack_idx + 1, reduction=current_stride, **block.feature_info(self.feature_location))
-                    module_name = f'blocks.{stack_idx}.{block_idx}'
                     leaf_name = feature_info.get('module', '')
-                    feature_info['module'] = '.'.join([module_name, leaf_name]) if leaf_name else module_name
+                    if leaf_name:
+                        feature_info['module'] = '.'.join([f'blocks.{stack_idx}.{block_idx}', leaf_name])
+                    else:
+                        feature_info['module'] = f'blocks.{stack_idx}'
                     self.features.append(feature_info)
 
                 total_block_idx += 1  # incr global block idx (across all stacks)

diff --git a/timm/models/efficientnet.py b/timm/models/efficientnet.py
@@ -269,15 +269,20 @@ def forward(self, x) -> List[torch.Tensor]:
 
 def _create_effnet(variant, pretrained=False, **kwargs):
     features_only = False
+    features_cls = False
     model_cls = EfficientNet
     kwargs_filter = None
     if kwargs.pop('features_only', False):
-        features_only = True
-        kwargs_filter = ('num_classes', 'num_features', 'head_conv', 'global_pool')
-        model_cls = EfficientNetFeatures
+        if 'feature_cfg' not in kwargs:
+            kwargs_filter = ('num_classes', 'num_features', 'head_conv', 'global_pool')
+            model_cls = EfficientNetFeatures
+            features_cls = True
+        else:
+            features_only = True
     model = build_model_with_cfg(
         model_cls, variant, pretrained,
-        pretrained_strict=not features_only,
+        features_only=features_only,
+        pretrained_strict=not features_cls,
         kwargs_filter=kwargs_filter,
         **kwargs)
     if features_only: