Spaces:
Runtime error
Runtime error
| import datetime | |
| import functools | |
| import os | |
| import torch | |
| import torch.distributed as dist | |
| import timm.models.hub as timm_hub | |
| def setup_for_distributed(is_master): | |
| """ | |
| This function disables printing when not in master process | |
| """ | |
| import builtins as __builtin__ | |
| builtin_print = __builtin__.print | |
| def print(*args, **kwargs): | |
| force = kwargs.pop("force", False) | |
| if is_master or force: | |
| builtin_print(*args, **kwargs) | |
| __builtin__.print = print | |
| def is_dist_avail_and_initialized(): | |
| if not dist.is_available(): | |
| return False | |
| if not dist.is_initialized(): | |
| return False | |
| return True | |
| def get_world_size(): | |
| if not is_dist_avail_and_initialized(): | |
| return 1 | |
| return dist.get_world_size() | |
| def get_rank(): | |
| if not is_dist_avail_and_initialized(): | |
| return 0 | |
| return dist.get_rank() | |
| def is_main_process(): | |
| return get_rank() == 0 | |
| def download_cached_file(url, check_hash=True, progress=False): | |
| """ | |
| Download a file from a URL and cache it locally. If the file already exists, it is not downloaded again. | |
| If distributed, only the main process downloads the file, and the other processes wait for the file to be downloaded. | |
| """ | |
| if not url.startswith('http'): | |
| return url | |
| def get_cached_file_path(): | |
| # a hack to sync the file path across processes | |
| parts = torch.hub.urlparse(url) | |
| filename = os.path.basename(parts.path) | |
| cached_file = os.path.join(timm_hub.get_cache_dir(), filename) | |
| return cached_file | |
| if is_main_process(): | |
| timm_hub.download_cached_file(url, check_hash, progress) | |
| return get_cached_file_path() | |