From e72c98997360a5ee8f271615ba02f8122b45fde0 Mon Sep 17 00:00:00 2001 From: nateraw Date: Wed, 1 Sep 2021 18:14:28 -0600 Subject: [PATCH 1/8] :sparkles: add ability to push to hf hub --- timm/models/hub.py | 120 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 118 insertions(+), 2 deletions(-) diff --git a/timm/models/hub.py b/timm/models/hub.py index 9a9b5530..9e0ac18a 100644 --- a/timm/models/hub.py +++ b/timm/models/hub.py @@ -1,6 +1,7 @@ import json import logging import os +from pathlib import Path from functools import partial from typing import Union, Optional @@ -13,8 +14,7 @@ except ImportError: from timm import __version__ try: - from huggingface_hub import hf_hub_url - from huggingface_hub import cached_download + from huggingface_hub import cached_download, hf_hub_url, HfFolder, HfApi, Repository cached_download = partial(cached_download, library_name="timm", library_version=__version__) except ImportError: hf_hub_url = None @@ -94,3 +94,119 @@ def load_state_dict_from_hf(model_id: str): cached_file = _download_from_hf(model_id, 'pytorch_model.bin') state_dict = torch.load(cached_file, map_location='cpu') return state_dict + + +def save_pretrained_for_hf(model, save_directory, **config_kwargs): + assert has_hf_hub(True) + save_directory = Path(save_directory) + save_directory.mkdir(exist_ok=True, parents=True) + + weights_path = save_directory / 'pytorch_model.bin' + torch.save(model.state_dict(), weights_path) + + config_path = save_directory / 'config.json' + config = model.default_cfg + config.update(config_kwargs) + + with config_path.open('w') as f: + json.dump(config, f, indent=4) + + +def push_to_hf_hub( + model, + repo_path_or_name: Optional[str] = None, + repo_url: Optional[str] = None, + commit_message: Optional[str] = "Add model", + organization: Optional[str] = None, + private: Optional[bool] = None, + api_endpoint: Optional[str] = None, + use_auth_token: Optional[Union[bool, str]] = None, + git_user: Optional[str] = None, + git_email: Optional[str] = None, + config: Optional[dict] = None, +): + """ + Upload model checkpoint and config to the 🤗 Model Hub while synchronizing a local clone of the repo in + :obj:`repo_path_or_name`. + + Parameters: + repo_path_or_name (:obj:`str`, `optional`): + Can either be a repository name for your model or tokenizer in the Hub or a path to a local folder (in + which case the repository will have the name of that local folder). If not specified, will default to + the name given by :obj:`repo_url` and a local directory with that name will be created. + repo_url (:obj:`str`, `optional`): + Specify this in case you want to push to an existing repository in the hub. If unspecified, a new + repository will be created in your namespace (unless you specify an :obj:`organization`) with + :obj:`repo_name`. + commit_message (:obj:`str`, `optional`): + Message to commit while pushing. Will default to :obj:`"add config"`, :obj:`"add tokenizer"` or + :obj:`"add model"` depending on the type of the class. + organization (:obj:`str`, `optional`): + Organization in which you want to push your model or tokenizer (you must be a member of this + organization). + private (:obj:`bool`, `optional`): + Whether or not the repository created should be private (requires a paying subscription). + api_endpoint (:obj:`str`, `optional`): + The API endpoint to use when pushing the model to the hub. + use_auth_token (:obj:`bool` or :obj:`str`, `optional`): + The token to use as HTTP bearer authorization for remote files. If :obj:`True`, will use the token + generated when running :obj:`transformers-cli login` (stored in :obj:`~/.huggingface`). Will default to + :obj:`True` if :obj:`repo_url` is not specified. + git_user (``str``, `optional`): + will override the ``git config user.name`` for committing and pushing files to the hub. + git_email (``str``, `optional`): + will override the ``git config user.email`` for committing and pushing files to the hub. + config (:obj:`dict`, `optional`): + Configuration object to be saved alongside the model weights. + + + Returns: + The url of the commit of your model in the given repository. + """ + assert has_hf_hub(True) + if repo_path_or_name is None and repo_url is None: + raise ValueError( + "You need to specify a `repo_path_or_name` or a `repo_url`." + ) + + if use_auth_token is None and repo_url is None: + token = HfFolder.get_token() + if token is None: + raise ValueError( + "You must login to the Hugging Face hub on this computer by typing `transformers-cli login` and " + "entering your credentials to use `use_auth_token=True`. Alternatively, you can pass your own " + "token as the `use_auth_token` argument." + ) + elif isinstance(use_auth_token, str): + token = use_auth_token + else: + token = None + + if repo_path_or_name is None: + repo_path_or_name = repo_url.split("/")[-1] + + # If no URL is passed and there's no path to a directory containing files, create a repo + if repo_url is None and not os.path.exists(repo_path_or_name): + repo_name = Path(repo_path_or_name).name + repo_url = HfApi(endpoint=api_endpoint).create_repo( + token, + repo_name, + organization=organization, + private=private, + repo_type=None, + exist_ok=True, + ) + + repo = Repository( + repo_path_or_name, + clone_from=repo_url, + use_auth_token=use_auth_token, + git_user=git_user, + git_email=git_email, + ) + repo.git_pull(rebase=True) + + save_config = model.default_cfg + save_config.update(config or {}) + with repo.commit(commit_message): + save_pretrained_for_hf(model, repo.local_dir, **save_config) From 28d2841acfad0dbdd720b9335f361bcc80cde1a8 Mon Sep 17 00:00:00 2001 From: nateraw Date: Wed, 1 Sep 2021 18:15:08 -0600 Subject: [PATCH 2/8] :lipstick: apply isort --- timm/models/hub.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/timm/models/hub.py b/timm/models/hub.py index 9e0ac18a..f5b041a6 100644 --- a/timm/models/hub.py +++ b/timm/models/hub.py @@ -1,20 +1,24 @@ import json import logging import os -from pathlib import Path from functools import partial -from typing import Union, Optional +from pathlib import Path +from typing import Optional, Union import torch -from torch.hub import load_state_dict_from_url, download_url_to_file, urlparse, HASH_REGEX +from torch.hub import (HASH_REGEX, download_url_to_file, + load_state_dict_from_url, urlparse) + try: from torch.hub import get_dir except ImportError: from torch.hub import _get_torch_home as get_dir from timm import __version__ + try: - from huggingface_hub import cached_download, hf_hub_url, HfFolder, HfApi, Repository + from huggingface_hub import (HfApi, HfFolder, Repository, cached_download, + hf_hub_url) cached_download = partial(cached_download, library_name="timm", library_version=__version__) except ImportError: hf_hub_url = None From abf9d51bc3d3b2ca69f569b9f226a3cdf4ec1d10 Mon Sep 17 00:00:00 2001 From: nateraw Date: Tue, 7 Sep 2021 18:39:26 -0600 Subject: [PATCH 3/8] :construction: wip --- timm/models/hub.py | 118 +++++++++++++-------------------------------- 1 file changed, 33 insertions(+), 85 deletions(-) diff --git a/timm/models/hub.py b/timm/models/hub.py index f5b041a6..4cdca7f0 100644 --- a/timm/models/hub.py +++ b/timm/models/hub.py @@ -3,11 +3,10 @@ import logging import os from functools import partial from pathlib import Path -from typing import Optional, Union +from typing import Union import torch -from torch.hub import (HASH_REGEX, download_url_to_file, - load_state_dict_from_url, urlparse) +from torch.hub import HASH_REGEX, download_url_to_file, urlparse, load_state_dict_from_url try: from torch.hub import get_dir @@ -17,8 +16,7 @@ except ImportError: from timm import __version__ try: - from huggingface_hub import (HfApi, HfFolder, Repository, cached_download, - hf_hub_url) + from huggingface_hub import HfApi, HfFolder, Repository, cached_download, hf_hub_url cached_download = partial(cached_download, library_name="timm", library_version=__version__) except ImportError: hf_hub_url = None @@ -110,107 +108,57 @@ def save_pretrained_for_hf(model, save_directory, **config_kwargs): config_path = save_directory / 'config.json' config = model.default_cfg + config['num_classes'] = config_kwargs.pop('num_classes', model.num_classes) + config['num_features'] = config_kwargs.pop('num_features', model.num_features) + config['labels'] = config_kwargs.pop('labels', [f"LABEL_{i}" for i in range(config['num_classes'])]) config.update(config_kwargs) with config_path.open('w') as f: - json.dump(config, f, indent=4) + json.dump(config, f, indent=2) def push_to_hf_hub( model, - repo_path_or_name: Optional[str] = None, - repo_url: Optional[str] = None, - commit_message: Optional[str] = "Add model", - organization: Optional[str] = None, - private: Optional[bool] = None, - api_endpoint: Optional[str] = None, - use_auth_token: Optional[Union[bool, str]] = None, - git_user: Optional[str] = None, - git_email: Optional[str] = None, - config: Optional[dict] = None, + local_dir, + repo_namespace_or_url=None, + commit_message='Add model', + use_auth_token=True, + git_email=None, + git_user=None, + revision=None, + **config_kwargs ): - """ - Upload model checkpoint and config to the 🤗 Model Hub while synchronizing a local clone of the repo in - :obj:`repo_path_or_name`. - - Parameters: - repo_path_or_name (:obj:`str`, `optional`): - Can either be a repository name for your model or tokenizer in the Hub or a path to a local folder (in - which case the repository will have the name of that local folder). If not specified, will default to - the name given by :obj:`repo_url` and a local directory with that name will be created. - repo_url (:obj:`str`, `optional`): - Specify this in case you want to push to an existing repository in the hub. If unspecified, a new - repository will be created in your namespace (unless you specify an :obj:`organization`) with - :obj:`repo_name`. - commit_message (:obj:`str`, `optional`): - Message to commit while pushing. Will default to :obj:`"add config"`, :obj:`"add tokenizer"` or - :obj:`"add model"` depending on the type of the class. - organization (:obj:`str`, `optional`): - Organization in which you want to push your model or tokenizer (you must be a member of this - organization). - private (:obj:`bool`, `optional`): - Whether or not the repository created should be private (requires a paying subscription). - api_endpoint (:obj:`str`, `optional`): - The API endpoint to use when pushing the model to the hub. - use_auth_token (:obj:`bool` or :obj:`str`, `optional`): - The token to use as HTTP bearer authorization for remote files. If :obj:`True`, will use the token - generated when running :obj:`transformers-cli login` (stored in :obj:`~/.huggingface`). Will default to - :obj:`True` if :obj:`repo_url` is not specified. - git_user (``str``, `optional`): - will override the ``git config user.name`` for committing and pushing files to the hub. - git_email (``str``, `optional`): - will override the ``git config user.email`` for committing and pushing files to the hub. - config (:obj:`dict`, `optional`): - Configuration object to be saved alongside the model weights. - - - Returns: - The url of the commit of your model in the given repository. - """ - assert has_hf_hub(True) - if repo_path_or_name is None and repo_url is None: - raise ValueError( - "You need to specify a `repo_path_or_name` or a `repo_url`." - ) - if use_auth_token is None and repo_url is None: - token = HfFolder.get_token() + if repo_namespace_or_url: + repo_owner, repo_name = repo_namespace_or_url.rstrip('/').split('/')[-2:] + else: + if isinstance(use_auth_token, str): + token = use_auth_token + else: + token = HfFolder.get_token() + if token is None: raise ValueError( "You must login to the Hugging Face hub on this computer by typing `transformers-cli login` and " "entering your credentials to use `use_auth_token=True`. Alternatively, you can pass your own " "token as the `use_auth_token` argument." ) - elif isinstance(use_auth_token, str): - token = use_auth_token - else: - token = None - - if repo_path_or_name is None: - repo_path_or_name = repo_url.split("/")[-1] - - # If no URL is passed and there's no path to a directory containing files, create a repo - if repo_url is None and not os.path.exists(repo_path_or_name): - repo_name = Path(repo_path_or_name).name - repo_url = HfApi(endpoint=api_endpoint).create_repo( - token, - repo_name, - organization=organization, - private=private, - repo_type=None, - exist_ok=True, - ) + + repo_owner = HfApi().whoami(token)['name'] + repo_name = Path(local_dir).name + + repo_url = f'https://huggingface.co/{repo_owner}/{repo_name}' repo = Repository( - repo_path_or_name, + local_dir, clone_from=repo_url, use_auth_token=use_auth_token, git_user=git_user, git_email=git_email, + revision=revision, ) - repo.git_pull(rebase=True) - save_config = model.default_cfg - save_config.update(config or {}) with repo.commit(commit_message): - save_pretrained_for_hf(model, repo.local_dir, **save_config) + save_pretrained_for_hf(model, repo.local_dir, **config_kwargs) + + return repo.git_remote_url() From 2b6ade24b3479a641f43aec7bded20bfbb19603e Mon Sep 17 00:00:00 2001 From: nateraw Date: Mon, 13 Sep 2021 23:31:28 -0400 Subject: [PATCH 4/8] :art: write model card to enable inference --- timm/models/hub.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/timm/models/hub.py b/timm/models/hub.py index 4cdca7f0..184e9b85 100644 --- a/timm/models/hub.py +++ b/timm/models/hub.py @@ -159,6 +159,13 @@ def push_to_hf_hub( ) with repo.commit(commit_message): + # Save model weights and config save_pretrained_for_hf(model, repo.local_dir, **config_kwargs) + # Save a model card if it doesn't exist, enabling inference. + readme_path = Path(repo.local_dir) / 'README.md' + readme_txt = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {repo_name}' + if not readme_path.exists(): + readme_path.write_text(readme_txt) + return repo.git_remote_url() From e65a2cba3d6d7bea6b95c4ed482872f4e2355e33 Mon Sep 17 00:00:00 2001 From: nateraw Date: Tue, 14 Sep 2021 01:07:04 -0400 Subject: [PATCH 5/8] :art: cleanup and add a couple comments --- timm/models/hub.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/timm/models/hub.py b/timm/models/hub.py index 184e9b85..31593f88 100644 --- a/timm/models/hub.py +++ b/timm/models/hub.py @@ -7,14 +7,12 @@ from typing import Union import torch from torch.hub import HASH_REGEX, download_url_to_file, urlparse, load_state_dict_from_url - try: from torch.hub import get_dir except ImportError: from torch.hub import _get_torch_home as get_dir from timm import __version__ - try: from huggingface_hub import HfApi, HfFolder, Repository, cached_download, hf_hub_url cached_download = partial(cached_download, library_name="timm", library_version=__version__) @@ -158,14 +156,15 @@ def push_to_hf_hub( revision=revision, ) + # Prepare a default model card that includes the necessary tags to enable inference. + readme_text = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {repo_name}' with repo.commit(commit_message): - # Save model weights and config + # Save model weights and config. save_pretrained_for_hf(model, repo.local_dir, **config_kwargs) - # Save a model card if it doesn't exist, enabling inference. + # Save a model card if it doesn't exist. readme_path = Path(repo.local_dir) / 'README.md' - readme_txt = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {repo_name}' if not readme_path.exists(): - readme_path.write_text(readme_txt) + readme_path.write_text(readme_text) return repo.git_remote_url() From adcb74f87f09cb9db75b62a8690b53ae1552dda5 Mon Sep 17 00:00:00 2001 From: nateraw Date: Tue, 14 Sep 2021 01:11:40 -0400 Subject: [PATCH 6/8] :art: Import load_state_dict_from_url directly --- timm/models/helpers.py | 4 ++-- timm/models/hub.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/timm/models/helpers.py b/timm/models/helpers.py index 662a7a48..281f2412 100644 --- a/timm/models/helpers.py +++ b/timm/models/helpers.py @@ -11,10 +11,10 @@ from typing import Any, Callable, Optional, Tuple import torch import torch.nn as nn - +from torch.hub import load_state_dict_from_url from .features import FeatureListNet, FeatureDictNet, FeatureHookNet -from .hub import has_hf_hub, download_cached_file, load_state_dict_from_hf, load_state_dict_from_url +from .hub import has_hf_hub, download_cached_file, load_state_dict_from_hf from .layers import Conv2dSame, Linear diff --git a/timm/models/hub.py b/timm/models/hub.py index 31593f88..a436aff6 100644 --- a/timm/models/hub.py +++ b/timm/models/hub.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import Union import torch -from torch.hub import HASH_REGEX, download_url_to_file, urlparse, load_state_dict_from_url +from torch.hub import HASH_REGEX, download_url_to_file, urlparse try: from torch.hub import get_dir except ImportError: From b18c9e323b472dbc01e2334d7c1e98e4d0ef4871 Mon Sep 17 00:00:00 2001 From: Nathan Raw Date: Mon, 22 Nov 2021 23:43:44 -0500 Subject: [PATCH 7/8] Update helpers.py --- timm/models/helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/helpers.py b/timm/models/helpers.py index 4e339e66..fd128252 100644 --- a/timm/models/helpers.py +++ b/timm/models/helpers.py @@ -15,7 +15,7 @@ from torch.hub import load_state_dict_from_url from .features import FeatureListNet, FeatureDictNet, FeatureHookNet from .fx_features import FeatureGraphNet -from .hub import has_hf_hub, download_cached_file, load_state_dict_from_hf, load_state_dict_from_url +from .hub import has_hf_hub, download_cached_file, load_state_dict_from_hf from .layers import Conv2dSame, Linear From d633a014e6362ae354f0aa0227c61033233dc260 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Tue, 23 Nov 2021 16:54:01 -0800 Subject: [PATCH 8/8] Post merge cleanup. Fix potential security issue passing kwargs directly through to serialized web data. --- timm/models/helpers.py | 8 ++++---- timm/models/hub.py | 27 ++++++++++++++------------- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/timm/models/helpers.py b/timm/models/helpers.py index fd128252..16ce64d0 100644 --- a/timm/models/helpers.py +++ b/timm/models/helpers.py @@ -184,12 +184,12 @@ def load_pretrained(model, default_cfg=None, num_classes=1000, in_chans=3, filte if not pretrained_url and not hf_hub_id: _logger.warning("No pretrained weights exist for this model. Using random initialization.") return - if hf_hub_id and has_hf_hub(necessary=not pretrained_url): - _logger.info(f'Loading pretrained weights from Hugging Face hub ({hf_hub_id})') - state_dict = load_state_dict_from_hf(hf_hub_id) - else: + if pretrained_url: _logger.info(f'Loading pretrained weights from url ({pretrained_url})') state_dict = load_state_dict_from_url(pretrained_url, progress=progress, map_location='cpu') + elif hf_hub_id and has_hf_hub(necessary=True): + _logger.info(f'Loading pretrained weights from Hugging Face hub ({hf_hub_id})') + state_dict = load_state_dict_from_hf(hf_hub_id) if filter_fn is not None: # for backwards compat with filter fn that take one arg, try one first, the two try: diff --git a/timm/models/hub.py b/timm/models/hub.py index a436aff6..65e7ba9a 100644 --- a/timm/models/hub.py +++ b/timm/models/hub.py @@ -16,9 +16,10 @@ from timm import __version__ try: from huggingface_hub import HfApi, HfFolder, Repository, cached_download, hf_hub_url cached_download = partial(cached_download, library_name="timm", library_version=__version__) + _has_hf_hub = True except ImportError: - hf_hub_url = None cached_download = None + _has_hf_hub = False _logger = logging.getLogger(__name__) @@ -53,11 +54,11 @@ def download_cached_file(url, check_hash=True, progress=False): def has_hf_hub(necessary=False): - if hf_hub_url is None and necessary: + if not _has_hf_hub and necessary: # if no HF Hub module installed and it is necessary to continue, raise error raise RuntimeError( 'Hugging Face hub model specified but package not installed. Run `pip install huggingface_hub`.') - return hf_hub_url is not None + return _has_hf_hub def hf_split(hf_id): @@ -96,8 +97,9 @@ def load_state_dict_from_hf(model_id: str): return state_dict -def save_pretrained_for_hf(model, save_directory, **config_kwargs): +def save_for_hf(model, save_directory, model_config=None): assert has_hf_hub(True) + model_config = model_config or {} save_directory = Path(save_directory) save_directory.mkdir(exist_ok=True, parents=True) @@ -105,14 +107,14 @@ def save_pretrained_for_hf(model, save_directory, **config_kwargs): torch.save(model.state_dict(), weights_path) config_path = save_directory / 'config.json' - config = model.default_cfg - config['num_classes'] = config_kwargs.pop('num_classes', model.num_classes) - config['num_features'] = config_kwargs.pop('num_features', model.num_features) - config['labels'] = config_kwargs.pop('labels', [f"LABEL_{i}" for i in range(config['num_classes'])]) - config.update(config_kwargs) + hf_config = model.default_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.update(model_config) with config_path.open('w') as f: - json.dump(config, f, indent=2) + json.dump(hf_config, f, indent=2) def push_to_hf_hub( @@ -124,9 +126,8 @@ def push_to_hf_hub( git_email=None, git_user=None, revision=None, - **config_kwargs + model_config=None, ): - if repo_namespace_or_url: repo_owner, repo_name = repo_namespace_or_url.rstrip('/').split('/')[-2:] else: @@ -160,7 +161,7 @@ def push_to_hf_hub( readme_text = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {repo_name}' with repo.commit(commit_message): # Save model weights and config. - save_pretrained_for_hf(model, repo.local_dir, **config_kwargs) + save_for_hf(model, repo.local_dir, model_config=model_config) # Save a model card if it doesn't exist. readme_path = Path(repo.local_dir) / 'README.md'