File size: 5,434 Bytes
b66adbc 98e0e3d b66adbc 98e0e3d b66adbc eef302c b66adbc eef302c 1a8429f b66adbc 5f25027 b66adbc 5f25027 b66adbc 5f25027 b66adbc 5f25027 b66adbc 98e0e3d bb19d2f fe028e9 b66adbc 98e0e3d 1a8429f 98e0e3d 1a8429f 98e0e3d 1a8429f 98e0e3d 1a8429f 98e0e3d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import os
import re
import shutil
import sys
import tempfile
import torch
import uuid
import warnings
import hashlib
from urllib.parse import urlparse
from urllib.error import HTTPError, URLError
from urllib.request import urlopen, Request
from typing import Optional
class _Faketqdm: # type: ignore[no-redef]
def __init__(self, total=None, disable=False,
unit=None, *args, **kwargs):
self.total = total
self.disable = disable
self.n = 0
# Ignore all extra *args and **kwargs lest you want to reinvent tqdm
def update(self, n):
if self.disable:
return
self.n += n
if self.total is None:
sys.stderr.write(f"\r{self.n:.1f} bytes")
else:
sys.stderr.write(f"\r{100 * self.n / float(self.total):.1f}%")
sys.stderr.flush()
# Don't bother implementing; use real tqdm if you want
def set_description(self, *args, **kwargs):
pass
def write(self, s):
sys.stderr.write(f"{s}\n")
def close(self):
self.disable = True
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.disable:
return
sys.stderr.write('\n')
try:
from tqdm import tqdm # If tqdm is installed use it, otherwise use the fake wrapper
except ImportError:
tqdm = _Faketqdm
def load_file_from_url(
url: str,
*,
model_dir: str,
progress: bool = True,
file_name: Optional[str] = None,
) -> str:
"""Download a file from `url` into `model_dir`, using the file present if possible.
Returns the path to the downloaded file.
"""
print(f'prepare load {url} to {model_dir}')
os.makedirs(model_dir, exist_ok=True)
print(f'created model_dir : {model_dir}')
if not file_name:
parts = urlparse(url)
print(f'URL parts : {parts}')
file_name = os.path.basename(parts.path)
print(f'file_name : {file_name}')
cached_file = os.path.abspath(os.path.join(model_dir, file_name))
if not os.path.exists(cached_file):
print(f'Downloading: "{url}" to {cached_file}\n')
from torch.hub import download_url_to_file
proxy_download_url_to_file(url, cached_file, progress=progress)
print ('DOWNLOADED FILE: ', url)
print(f'Using cached file: {cached_file}')
return cached_file
def proxy_download_url_to_file(url: str, dst: str, hash_prefix: Optional[str] = None,
progress: bool = True) -> None:
r"""Download object at the given URL to a local path.
Args:
url (str): URL of the object to download
dst (str): Full path where object will be saved, e.g. ``/tmp/temporary_file``
hash_prefix (str, optional): If not None, the SHA256 downloaded file should start with ``hash_prefix``.
Default: None
progress (bool, optional): whether or not to display a progress bar to stderr
Default: True
Example:
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_HUB)
>>> # xdoctest: +REQUIRES(POSIX)
>>> torch.hub.download_url_to_file('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth', '/tmp/temporary_file')
"""
print('PROXY DOWNLOAD')
file_size = None
req = Request(url, headers={"User-Agent": "torch.hub"})
u = urlopen(req)
meta = u.info()
if hasattr(meta, 'getheaders'):
content_length = meta.getheaders("Content-Length")
else:
content_length = meta.get_all("Content-Length")
if content_length is not None and len(content_length) > 0:
file_size = int(content_length[0])
# We deliberately save it in a temp file and move it after
# download is complete. This prevents a local working checkpoint
# being overridden by a broken download.
# We deliberately do not use NamedTemporaryFile to avoid restrictive
# file permissions being applied to the downloaded file.
dst = os.path.expanduser(dst)
print(f'PROXY DOWNLOAD: {dst}')
for seq in range(tempfile.TMP_MAX):
tmp_dst = dst + '.' + uuid.uuid4().hex + '.partial'
try:
f = open(tmp_dst, 'w+b')
except FileExistsError:
continue
break
else:
raise FileExistsError(errno.EEXIST, 'No usable temporary file name found')
try:
if hash_prefix is not None:
sha256 = hashlib.sha256()
with tqdm(total=file_size, disable=not progress,
unit='B', unit_scale=True, unit_divisor=1024) as pbar:
while True:
buffer = u.read(8192)
if len(buffer) == 0:
break
f.write(buffer)
if hash_prefix is not None:
sha256.update(buffer)
pbar.update(len(buffer))
f.close()
print(f'PROXY DOWNLOAD: closed file {f.name}')
if hash_prefix is not None:
digest = sha256.hexdigest()
if digest[:len(hash_prefix)] != hash_prefix:
raise RuntimeError(f'invalid hash value (expected "{hash_prefix}", got "{digest}")')
shutil.move(f.name, dst)
print(f'PROXY DOWNLOAD: moved file {f.name} to {dst}')
finally:
print(f'PROXY DOWNLOAD: finally closing {f.name}')
f.close()
if os.path.exists(f.name):
os.remove(f.name)
|