Add ported Tensorflow MaxVit weights. Add a few more CLIP ViT fine-tunes. Tweak some model tag names. Improve model tag name sorting. Update HF hub push config layout.

Author: rwightman

timm/__init__.py

@@ -1,4 +1,4 @@
 from .version import __version__
-from .models import create_model, list_models, is_model, list_modules, model_entrypoint, \
+from .models import create_model, list_models, list_pretrained, is_model, list_modules, model_entrypoint, \
     is_scriptable, is_exportable, set_scriptable, set_exportable, \
     is_model_pretrained, get_pretrained_cfg, get_pretrained_cfg_value

timm/models/__init__.py

@@ -70,5 +70,6 @@
 from .layers import convert_splitbn_model, convert_sync_batchnorm
 from .layers import is_scriptable, is_exportable, set_scriptable, set_exportable, is_no_jit, set_no_jit
 from .layers import set_fast_norm
-from .registry import register_model, model_entrypoint, list_models, is_model, list_modules, is_model_in_modules,\
-    is_model_pretrained, get_pretrained_cfg, get_pretrained_cfg_value
+from ._pretrained import PretrainedCfg, filter_pretrained_cfg, generate_default_cfgs, split_model_name_tag
+from .registry import register_model, model_entrypoint, list_models, list_pretrained, is_model, list_modules,\
+    is_model_in_modules, is_model_pretrained, get_pretrained_cfg, get_pretrained_cfg_value

timm/models/_pretrained.py

@@ -1,5 +1,6 @@
+import copy
 from collections import deque, defaultdict
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, field, replace, asdict
 from typing import Any, Deque, Dict, Tuple, Optional, Union
 
 
@@ -8,13 +9,13 @@ class PretrainedCfg:
     """
     """
     # weight locations
-    url: str = ''
-    file: str = ''
-    hf_hub_id: str = ''
-    hf_hub_filename: str = ''
+    url: Optional[Union[str, Tuple[str, str]]] = None
+    file: Optional[str] = None
+    hf_hub_id: Optional[str] = None
+    hf_hub_filename: Optional[str] = None
 
-    source: str = ''  # source of cfg / weight location used (url, file, hf-hub)
-    architecture: str = ''  # architecture variant can be set when not implicit
+    source: Optional[str] = None  # source of cfg / weight location used (url, file, hf-hub)
+    architecture: Optional[str] = None  # architecture variant can be set when not implicit
     custom_load: bool = False  # use custom model specific model.load_pretrained() (ie for npz files)
 
     # input / data config
@@ -31,22 +32,40 @@ class PretrainedCfg:
 
     # head config
     num_classes: int = 1000
-    label_offset: int = 0
+    label_offset: Optional[int] = None
 
     # model attributes that vary with above or required for pretrained adaptation
     pool_size: Optional[Tuple[int, ...]] = None
     test_pool_size: Optional[Tuple[int, ...]] = None
-    first_conv: str = ''
-    classifier: str = ''
+    first_conv: Optional[str] = None
+    classifier: Optional[str] = None
 
-    license: str = ''
-    source_url: str = ''
-    paper: str = ''
-    notes: str = ''
+    license: Optional[str] = None
+    source_url: Optional[str] = None
+    paper: Optional[str] = None
+    notes: Optional[str] = None
 
     @property
     def has_weights(self):
-        return self.url.startswith('http') or self.file or self.hf_hub_id
+        return self.url or self.file or self.hf_hub_id
+
+    def to_dict(self, remove_source=False, remove_null=True):
+        return filter_pretrained_cfg(
+            asdict(self),
+            remove_source=remove_source,
+            remove_null=remove_null
+        )
+
+
+def filter_pretrained_cfg(cfg, remove_source=False, remove_null=True):
+    filtered_cfg = {}
+    for k, v in cfg.items():
+        if remove_source and k in {'url', 'file', 'hf_hub_id', 'hf_hub_id', 'hf_hub_filename', 'source'}:
+            continue
+        if remove_null and v is None:
+            continue
+        filtered_cfg[k] = v
+    return filtered_cfg
 
 
 @dataclass
@@ -71,7 +90,7 @@ def split_model_name_tag(model_name: str, no_tag=''):
     return model_name, tag
 
 
-def generate_defaults(cfgs: Dict[str, Union[Dict[str, Any], PretrainedCfg]]):
+def generate_default_cfgs(cfgs: Dict[str, Union[Dict[str, Any], PretrainedCfg]]):
     out = defaultdict(DefaultCfg)
     default_set = set()  # no tag and tags ending with * are prioritized as default
 
@@ -82,21 +101,22 @@ def generate_defaults(cfgs: Dict[str, Union[Dict[str, Any], PretrainedCfg]]):
 
         model, tag = split_model_name_tag(k)
         is_default_set = model in default_set
-        priority = not tag or (tag.endswith('*') and not is_default_set)
+        priority = (has_weights and not tag) or (tag.endswith('*') and not is_default_set)
         tag = tag.strip('*')
 
         default_cfg = out[model]
-        if has_weights:
-            default_cfg.is_pretrained = True
 
         if priority:
             default_cfg.tags.appendleft(tag)
             default_set.add(model)
-        elif has_weights and not default_set:
+        elif has_weights and not default_cfg.is_pretrained:
             default_cfg.tags.appendleft(tag)
         else:
             default_cfg.tags.append(tag)
 
+        if has_weights:
+            default_cfg.is_pretrained = True
+
         default_cfg.cfgs[tag] = v
 
     return out

timm/models/convnext.py

@@ -21,7 +21,7 @@
 from .helpers import named_apply, build_model_with_cfg, checkpoint_seq
 from .layers import trunc_normal_, SelectAdaptivePool2d, DropPath, ConvMlp, Mlp, LayerNorm2d, LayerNorm, \
     create_conv2d, get_act_layer, make_divisible, to_ntuple
-from ._pretrained import generate_defaults
+from ._pretrained import generate_default_cfgs
 from .registry import register_model
 
 
@@ -373,7 +373,7 @@ def _cfg(url='', **kwargs):
     }
 
 
-default_cfgs = generate_defaults({
+default_cfgs = generate_default_cfgs({
     # timm specific variants
     'convnext_atto.timm_in1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-rsb-weights/convnext_atto_d2-01bb0f51.pth',

timm/models/helpers.py

@@ -575,7 +575,7 @@ def build_model_with_cfg(
     )
 
     # FIXME converting back to dict, PretrainedCfg use should be propagated further, but not into model
-    pretrained_cfg = dataclasses.asdict(pretrained_cfg)
+    pretrained_cfg = pretrained_cfg.to_dict()
 
     _update_default_kwargs(pretrained_cfg, kwargs, kwargs_filter)
 

timm/models/hub.py

@@ -15,11 +15,13 @@
     from torch.hub import _get_torch_home as get_dir
 
 from timm import __version__
+from timm.models._pretrained import filter_pretrained_cfg
 
 try:
-    from huggingface_hub import (create_repo, get_hf_file_metadata,
-                                 hf_hub_download, hf_hub_url,
-                                 repo_type_and_id_from_hf_id, upload_folder)
+    from huggingface_hub import (
+        create_repo, get_hf_file_metadata,
+        hf_hub_download, hf_hub_url,
+        repo_type_and_id_from_hf_id, upload_folder)
     from huggingface_hub.utils import EntryNotFoundError
     hf_hub_download = partial(hf_hub_download, library_name="timm", library_version=__version__)
     _has_hf_hub = True
@@ -46,8 +48,11 @@ def get_cache_dir(child_dir=''):
 
 
 def download_cached_file(url, check_hash=True, progress=False):
-    parts = urlparse(url)
-    filename = os.path.basename(parts.path)
+    if isinstance(url, (list, tuple)):
+        url, filename = url
+    else:
+        parts = urlparse(url)
+        filename = os.path.basename(parts.path)
     cached_file = os.path.join(get_cache_dir(), filename)
     if not os.path.exists(cached_file):
         _logger.info('Downloading: "{}" to {}\n'.format(url, cached_file))
@@ -90,10 +95,27 @@ def _download_from_hf(model_id: str, filename: str):
 def load_model_config_from_hf(model_id: str):
     assert has_hf_hub(True)
     cached_file = _download_from_hf(model_id, 'config.json')
-    pretrained_cfg = load_cfg_from_json(cached_file)
+
+    hf_config = load_cfg_from_json(cached_file)
+    if 'pretrained_cfg' not in hf_config:
+        # old form, pull pretrain_cfg out of the base dict
+        pretrained_cfg = hf_config
+        hf_config = {}
+        hf_config['architecture'] = pretrained_cfg.pop('architecture')
+        hf_config['num_features'] = pretrained_cfg.pop('num_features', None)
+        if 'labels' in pretrained_cfg:
+            hf_config['label_name'] = pretrained_cfg.pop('labels')
+        hf_config['pretrained_cfg'] = pretrained_cfg
+
+    # NOTE currently discarding parent config as only arch name and pretrained_cfg used in timm right now
+    pretrained_cfg = hf_config['pretrained_cfg']
     pretrained_cfg['hf_hub_id'] = model_id  # insert hf_hub id for pretrained weight load during model creation
     pretrained_cfg['source'] = 'hf-hub'
-    model_name = pretrained_cfg.get('architecture')
+    if 'num_classes' in hf_config:
+        # model should be created with parent num_classes if they exist
+        pretrained_cfg['num_classes'] = hf_config['num_classes']
+    model_name = hf_config['architecture']
+
     return pretrained_cfg, model_name
 
 
@@ -114,10 +136,34 @@ def save_for_hf(model, save_directory, model_config=None):
     torch.save(model.state_dict(), weights_path)
 
     config_path = save_directory / 'config.json'
-    hf_config = model.pretrained_cfg
-    hf_config['num_classes'] = model_config.pop('num_classes', model.num_classes)
-    hf_config['num_features'] = model_config.pop('num_features', model.num_features)
-    hf_config['labels'] = model_config.pop('labels', [f"LABEL_{i}" for i in range(hf_config['num_classes'])])
+    hf_config = {}
+    pretrained_cfg = filter_pretrained_cfg(model.pretrained_cfg, remove_source=True, remove_null=True)
+    # set some values at root config level
+    hf_config['architecture'] = pretrained_cfg.pop('architecture')
+    hf_config['num_classes'] = model_config.get('num_classes', model.num_classes)
+    hf_config['num_features'] = model_config.get('num_features', model.num_features)
+    hf_config['global_pool'] = model_config.get('global_pool', getattr(model, 'global_pool', None))
+
+    if 'label' in model_config:
+        _logger.warning(
+            "'label' as a config field for timm models is deprecated. Please use 'label_name' and 'display_name'. "
+            "Using provided 'label' field as 'label_name'.")
+        model_config['label_name'] = model_config.pop('label')
+
+    label_name = model_config.pop('label_name', None)
+    if label_name:
+        assert isinstance(label_name, (dict, list, tuple))
+        # map label id (classifier index) -> unique label name (ie synset for ImageNet, MID for OpenImages)
+        # can be a dict id: name if there are id gaps, or tuple/list if no gaps.
+        hf_config['label_name'] = model_config['label_name']
+
+    display_name = model_config.pop('display_name', None)
+    if display_name:
+        assert isinstance(display_name, dict)
+        # map label_name -> user interface display name
+        hf_config['display_name'] = model_config['display_name']
+
+    hf_config['pretrained_cfg'] = pretrained_cfg
     hf_config.update(model_config)
 
     with config_path.open('w') as f:
@@ -127,14 +173,14 @@ def save_for_hf(model, save_directory, model_config=None):
 def push_to_hf_hub(
     model,
     repo_id: str,
-    commit_message: str ='Add model',
+    commit_message: str = 'Add model',
     token: Optional[str] = None,
     revision: Optional[str] = None,
     private: bool = False,
     create_pr: bool = False,
     model_config: Optional[dict] = None,
 ):
-    # Create repo if doesn't exist yet
+    # Create repo if it doesn't exist yet
     repo_url = create_repo(repo_id, token=token, private=private, exist_ok=True)
 
     # Infer complete repo_id from repo_url
@@ -154,10 +200,11 @@ def push_to_hf_hub(
         # Save model weights and config.
         save_for_hf(model, tmpdir, model_config=model_config)
 
-        # Add readme if does not exist
+        # Add readme if it does not exist
         if not has_readme:
+            model_name = repo_id.split('/')[-1]
             readme_path = Path(tmpdir) / "README.md"
-            readme_text = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {repo_id}'
+            readme_text = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {model_name}'
             readme_path.write_text(readme_text)
 
         # Upload model and return