You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
pytorch-image-models/timm/models/hub.py

172 lines
5.8 KiB

import json
import logging
import os
from functools import partial
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
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__)
except ImportError:
hf_hub_url = None
cached_download = None
_logger = logging.getLogger(__name__)
def get_cache_dir(child_dir=''):
"""
Returns the location of the directory where models are cached (and creates it if necessary).
"""
# Issue warning to move data if old env is set
if os.getenv('TORCH_MODEL_ZOO'):
_logger.warning('TORCH_MODEL_ZOO is deprecated, please use env TORCH_HOME instead')
hub_dir = get_dir()
child_dir = () if not child_dir else (child_dir,)
model_dir = os.path.join(hub_dir, 'checkpoints', *child_dir)
os.makedirs(model_dir, exist_ok=True)
return model_dir
def download_cached_file(url, check_hash=True, progress=False):
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))
hash_prefix = None
if check_hash:
r = HASH_REGEX.search(filename) # r is Optional[Match[str]]
hash_prefix = r.group(1) if r else None
download_url_to_file(url, cached_file, hash_prefix, progress=progress)
return cached_file
def has_hf_hub(necessary=False):
if hf_hub_url is None 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
def hf_split(hf_id):
rev_split = hf_id.split('@')
assert 0 < len(rev_split) <= 2, 'hf_hub id should only contain one @ character to identify revision.'
hf_model_id = rev_split[0]
hf_revision = rev_split[-1] if len(rev_split) > 1 else None
return hf_model_id, hf_revision
def load_cfg_from_json(json_file: Union[str, os.PathLike]):
with open(json_file, "r", encoding="utf-8") as reader:
text = reader.read()
return json.loads(text)
def _download_from_hf(model_id: str, filename: str):
hf_model_id, hf_revision = hf_split(model_id)
url = hf_hub_url(hf_model_id, filename, revision=hf_revision)
return cached_download(url, cache_dir=get_cache_dir('hf'))
def load_model_config_from_hf(model_id: str):
assert has_hf_hub(True)
cached_file = _download_from_hf(model_id, 'config.json')
default_cfg = load_cfg_from_json(cached_file)
default_cfg['hf_hub'] = model_id # insert hf_hub id for pretrained weight load during model creation
model_name = default_cfg.get('architecture')
return default_cfg, model_name
def load_state_dict_from_hf(model_id: str):
assert has_hf_hub(True)
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['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=2)
def push_to_hf_hub(
model,
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
):
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."
)
repo_owner = HfApi().whoami(token)['name']
repo_name = Path(local_dir).name
repo_url = f'https://huggingface.co/{repo_owner}/{repo_name}'
repo = Repository(
local_dir,
clone_from=repo_url,
use_auth_token=use_auth_token,
git_user=git_user,
git_email=git_email,
revision=revision,
)
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()