mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Support other features only modes for EfficientNet

Browse Source
This commit is contained in:
Ross Wightman 2023-06-14 12:57:39 -07:00
parent 2d597b126d
commit a09c88ed0f
2 changed files with 13 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
timm/models/_efficientnet_builder.py
View File

@ -419,9 +419,11 @@ class EfficientNetBuilder:
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)

13
timm/models/efficientnet.py
View File

@ -269,15 +269,20 @@ class EfficientNetFeatures(nn.Module):
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: