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

Merge pull request #2142 from huggingface/fix_hrnet_head

Browse Source 
Fix #2139, disable strict weight loading when head changes from classification
This commit is contained in:
Ross Wightman 2024-04-09 10:04:09 -07:00 committed by GitHub
commit 9531eb793c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
timm/models/hrnet.py
View File

@ -862,12 +862,17 @@ def _create_hrnet(variant, pretrained=False, cfg_variant=None, **model_kwargs):
kwargs_filter = ('num_classes', 'global_pool') kwargs_filter = ('num_classes', 'global_pool')
features_only = True features_only = True
cfg_variant = cfg_variant or variant cfg_variant = cfg_variant or variant
pretrained_strict = model_kwargs.pop(
'pretrained_strict',
not features_only and model_kwargs.get('head', 'classification') == 'classification'
)
model = build_model_with_cfg( model = build_model_with_cfg(
model_cls, model_cls,
variant, variant,
pretrained, pretrained,
model_cfg=cfg_cls[cfg_variant], model_cfg=cfg_cls[cfg_variant],
pretrained_strict=not features_only, pretrained_strict=pretrained_strict,
kwargs_filter=kwargs_filter, kwargs_filter=kwargs_filter,
**model_kwargs, **model_kwargs,
) )