| import difflib |
| import os |
| import io |
| import shutil |
| import struct |
| import sys |
| import torch |
| import tarfile |
| import tempfile |
| import warnings |
| from contextlib import closing, contextmanager |
| from ._utils import _import_dotted_name |
| from ._six import string_classes as _string_classes |
| from torch._utils_internal import get_source_lines_and_file |
| from torch.types import Storage |
| from typing import Any, BinaryIO, cast, Dict, Optional, Type, Tuple, Union, IO |
| import copyreg |
| import pickle |
| import pathlib |
| |
| DEFAULT_PROTOCOL = 2 |
| |
| LONG_SIZE = struct.Struct('=l').size |
| INT_SIZE = struct.Struct('=i').size |
| SHORT_SIZE = struct.Struct('=h').size |
| |
| MAGIC_NUMBER = 0x1950a86a20f9469cfc6c |
| PROTOCOL_VERSION = 1001 |
| STORAGE_KEY_SEPARATOR = ',' |
| |
| class SourceChangeWarning(Warning): |
| pass |
| |
| |
| @contextmanager |
| def mkdtemp(): |
| path = tempfile.mkdtemp() |
| yield path |
| shutil.rmtree(path) |
| |
| |
| _package_registry = [] |
| |
| |
| def _is_zipfile(f) -> bool: |
| # This is a stricter implementation than zipfile.is_zipfile(). |
| # zipfile.is_zipfile() is True if the magic number appears anywhere in the |
| # binary. Since we expect the files here to be generated by torch.save or |
| # torch.jit.save, it's safe to only check the start bytes and avoid |
| # collisions and assume the zip has only 1 file. |
| # See bugs.python.org/issue28494. |
| |
| # Read the first 4 bytes of the file |
| read_bytes = [] |
| start = f.tell() |
| |
| byte = f.read(1) |
| while byte != "": |
| read_bytes.append(byte) |
| if len(read_bytes) == 4: |
| break |
| byte = f.read(1) |
| f.seek(start) |
| |
| local_header_magic_number = [b'P', b'K', b'\x03', b'\x04'] |
| return read_bytes == local_header_magic_number |
| |
| |
| def register_package(priority, tagger, deserializer): |
| queue_elem = (priority, tagger, deserializer) |
| _package_registry.append(queue_elem) |
| _package_registry.sort() |
| |
| |
| def check_module_version_greater_or_equal(module, req_version_tuple, error_if_malformed=True): |
| ''' |
| Check if a module's version satisfies requirements |
| |
| Usually, a module's version string will be like 'x.y.z', which would be represented |
| as a tuple (x, y, z), but sometimes it could be an unexpected format. If the version |
| string does not match the given tuple's format up to the length of the tuple, then |
| error and exit or emit a warning. |
| |
| Args: |
| module: the module to check the version of |
| req_version_tuple: tuple (usually of ints) representing the required version |
| error_if_malformed: whether we should exit if module version string is malformed |
| |
| Returns: |
| requirement_is_met: bool |
| ''' |
| try: |
| version_strs = module.__version__.split('.') |
| # Cast module version fields to match the types of the required version |
| module_version = tuple( |
| type(req_field)(version_strs[idx]) for idx, req_field in enumerate(req_version_tuple) |
| ) |
| requirement_is_met = module_version >= req_version_tuple |
| |
| except Exception as e: |
| message = ( |
| "'%s' module version string is malformed '%s' and cannot be compared" |
| " with tuple %s" |
| ) % ( |
| module.__name__, module.__version__, str(req_version_tuple) |
| ) |
| if error_if_malformed: |
| raise RuntimeError(message) from e |
| else: |
| warnings.warn(message + ', but continuing assuming that requirement is met') |
| requirement_is_met = True |
| |
| return requirement_is_met |
| |
| |
| def _cpu_tag(obj): |
| if type(obj).__module__ == 'torch': |
| return 'cpu' |
| |
| |
| def _cuda_tag(obj): |
| if type(obj).__module__ == 'torch.cuda': |
| return 'cuda:' + str(obj.get_device()) |
| |
| |
| def _cpu_deserialize(obj, location): |
| if location == 'cpu': |
| return obj |
| |
| |
| def validate_cuda_device(location): |
| device = torch.cuda._utils._get_device_index(location, True) |
| |
| if not torch.cuda.is_available(): |
| raise RuntimeError('Attempting to deserialize object on a CUDA ' |
| 'device but torch.cuda.is_available() is False. ' |
| 'If you are running on a CPU-only machine, ' |
| 'please use torch.load with map_location=torch.device(\'cpu\') ' |
| 'to map your storages to the CPU.') |
| device_count = torch.cuda.device_count() |
| if device >= device_count: |
| raise RuntimeError('Attempting to deserialize object on CUDA device ' |
| f'{device} but torch.cuda.device_count() is {device_count}. Please use ' |
| 'torch.load with map_location to map your storages ' |
| 'to an existing device.') |
| return device |
| |
| |
| def _cuda_deserialize(obj, location): |
| if location.startswith('cuda'): |
| device = validate_cuda_device(location) |
| if getattr(obj, "_torch_load_uninitialized", False): |
| storage_type = getattr(torch.cuda, type(obj).__name__) |
| with torch.cuda.device(device): |
| return storage_type(obj.size()) |
| else: |
| return obj.cuda(device) |
| |
| |
| register_package(10, _cpu_tag, _cpu_deserialize) |
| register_package(20, _cuda_tag, _cuda_deserialize) |
| |
| |
| def location_tag(storage: Storage): |
| for _, tagger, _ in _package_registry: |
| location = tagger(storage) |
| if location: |
| return location |
| raise RuntimeError("don't know how to determine data location of " |
| + torch.typename(storage)) |
| |
| |
| def default_restore_location(storage, location): |
| for _, _, fn in _package_registry: |
| result = fn(storage, location) |
| if result is not None: |
| return result |
| raise RuntimeError("don't know how to restore data location of " |
| + torch.typename(storage) + " (tagged with " |
| + location + ")") |
| |
| |
| def normalize_storage_type(storage_type): |
| return getattr(torch, storage_type.__name__) |
| |
| |
| def storage_to_tensor_type(storage): |
| storage_type = type(storage) |
| module = _import_dotted_name(storage_type.__module__) |
| return getattr(module, storage_type.__name__.replace('Storage', 'Tensor')) |
| |
| |
| def _is_path(name_or_buffer): |
| return isinstance(name_or_buffer, str) or \ |
| isinstance(name_or_buffer, pathlib.Path) |
| |
| |
| class _opener(object): |
| def __init__(self, file_like): |
| self.file_like = file_like |
| |
| def __enter__(self): |
| return self.file_like |
| |
| def __exit__(self, *args): |
| pass |
| |
| |
| class _open_file(_opener): |
| def __init__(self, name, mode): |
| super(_open_file, self).__init__(open(name, mode)) |
| |
| def __exit__(self, *args): |
| self.file_like.close() |
| |
| |
| class _open_buffer_reader(_opener): |
| def __init__(self, buffer): |
| super(_open_buffer_reader, self).__init__(buffer) |
| _check_seekable(buffer) |
| |
| |
| class _open_buffer_writer(_opener): |
| def __exit__(self, *args): |
| self.file_like.flush() |
| |
| |
| def _open_file_like(name_or_buffer, mode): |
| if _is_path(name_or_buffer): |
| return _open_file(name_or_buffer, mode) |
| else: |
| if 'w' in mode: |
| return _open_buffer_writer(name_or_buffer) |
| elif 'r' in mode: |
| return _open_buffer_reader(name_or_buffer) |
| else: |
| raise RuntimeError(f"Expected 'r' or 'w' in mode but got {mode}") |
| |
| |
| class _open_zipfile_reader(_opener): |
| def __init__(self, name_or_buffer) -> None: |
| super(_open_zipfile_reader, self).__init__(torch._C.PyTorchFileReader(name_or_buffer)) |
| |
| |
| class _open_zipfile_writer_file(_opener): |
| def __init__(self, name) -> None: |
| super(_open_zipfile_writer_file, self).__init__(torch._C.PyTorchFileWriter(str(name))) |
| |
| def __exit__(self, *args) -> None: |
| self.file_like.write_end_of_file() |
| |
| |
| class _open_zipfile_writer_buffer(_opener): |
| def __init__(self, buffer) -> None: |
| self.buffer = buffer |
| super(_open_zipfile_writer_buffer, self).__init__(torch._C.PyTorchFileWriter(buffer)) |
| |
| def __exit__(self, *args) -> None: |
| self.file_like.write_end_of_file() |
| self.buffer.flush() |
| |
| |
| def _open_zipfile_writer(name_or_buffer): |
| container: Type[_opener] |
| if _is_path(name_or_buffer): |
| container = _open_zipfile_writer_file |
| else: |
| container = _open_zipfile_writer_buffer |
| return container(name_or_buffer) |
| |
| |
| def _is_compressed_file(f) -> bool: |
| compress_modules = ['gzip'] |
| try: |
| return f.__module__ in compress_modules |
| except AttributeError: |
| return False |
| |
| |
| def _should_read_directly(f): |
| """ |
| Checks if f is a file that should be read directly. It should be read |
| directly if it is backed by a real file (has a fileno) and is not a |
| a compressed file (e.g. gzip) |
| """ |
| if _is_compressed_file(f): |
| return False |
| try: |
| return f.fileno() >= 0 |
| except io.UnsupportedOperation: |
| return False |
| except AttributeError: |
| return False |
| |
| |
| def _check_seekable(f) -> bool: |
| |
| def raise_err_msg(patterns, e): |
| for p in patterns: |
| if p in str(e): |
| msg = (str(e) + ". You can only torch.load from a file that is seekable." |
| + " Please pre-load the data into a buffer like io.BytesIO and" |
| + " try to load from it instead.") |
| raise type(e)(msg) |
| raise e |
| |
| try: |
| f.seek(f.tell()) |
| return True |
| except (io.UnsupportedOperation, AttributeError) as e: |
| raise_err_msg(["seek", "tell"], e) |
| return False |
| |
| def _check_dill_version(pickle_module) -> None: |
| '''Checks if using dill as the pickle module, and if so, checks if it is the correct version. |
| If dill version is lower than 0.3.1, a ValueError is raised. |
| |
| Args: |
| pickle_module: module used for pickling metadata and objects |
| |
| ''' |
| if pickle_module.__name__ == 'dill': |
| required_dill_version = (0, 3, 1) |
| if not check_module_version_greater_or_equal(pickle_module, required_dill_version, False): |
| raise ValueError(( |
| "'torch' supports dill >= %s, but you have dill %s." |
| " Please upgrade dill or switch to 'pickle'" |
| ) % ( |
| '.'.join([str(num) for num in required_dill_version]), |
| pickle_module.__version__ |
| )) |
| |
| def save(obj, f: Union[str, os.PathLike, BinaryIO, IO[bytes]], |
| pickle_module=pickle, pickle_protocol=DEFAULT_PROTOCOL, _use_new_zipfile_serialization=True) -> None: |
| """Saves an object to a disk file. |
| |
| See also: `saving-loading-tensors` |
| |
| Args: |
| obj: saved object |
| f: a file-like object (has to implement write and flush) or a string or |
| os.PathLike object containing a file name |
| pickle_module: module used for pickling metadata and objects |
| pickle_protocol: can be specified to override the default protocol |
| |
| .. note:: |
| A common PyTorch convention is to save tensors using .pt file extension. |
| |
| .. note:: |
| PyTorch preserves storage sharing across serialization. See |
| `preserve-storage-sharing` for more details. |
| |
| .. note:: |
| The 1.6 release of PyTorch switched ``torch.save`` to use a new |
| zipfile-based file format. ``torch.load`` still retains the ability to |
| load files in the old format. If for any reason you want ``torch.save`` |
| to use the old format, pass the kwarg ``_use_new_zipfile_serialization=False``. |
| |
| Example: |
| >>> # Save to file |
| >>> x = torch.tensor([0, 1, 2, 3, 4]) |
| >>> torch.save(x, 'tensor.pt') |
| >>> # Save to io.BytesIO buffer |
| >>> buffer = io.BytesIO() |
| >>> torch.save(x, buffer) |
| """ |
| _check_dill_version(pickle_module) |
| |
| with _open_file_like(f, 'wb') as opened_file: |
| if _use_new_zipfile_serialization: |
| with _open_zipfile_writer(opened_file) as opened_zipfile: |
| _save(obj, opened_zipfile, pickle_module, pickle_protocol) |
| return |
| _legacy_save(obj, opened_file, pickle_module, pickle_protocol) |
| |
| |
| def _legacy_save(obj, f, pickle_module, pickle_protocol) -> None: |
| import torch.nn as nn |
| serialized_container_types = {} |
| serialized_storages = {} |
| |
| def persistent_id(obj: Any) -> Optional[Tuple]: |
| # FIXME: the docs say that persistent_id should only return a string |
| # but torch store returns tuples. This works only in the binary protocol |
| # see |
| # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects |
| # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537 |
| if isinstance(obj, type) and issubclass(obj, nn.Module): |
| if obj in serialized_container_types: |
| return None |
| serialized_container_types[obj] = True |
| source_file = source = None |
| try: |
| source_lines, _, source_file = get_source_lines_and_file(obj) |
| source = ''.join(source_lines) |
| except Exception: # saving the source is optional, so we can ignore any errors |
| warnings.warn("Couldn't retrieve source code for container of " |
| "type " + obj.__name__ + ". It won't be checked " |
| "for correctness upon loading.") |
| return ('module', obj, source_file, source) |
| |
| elif torch.is_storage(obj): |
| view_metadata: Optional[Tuple[str, int, int]] |
| obj = cast(Storage, obj) |
| storage_type = normalize_storage_type(type(obj)) |
| # Offset is always 0, but we keep it for backwards compatibility |
| # with the old serialization format (which supported storage views) |
| offset = 0 |
| obj_key = str(obj._cdata) |
| location = location_tag(obj) |
| serialized_storages[obj_key] = obj |
| is_view = obj._cdata != obj._cdata |
| if is_view: |
| view_metadata = (str(obj._cdata), offset, obj.size()) |
| else: |
| view_metadata = None |
| |
| return ('storage', |
| storage_type, |
| obj_key, |
| location, |
| obj.size(), |
| view_metadata) |
| return None |
| |
| sys_info = dict( |
| protocol_version=PROTOCOL_VERSION, |
| little_endian=sys.byteorder == 'little', |
| type_sizes=dict( |
| short=SHORT_SIZE, |
| int=INT_SIZE, |
| long=LONG_SIZE, |
| ), |
| ) |
| |
| pickle_module.dump(MAGIC_NUMBER, f, protocol=pickle_protocol) |
| pickle_module.dump(PROTOCOL_VERSION, f, protocol=pickle_protocol) |
| pickle_module.dump(sys_info, f, protocol=pickle_protocol) |
| pickler = pickle_module.Pickler(f, protocol=pickle_protocol) |
| pickler.persistent_id = persistent_id |
| pickler.dump(obj) |
| |
| serialized_storage_keys = sorted(serialized_storages.keys()) |
| pickle_module.dump(serialized_storage_keys, f, protocol=pickle_protocol) |
| f.flush() |
| for key in serialized_storage_keys: |
| serialized_storages[key]._write_file(f, _should_read_directly(f), True) |
| |
| |
| def _save(obj, zip_file, pickle_module, pickle_protocol): |
| serialized_storages = {} |
| |
| def persistent_id(obj): |
| # FIXME: the docs say that persistent_id should only return a string |
| # but torch store returns tuples. This works only in the binary protocol |
| # see |
| # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects |
| # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537 |
| if torch.is_storage(obj): |
| storage_type = normalize_storage_type(type(obj)) |
| obj_key = str(obj._cdata) |
| location = location_tag(obj) |
| serialized_storages[obj_key] = obj |
| |
| return ('storage', |
| storage_type, |
| obj_key, |
| location, |
| obj.size()) |
| return None |
| |
| # Write the pickle data for `obj` |
| data_buf = io.BytesIO() |
| pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol) |
| pickler.persistent_id = persistent_id |
| pickler.dump(obj) |
| data_value = data_buf.getvalue() |
| zip_file.write_record('data.pkl', data_value, len(data_value)) |
| |
| # Write each tensor to a file named tensor/the_tensor_key in the zip archive |
| for key in sorted(serialized_storages.keys()): |
| name = f'data/{key}' |
| storage = serialized_storages[key] |
| # given that we copy things around anyway, we might use storage.cpu() |
| # this means to that to get tensors serialized, you need to implement |
| # .cpu() on the underlying Storage |
| if storage.device.type != 'cpu': |
| storage = storage.cpu() |
| # Now that it is on the CPU we can directly copy it into the zip file |
| num_bytes = storage.size() * storage.element_size() |
| zip_file.write_record(name, storage.data_ptr(), num_bytes) |
| |
| |
| def load(f, map_location=None, pickle_module=pickle, **pickle_load_args): |
| """Loads an object saved with :func:`torch.save` from a file. |
| |
| :func:`torch.load` uses Python's unpickling facilities but treats storages, |
| which underlie tensors, specially. They are first deserialized on the |
| CPU and are then moved to the device they were saved from. If this fails |
| (e.g. because the run time system doesn't have certain devices), an exception |
| is raised. However, storages can be dynamically remapped to an alternative |
| set of devices using the :attr:`map_location` argument. |
| |
| If :attr:`map_location` is a callable, it will be called once for each serialized |
| storage with two arguments: storage and location. The storage argument |
| will be the initial deserialization of the storage, residing on the CPU. |
| Each serialized storage has a location tag associated with it which |
| identifies the device it was saved from, and this tag is the second |
| argument passed to :attr:`map_location`. The builtin location tags are ``'cpu'`` |
| for CPU tensors and ``'cuda:device_id'`` (e.g. ``'cuda:2'``) for CUDA tensors. |
| :attr:`map_location` should return either ``None`` or a storage. If |
| :attr:`map_location` returns a storage, it will be used as the final deserialized |
| object, already moved to the right device. Otherwise, :func:`torch.load` will |
| fall back to the default behavior, as if :attr:`map_location` wasn't specified. |
| |
| If :attr:`map_location` is a :class:`torch.device` object or a string containing |
| a device tag, it indicates the location where all tensors should be loaded. |
| |
| Otherwise, if :attr:`map_location` is a dict, it will be used to remap location tags |
| appearing in the file (keys), to ones that specify where to put the |
| storages (values). |
| |
| User extensions can register their own location tags and tagging and |
| deserialization methods using :func:`torch.serialization.register_package`. |
| |
| Args: |
| f: a file-like object (has to implement :meth:`read`, :meth:`readline`, :meth:`tell`, and :meth:`seek`), |
| or a string or os.PathLike object containing a file name |
| map_location: a function, :class:`torch.device`, string or a dict specifying how to remap storage |
| locations |
| pickle_module: module used for unpickling metadata and objects (has to |
| match the :attr:`pickle_module` used to serialize file) |
| pickle_load_args: (Python 3 only) optional keyword arguments passed over to |
| :func:`pickle_module.load` and :func:`pickle_module.Unpickler`, e.g., |
| :attr:`errors=...`. |
| |
| .. warning:: |
| :func:`torch.load()` uses ``pickle`` module implicitly, which is known to be insecure. |
| It is possible to construct malicious pickle data which will execute arbitrary code |
| during unpickling. Never load data that could have come from an untrusted |
| source, or that could have been tampered with. **Only load data you trust**. |
| |
| .. note:: |
| When you call :func:`torch.load()` on a file which contains GPU tensors, those tensors |
| will be loaded to GPU by default. You can call ``torch.load(.., map_location='cpu')`` |
| and then :meth:`load_state_dict` to avoid GPU RAM surge when loading a model checkpoint. |
| |
| .. note:: |
| By default, we decode byte strings as ``utf-8``. This is to avoid a common error |
| case ``UnicodeDecodeError: 'ascii' codec can't decode byte 0x...`` |
| when loading files saved by Python 2 in Python 3. If this default |
| is incorrect, you may use an extra :attr:`encoding` keyword argument to specify how |
| these objects should be loaded, e.g., :attr:`encoding='latin1'` decodes them |
| to strings using ``latin1`` encoding, and :attr:`encoding='bytes'` keeps them |
| as byte arrays which can be decoded later with ``byte_array.decode(...)``. |
| |
| Example: |
| >>> torch.load('tensors.pt') |
| # Load all tensors onto the CPU |
| >>> torch.load('tensors.pt', map_location=torch.device('cpu')) |
| # Load all tensors onto the CPU, using a function |
| >>> torch.load('tensors.pt', map_location=lambda storage, loc: storage) |
| # Load all tensors onto GPU 1 |
| >>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1)) |
| # Map tensors from GPU 1 to GPU 0 |
| >>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'}) |
| # Load tensor from io.BytesIO object |
| >>> with open('tensor.pt', 'rb') as f: |
| ... buffer = io.BytesIO(f.read()) |
| >>> torch.load(buffer) |
| # Load a module with 'ascii' encoding for unpickling |
| >>> torch.load('module.pt', encoding='ascii') |
| """ |
| _check_dill_version(pickle_module) |
| |
| if 'encoding' not in pickle_load_args.keys(): |
| pickle_load_args['encoding'] = 'utf-8' |
| |
| with _open_file_like(f, 'rb') as opened_file: |
| if _is_zipfile(opened_file): |
| # The zipfile reader is going to advance the current file position. |
| # If we want to actually tail call to torch.jit.load, we need to |
| # reset back to the original position. |
| orig_position = opened_file.tell() |
| with _open_zipfile_reader(opened_file) as opened_zipfile: |
| if _is_torchscript_zip(opened_zipfile): |
| warnings.warn("'torch.load' received a zip file that looks like a TorchScript archive" |
| " dispatching to 'torch.jit.load' (call 'torch.jit.load' directly to" |
| " silence this warning)", UserWarning) |
| opened_file.seek(orig_position) |
| return torch.jit.load(opened_file) |
| return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args) |
| return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args) |
| |
| |
| # Register pickling support for layout instances such as |
| # torch.sparse_coo, etc |
| def _get_layout(name): |
| """Get layout extension object from its string representation. |
| """ |
| cache = _get_layout.cache # type: ignore[attr-defined] |
| if not cache: |
| for v in torch.__dict__.values(): |
| if isinstance(v, torch.layout): |
| cache[str(v)] = v |
| return cache[name] |
| |
| # There are yet not good way to type annotate function attributes https://github.com/python/mypy/issues/2087 |
| _get_layout.cache = {} # type: ignore[attr-defined] |
| copyreg.pickle(torch.layout, lambda obj: (_get_layout, (str(obj),))) |
| |
| |
| def _legacy_load(f, map_location, pickle_module, **pickle_load_args): |
| deserialized_objects: Dict[int, Any] = {} |
| |
| restore_location = _get_restore_location(map_location) |
| |
| def _check_container_source(container_type, source_file, original_source): |
| try: |
| current_source = ''.join(get_source_lines_and_file(container_type)[0]) |
| except Exception: # saving the source is optional, so we can ignore any errors |
| warnings.warn("Couldn't retrieve source code for container of " |
| "type " + container_type.__name__ + ". It won't be checked " |
| "for correctness upon loading.") |
| return |
| if original_source != current_source: |
| if container_type.dump_patches: |
| file_name = container_type.__name__ + '.patch' |
| diff = difflib.unified_diff(current_source.split('\n'), |
| original_source.split('\n'), |
| source_file, |
| source_file, lineterm="") |
| lines = '\n'.join(diff) |
| try: |
| with open(file_name, 'a+') as f: |
| file_size = f.seek(0, 2) |
| f.seek(0) |
| if file_size == 0: |
| f.write(lines) |
| elif file_size != len(lines) or f.read() != lines: |
| raise IOError |
| msg = ("Saved a reverse patch to " + file_name + ". " |
| "Run `patch -p0 < " + file_name + "` to revert your " |
| "changes.") |
| except IOError: |
| msg = ("Tried to save a patch, but couldn't create a " |
| "writable file " + file_name + ". Make sure it " |
| "doesn't exist and your working directory is " |
| "writable.") |
| else: |
| msg = ("you can retrieve the original source code by " |
| "accessing the object's source attribute or set " |
| "`torch.nn.Module.dump_patches = True` and use the " |
| "patch tool to revert the changes.") |
| msg = f"source code of class '{torch.typename(container_type)}' has changed. {msg}" |
| warnings.warn(msg, SourceChangeWarning) |
| |
| def legacy_load(f): |
| deserialized_objects: Dict[int, Any] = {} |
| |
| def persistent_load(saved_id): |
| if isinstance(saved_id, tuple): |
| # Ignore containers that don't have any sources saved |
| if all(saved_id[1:]): |
| _check_container_source(*saved_id) |
| return saved_id[0] |
| return deserialized_objects[int(saved_id)] |
| |
| with closing(tarfile.open(fileobj=f, mode='r:', format=tarfile.PAX_FORMAT)) as tar, \ |
| mkdtemp() as tmpdir: |
| |
| tar.extract('storages', path=tmpdir) |
| with open(os.path.join(tmpdir, 'storages'), 'rb', 0) as f: |
| num_storages = pickle_module.load(f, **pickle_load_args) |
| for i in range(num_storages): |
| args = pickle_module.load(f, **pickle_load_args) |
| key, location, storage_type = args |
| obj = storage_type._new_with_file(f) |
| obj = restore_location(obj, location) |
| deserialized_objects[key] = obj |
| |
| storage_views = pickle_module.load(f, **pickle_load_args) |
| for target_cdata, root_cdata, offset, size in storage_views: |
| root = deserialized_objects[root_cdata] |
| deserialized_objects[target_cdata] = root[offset:offset + size] |
| |
| tar.extract('tensors', path=tmpdir) |
| with open(os.path.join(tmpdir, 'tensors'), 'rb', 0) as f: |
| num_tensors = pickle_module.load(f, **pickle_load_args) |
| for _ in range(num_tensors): |
| args = pickle_module.load(f, **pickle_load_args) |
| key, storage_id, original_tensor_type = args |
| storage = deserialized_objects[storage_id] |
| tensor_type = storage_to_tensor_type(storage) |
| ndim, = struct.unpack('<i', f.read(4)) |
| # skip next 4 bytes; legacy encoding treated ndim as 8 bytes |
| f.read(4) |
| size = struct.unpack(f'<{ndim}q', f.read(8 * ndim)) |
| stride = struct.unpack(f'<{ndim}q', f.read(8 * ndim)) |
| storage_offset, = struct.unpack('<q', f.read(8)) |
| tensor = tensor_type().set_(storage, storage_offset, size, stride) |
| deserialized_objects[key] = tensor |
| |
| pickle_file = tar.extractfile('pickle') |
| unpickler = pickle_module.Unpickler(pickle_file, **pickle_load_args) |
| unpickler.persistent_load = persistent_load |
| result = unpickler.load() |
| return result |
| |
| deserialized_objects = {} |
| |
| def persistent_load(saved_id): |
| assert isinstance(saved_id, tuple) |
| typename = _maybe_decode_ascii(saved_id[0]) |
| data = saved_id[1:] |
| |
| if typename == 'module': |
| # Ignore containers that don't have any sources saved |
| if all(data[1:]): |
| _check_container_source(*data) |
| return data[0] |
| elif typename == 'storage': |
| data_type, root_key, location, size, view_metadata = data |
| location = _maybe_decode_ascii(location) |
| if root_key not in deserialized_objects: |
| obj = data_type(size) |
| obj._torch_load_uninitialized = True |
| deserialized_objects[root_key] = restore_location(obj, location) |
| storage = deserialized_objects[root_key] |
| if view_metadata is not None: |
| view_key, offset, view_size = view_metadata |
| if view_key not in deserialized_objects: |
| deserialized_objects[view_key] = storage[offset:offset + view_size] |
| return deserialized_objects[view_key] |
| else: |
| return storage |
| else: |
| raise RuntimeError("Unknown saved id type: %s" % saved_id[0]) |
| |
| _check_seekable(f) |
| f_should_read_directly = _should_read_directly(f) |
| |
| if f_should_read_directly and f.tell() == 0: |
| # legacy_load requires that f has fileno() |
| # only if offset is zero we can attempt the legacy tar file loader |
| try: |
| return legacy_load(f) |
| except tarfile.TarError: |
| if _is_zipfile(f): |
| # .zip is used for torch.jit.save and will throw an un-pickling error here |
| raise RuntimeError( |
| f"{f.name} is a zip archive (did you mean to use torch.jit.load()?)") from None |
| # if not a tarfile, reset file offset and proceed |
| f.seek(0) |
| |
| if not hasattr(f, 'readinto') and (3, 8, 0) <= sys.version_info < (3, 8, 2): |
| raise RuntimeError( |
| "torch.load does not work with file-like objects that do not implement readinto on Python 3.8.0 and 3.8.1. " |
| f"Received object of type \"{type(f)}\". Please update to Python 3.8.2 or newer to restore this " |
| "functionality.") |
| |
| magic_number = pickle_module.load(f, **pickle_load_args) |
| if magic_number != MAGIC_NUMBER: |
| raise RuntimeError("Invalid magic number; corrupt file?") |
| protocol_version = pickle_module.load(f, **pickle_load_args) |
| if protocol_version != PROTOCOL_VERSION: |
| raise RuntimeError("Invalid protocol version: %s" % protocol_version) |
| |
| _sys_info = pickle_module.load(f, **pickle_load_args) |
| unpickler = pickle_module.Unpickler(f, **pickle_load_args) |
| unpickler.persistent_load = persistent_load |
| result = unpickler.load() |
| |
| deserialized_storage_keys = pickle_module.load(f, **pickle_load_args) |
| |
| offset = f.tell() if f_should_read_directly else None |
| for key in deserialized_storage_keys: |
| assert key in deserialized_objects |
| deserialized_objects[key]._set_from_file(f, offset, f_should_read_directly) |
| if offset is not None: |
| offset = f.tell() |
| |
| torch._utils._validate_loaded_sparse_tensors() |
| |
| return result |
| |
| |
| def _maybe_decode_ascii(bytes_str: Union[bytes, str]) -> str: |
| # When using encoding='bytes' in Py3, some **internal** keys stored as |
| # strings in Py2 are loaded as bytes. This function decodes them with |
| # ascii encoding, one that Py3 uses by default. |
| # |
| # NOTE: This should only be used on internal keys (e.g., `typename` and |
| # `location` in `persistent_load` below! |
| if isinstance(bytes_str, bytes): |
| return bytes_str.decode('ascii') |
| return bytes_str |
| |
| |
| def _get_restore_location(map_location): |
| if map_location is None: |
| restore_location = default_restore_location |
| elif isinstance(map_location, dict): |
| def restore_location(storage, location): |
| location = map_location.get(location, location) |
| return default_restore_location(storage, location) |
| elif isinstance(map_location, _string_classes): |
| def restore_location(storage, location): |
| return default_restore_location(storage, map_location) |
| elif isinstance(map_location, torch.device): |
| def restore_location(storage, location): |
| return default_restore_location(storage, str(map_location)) |
| else: |
| def restore_location(storage, location): |
| result = map_location(storage, location) |
| if result is None: |
| result = default_restore_location(storage, location) |
| return result |
| return restore_location |
| |
| |
| def _load(zip_file, map_location, pickle_module, pickle_file='data.pkl', **pickle_load_args): |
| restore_location = _get_restore_location(map_location) |
| |
| loaded_storages = {} |
| |
| def load_tensor(data_type, size, key, location): |
| name = f'data/{key}' |
| dtype = data_type(0).dtype |
| |
| storage = zip_file.get_storage_from_record(name, size, dtype).storage() |
| loaded_storages[key] = restore_location(storage, location) |
| |
| def persistent_load(saved_id): |
| assert isinstance(saved_id, tuple) |
| typename = _maybe_decode_ascii(saved_id[0]) |
| data = saved_id[1:] |
| |
| assert typename == 'storage', \ |
| f"Unknown typename for persistent_load, expected 'storage' but got '{typename}'" |
| data_type, key, location, size = data |
| if key not in loaded_storages: |
| load_tensor(data_type, size, key, _maybe_decode_ascii(location)) |
| storage = loaded_storages[key] |
| return storage |
| |
| # Load the data (which may in turn use `persistent_load` to load tensors) |
| data_file = io.BytesIO(zip_file.get_record(pickle_file)) |
| unpickler = pickle_module.Unpickler(data_file, **pickle_load_args) |
| unpickler.persistent_load = persistent_load |
| result = unpickler.load() |
| |
| torch._utils._validate_loaded_sparse_tensors() |
| |
| return result |
| |
| |
| def _is_torchscript_zip(zip_file): |
| return 'constants.pkl' in zip_file.get_all_records() |
