-
Notifications
You must be signed in to change notification settings - Fork 27.6k
Expand file tree
/
Copy pathserialization.py
More file actions
2201 lines (1845 loc) · 85.3 KB
/
serialization.py
File metadata and controls
2201 lines (1845 loc) · 85.3 KB
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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# mypy: allow-untyped-defs
import copyreg
import difflib
import functools
import io
import os
import pickle
import re
import shutil
import struct
import sys
import tarfile
import tempfile
import threading
import warnings
from collections.abc import Callable
from contextlib import closing, contextmanager
from enum import Enum
from typing import Any, cast, Generic, IO, TypeAlias, TypeVar
from typing_extensions import TypeIs
import torch
import torch._weights_only_unpickler as _weights_only_unpickler
from torch._sources import get_source_lines_and_file
from torch._utils import _import_dotted_name
from torch.storage import _get_dtype_from_pickle_storage_type
from torch.types import FileLike, Storage
__all__ = [
"SourceChangeWarning",
"mkdtemp",
"register_package",
"check_module_version_greater_or_equal",
"validate_cuda_device",
"validate_hpu_device",
"location_tag",
"default_restore_location",
"normalize_storage_type",
"storage_to_tensor_type",
"save",
"load",
"StorageType",
"LoadEndianness",
"get_crc32_options",
"set_crc32_options",
"get_default_load_endianness",
"set_default_load_endianness",
"get_default_mmap_options",
"set_default_mmap_options",
"clear_safe_globals",
"get_safe_globals",
"add_safe_globals",
"safe_globals",
"get_unsafe_globals_in_checkpoint",
"skip_data",
]
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 = ","
MAP_LOCATION: TypeAlias = (
Callable[[Storage, str], Storage] | torch.device | str | dict[str, str] | None
)
STORAGE: TypeAlias = Storage | torch.storage.TypedStorage | torch.UntypedStorage
IS_WINDOWS = sys.platform == "win32"
UNSAFE_MESSAGE = (
"In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` "
"from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, "
"but it can result in arbitrary code execution. Do it only if you got the file from a "
"trusted source."
)
if not IS_WINDOWS:
from mmap import MAP_PRIVATE, MAP_SHARED
else:
MAP_SHARED, MAP_PRIVATE = None, None # type: ignore[assignment]
def _default_to_weights_only(pickle_module):
is_fbcode = not hasattr(torch.version, "git_version")
return pickle_module is None and not is_fbcode
# _serialization_tls is used to store thread local state specific to serialization
# that needs to be propagated to other files, in particular we use this for
# (1) map_location (needed for wrapper subclasses/third party devices to torch._utils)
# (2) skip_data (needed for torch.Tensor.__reduce_ex__ for skip_data ctx)
# (3) materialize_fake_tensors (needed for torch.Tensor.__reduce_ex__ for skip_data ctx)
class _SerializationLocal(threading.local):
def __init__(self):
super().__init__()
self.map_location: MAP_LOCATION | None = None
self.skip_data: bool = False
self.materialize_fake_tensors: bool = False
_serialization_tls = _SerializationLocal()
class SourceChangeWarning(Warning):
pass
@contextmanager
def mkdtemp():
path = tempfile.mkdtemp()
try:
yield path
finally:
shutil.rmtree(path)
_package_registry: list[
tuple[
int,
Callable[[STORAGE], str | None],
Callable[[STORAGE, str], STORAGE | None],
]
] = []
class LoadEndianness(Enum):
NATIVE = 1
LITTLE = 2
BIG = 3
def get_default_load_endianness() -> LoadEndianness | None:
"""
Get fallback byte order for loading files
If byteorder mark is not present in saved checkpoint,
this byte order is used as fallback.
By default, it's "native" byte order.
Returns:
default_load_endian: Optional[LoadEndianness]
"""
from torch.utils.serialization import config
return config.load.endianness
def set_default_load_endianness(endianness):
"""
Set fallback byte order for loading files
If byteorder mark is not present in saved checkpoint,
this byte order is used as fallback.
By default, it's "native" byte order.
Args:
endianness: the new fallback byte order
"""
if not isinstance(endianness, LoadEndianness) and endianness is not None:
raise TypeError("Invalid argument type in function set_default_load_endianness")
from torch.utils.serialization import config
config.load.endianness = endianness
def get_crc32_options() -> bool:
"""
Get whether :func:`torch.save` computes and writes crc32 for each record.
Defaults to ``True``.
"""
from torch.utils.serialization import config
return config.save.compute_crc32
def set_crc32_options(compute_crc32: bool):
"""
Set whether :func:`torch.save` computes and writes crc32 for each record.
.. note::
Setting this to ``False`` may make unzipping of the ``torch.save`` output
fail or warn due to corrupted CRC32. However ``torch.load`` will be
able to load the file.
Args:
compute_crc32 (bool): set crc32 computation flag
"""
from torch.utils.serialization import config
config.save.compute_crc32 = compute_crc32
def get_default_mmap_options() -> int | None:
"""
Get default mmap options for :func:`torch.load` with ``mmap=True``.
Defaults to ``mmap.MAP_PRIVATE``.
Returns:
default_mmap_options: int
"""
from torch.utils.serialization import config
return config.load.mmap_flags
def _get_storage_alignment() -> int:
"""
Gets alignment for storages in torch.save files/
Defaults to 64.
Returns:
storage_alginment: int
"""
from torch.utils.serialization import config
return config.save.storage_alignment
class set_default_mmap_options:
"""
Context manager or function to set default mmap options for :func:`torch.load` with ``mmap=True`` to flags.
For now, only either ``mmap.MAP_PRIVATE`` or ``mmap.MAP_SHARED`` are supported.
Please open an issue if you need any other option to be added here.
.. note::
This feature is currently not supported for Windows.
Args:
flags: ``mmap.MAP_PRIVATE`` or ``mmap.MAP_SHARED``
"""
def __init__(self, flags: int) -> None:
if IS_WINDOWS:
raise RuntimeError(
"Changing the default mmap options is currently not supported for Windows"
)
if flags != MAP_PRIVATE and flags != MAP_SHARED:
raise ValueError(
"Invalid argument in function set_default_mmap_options, "
f"expected mmap.MAP_PRIVATE or mmap.MAP_SHARED, but got {flags}"
)
# global config
from torch.utils.serialization import config
self.prev = config.load.mmap_flags
config.load.mmap_flags = flags
def __enter__(self) -> None:
pass
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
from torch.utils.serialization import config
config.load.mmap_flags = self.prev
def clear_safe_globals() -> None:
"""
Clears the list of globals that are safe for ``weights_only`` load.
"""
_weights_only_unpickler._clear_safe_globals()
def get_safe_globals() -> list[Callable | tuple[Callable, str]]:
"""
Returns the list of user-added globals that are safe for ``weights_only`` load.
"""
return _weights_only_unpickler._get_safe_globals()
def add_safe_globals(safe_globals: list[Callable | tuple[Callable, str]]) -> None:
"""
Marks the given globals as safe for ``weights_only`` load. For example, functions
added to this list can be called during unpickling, classes could be instantiated
and have state set.
Each item in the list can either be a function/class or a tuple of the form
(function/class, string) where string is the full path of the function/class.
Within the serialized format, each function is identified with its full
path as ``{__module__}.{__qualname__}``. When calling this API, you can provide this
full path that should match the one in the checkpoint otherwise the default
``{fn.__module__}.{fn.__qualname__}`` will be used.
Args:
safe_globals (List[Union[Callable, Tuple[Callable, str]]]): list of globals to mark as safe
Example:
>>> # xdoctest: +SKIP("Can't torch.save(t, ...) as doctest thinks MyTensor is defined on torch.serialization")
>>> import tempfile
>>> class MyTensor(torch.Tensor):
... pass
>>> t = MyTensor(torch.randn(2, 3))
>>> with tempfile.NamedTemporaryFile() as f:
... torch.save(t, f.name)
# Running `torch.load(f.name, weights_only=True)` will fail with
# Unsupported global: GLOBAL __main__.MyTensor was not an allowed global by default.
# Check the code and make sure MyTensor is safe to be used when loaded from an arbitrary checkpoint.
... torch.serialization.add_safe_globals([MyTensor])
... torch.load(f.name, weights_only=True)
# MyTensor([[-0.5024, -1.8152, -0.5455],
# [-0.8234, 2.0500, -0.3657]])
"""
_weights_only_unpickler._add_safe_globals(safe_globals)
class safe_globals(_weights_only_unpickler._safe_globals):
r"""Context-manager that adds certain globals as safe for ``weights_only`` load.
Args:
safe_globals: List of globals for weights_only load.
Example:
>>> # xdoctest: +SKIP("Can't torch.save(t, ...) as doctest thinks MyTensor is defined on torch.serialization")
>>> import tempfile
>>> class MyTensor(torch.Tensor):
... pass
>>> t = MyTensor(torch.randn(2, 3))
>>> with tempfile.NamedTemporaryFile() as f:
... torch.save(t, f.name)
# Running `torch.load(f.name, weights_only=True)` will fail with
# Unsupported global: GLOBAL __main__.MyTensor was not an allowed global by default.
# Check the code and make sure MyTensor is safe to be used when loaded from an arbitrary checkpoint.
... with torch.serialization.safe_globals([MyTensor]):
... torch.load(f.name, weights_only=True)
# MyTensor([[-0.5024, -1.8152, -0.5455],
# [-0.8234, 2.0500, -0.3657]])
>>> assert torch.serialization.get_safe_globals() == []
"""
def get_unsafe_globals_in_checkpoint(f: FileLike) -> list[str]:
"""Returns a list of strings of functions/classes in a ``torch.save`` object that are not safe for ``weights_only``.
For a given function or class ``f``, the corresponding string will be of the form
``{f.__module__}.{f.__name__}``.
This function will return any GLOBALs in the checkpoint that are not in the set marked safe
for ``weights_only`` (either via :func:`add_safe_globals` or :class:`safe_globals` context or
allowlisted by ``torch`` by default).
.. note::
This function will statically disassemble the pickle file in the checkpoint.
The implication is any classes dynamically pushed onto the stack during unpickling
will not be included in the output.
Args:
f: File-like object or string containing the checkpoint object saved via ``torch.save``
Returns:
A list of strings of pickle GLOBALs in the checkpoint that are not allowlisted for ``weights_only``.
"""
default_safe_globals_strings = set(
_weights_only_unpickler._get_allowed_globals().keys()
)
user_safe_global_strings = set(
_weights_only_unpickler._get_user_allowed_globals().keys()
)
safe_global_strings = default_safe_globals_strings.union(user_safe_global_strings)
with _open_file_like(f, "rb") as opened_file:
if not _is_zipfile(opened_file):
raise ValueError("Expected input to be a checkpoint returned by torch.save")
with _open_zipfile_reader(opened_file) as zip_file:
if _is_torchscript_zip(zip_file):
raise ValueError(
"Expected input to be a checkpoint returned by torch.save but got a torchscript checkpoint"
)
data_file = io.BytesIO(zip_file.get_record("data.pkl"))
all_globals = _weights_only_unpickler.get_globals_in_pkl(data_file)
return list(all_globals.difference(safe_global_strings))
class skip_data:
"""
Context-manager that skips writing/reading storage bytes for ``torch.save`` / ``torch.load`` calls.
For the save path, storages will still be saved, but the space that their bytes would usually be written to
will be empty space. The storage bytes can then be populated in a separate pass.
For the load path, tensors will be loaded per the checkpoint but their storages will not be populated with data.
.. warning::
The ``skip_data`` context manager is an early prototype and is subject to change.
Args:
materialize_fake_tensors: Whether to materialize FakeTensors during save. This is a no-op for the load path.
Example:
>>> # xdoctest: +SKIP("NamedTemporaryFile on Windows")
>>> import tempfile
>>> t = torch.randn(2, 3)
>>> with tempfile.NamedTemporaryFile() as f:
... with torch.serialization.skip_data():
... torch.save(t, f.name)
... torch.load(f.name, weights_only=True)
tensor([[0., 0., 0.],
[0., 0., 0.]])
"""
def __init__(self, materialize_fake_tensors: bool = False):
self.materialize_fake_tensors = materialize_fake_tensors
def __enter__(self):
global _serialization_tls
self._old_skip_data = _serialization_tls.skip_data
self._old_materialize_fake_tensors = _serialization_tls.materialize_fake_tensors
_serialization_tls.skip_data = True
_serialization_tls.materialize_fake_tensors = self.materialize_fake_tensors
def __exit__(self, type, value, tb):
global _serialization_tls
_serialization_tls.skip_data = self._old_skip_data
_serialization_tls.materialize_fake_tensors = self._old_materialize_fake_tensors
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.
start = f.tell()
# Read the first few bytes and match against the ZIP file signature
local_header_magic_number = b"PK\x03\x04"
read_bytes = f.read(len(local_header_magic_number))
f.seek(start)
return read_bytes == local_header_magic_number
def register_package(
priority: int,
tagger: Callable[[STORAGE], str | None],
deserializer: Callable[[STORAGE, str], STORAGE | None],
):
"""
Registers callables for tagging and deserializing storage objects with an associated priority.
Tagging associates a device with a storage object at save time while deserializing moves a
storage object to an appropriate device at load time. :attr:`tagger` and :attr:`deserializer`
are run in the order given by their :attr:`priority` until a tagger/deserializer returns a
value that is not `None`.
To override the deserialization behavior for a device in the global registry, one can register a
tagger with a higher priority than the existing tagger.
This function can also be used to register a tagger and deserializer for new devices.
Args:
priority: Indicates the priority associated with the tagger and deserializer, where a lower
value indicates higher priority.
tagger: Callable that takes in a storage object and returns its tagged device as a string
or None.
deserializer: Callable that takes in storage object and a device string and returns a storage
object on the appropriate device or None.
Returns:
`None`
Example:
>>> def ipu_tag(obj):
>>> if obj.device.type == 'ipu':
>>> return 'ipu'
>>> def ipu_deserialize(obj, location):
>>> if location.startswith('ipu'):
>>> ipu = getattr(torch, "ipu", None)
>>> assert ipu is not None, "IPU device module is not loaded"
>>> assert torch.ipu.is_available(), "ipu is not available"
>>> return obj.ipu(location)
>>> torch.serialization.register_package(11, ipu_tag, ipu_deserialize)
"""
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 = (
f"'{module.__name__}' module version string is malformed '{module.__version__}' and cannot be compared"
f" with tuple {str(req_version_tuple)}"
)
if error_if_malformed:
raise RuntimeError(message) from e
else:
warnings.warn(
message + ", but continuing assuming that requirement is met",
stacklevel=2,
)
requirement_is_met = True
return requirement_is_met
def _cpu_tag(obj):
if obj.device.type == "cpu":
return "cpu"
def _mps_tag(obj):
if obj.device.type == "mps":
return "mps"
def _meta_tag(obj):
if obj.device.type == "meta":
return "meta"
def _backend_tag(backend_name, obj):
if backend_name == "privateuse1":
backend_name = torch._C._get_privateuse1_backend_name()
if obj.device.type == backend_name:
if obj.device.index is None:
return backend_name
else:
return backend_name + ":" + str(obj.device.index)
def _cpu_deserialize(obj, location):
if location == "cpu":
return obj
def _mps_deserialize(obj, location):
if location.startswith("mps"):
return obj.mps()
def _meta_deserialize(obj, location):
if location == "meta":
return torch.UntypedStorage(obj.nbytes(), device="meta")
def _is_meta_location(map_location):
"""
Check if map_location specifies the meta device.
This is used to skip reading storage data from disk when loading
to the meta device, since meta tensors don't hold actual data.
Args:
map_location: The map_location argument passed to torch.load
Returns:
True if map_location is definitively the meta device, False otherwise.
For dict or callable map_location, returns False since we can't
easily determine the target device without evaluating it.
"""
if map_location is None:
return False
if isinstance(map_location, str):
return map_location == "meta"
if isinstance(map_location, torch.device):
return map_location.type == "meta"
# dict or callable - can't easily determine
return False
def _validate_device(location, backend_name):
"""
Check whether the device index of specified backend is valid
In case of privateuse1 backend, your must first register a device_module for
privateuse1 using torch._register_device_module. Implement the following
methods in device_module like cuda: device_module._utils._get_device_index(location, True),
device_module.device_count().
Args:
location: string of device
backend_name: the backend name or the name of privateuse1, which can be renamed
Returns:
device_index: int
"""
if not hasattr(torch, backend_name):
raise RuntimeError(
f"The {backend_name.upper()} device module is not registered. "
"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_module = getattr(torch, backend_name)
if hasattr(device_module, "_utils") and hasattr(
device_module._utils, "_get_device_index"
):
device_index = device_module._utils._get_device_index(location, True)
device = torch.device(backend_name, device_index)
else:
device = torch.device(location)
device_index = device.index if device.index else 0
if hasattr(device_module, "is_available") and not device_module.is_available():
raise RuntimeError(
f"Attempting to deserialize object on a {backend_name.upper()} "
f"device but torch.{backend_name}.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."
)
if hasattr(device_module, "device_count"):
device_count = device_module.device_count()
if device_index >= device_count:
raise RuntimeError(
f"Attempting to deserialize object on {backend_name.upper()} device "
f"{device_index} but torch.{backend_name}.device_count() is {device_count}. "
"Please use torch.load with map_location to map your storages "
"to an existing device."
)
return device
def validate_cuda_device(location):
return _validate_device(location, "cuda").index
def validate_hpu_device(location):
return _validate_device(location, "hpu").index
def _deserialize(backend_name, obj, location):
if backend_name == "privateuse1":
backend_name = torch._C._get_privateuse1_backend_name()
if location == backend_name or bool(
re.match(f"{backend_name}(:|[0-9]+)", location)
):
device = _validate_device(location, backend_name)
return obj.to(device=device)
register_package(10, _cpu_tag, _cpu_deserialize)
register_package(
20,
functools.partial(_backend_tag, "cuda"),
functools.partial(_deserialize, "cuda"),
)
register_package(21, _mps_tag, _mps_deserialize)
register_package(22, _meta_tag, _meta_deserialize)
register_package(
23,
functools.partial(_backend_tag, "privateuse1"),
functools.partial(_deserialize, "privateuse1"),
)
register_package(
24,
functools.partial(_backend_tag, "hpu"),
functools.partial(_deserialize, "hpu"),
)
register_package(
25,
functools.partial(_backend_tag, "xpu"),
functools.partial(_deserialize, "xpu"),
)
register_package(
26,
functools.partial(_backend_tag, "mtia"),
functools.partial(_deserialize, "mtia"),
)
def location_tag(
storage: Storage | torch.storage.TypedStorage | torch.UntypedStorage,
):
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):
"""
Restores `storage` using a deserializer function registered for the `location`.
This function looks in the registry for deserializer functions that match the `location`.
If found, it attempts to use them, in priority order, to restore `storage` until one
returns a not `None` result. If no deserializer can be found in the registry, or all found fail
to bear a result, it raises a `RuntimeError`.
Args:
storage (STORAGE): the storage object to restore
location (str): the location tag associated with the storage object
Returns:
storage: Optional[STORAGE]
Raises:
RuntimeError: If no deserializer matching `location` is found in the registry or if
all matching ones return `None`.
"""
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: object) -> TypeIs[str | os.PathLike]:
return isinstance(name_or_buffer, (str, os.PathLike))
T = TypeVar("T")
class _opener(Generic[T]):
def __init__(self, file_like: T) -> None:
self.file_like: T = file_like
def __enter__(self):
return self.file_like
def __exit__(self, *args):
pass
class _open_file(_opener[IO[bytes]]):
def __init__(self, name: str | os.PathLike[str], mode: str) -> None:
super().__init__(open(name, mode)) # noqa: SIM115
def __exit__(self, *args):
self.file_like.close()
class _open_buffer_reader(_opener[IO[bytes]]):
def __init__(self, buffer: IO[bytes]) -> None:
super().__init__(buffer)
_check_seekable(buffer)
class _open_buffer_writer(_opener[IO[bytes]]):
def __exit__(self, *args):
self.file_like.flush()
def _open_file_like(name_or_buffer: FileLike, mode: str) -> _opener[IO[bytes]]:
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[torch._C.PyTorchFileReader]):
def __init__(self, name_or_buffer: str | IO[bytes]) -> None:
super().__init__(torch._C.PyTorchFileReader(name_or_buffer))
class _open_zipfile_writer_file(_opener[torch._C.PyTorchFileWriter]):
def __init__(self, name: str) -> None:
self.file_stream = None
self.name = name
try:
self.name.encode("ascii")
except UnicodeEncodeError:
# PyTorchFileWriter only supports ascii filename.
# For filenames with non-ascii characters, we rely on Python
# for writing out the file.
# pyrefly: ignore [bad-assignment]
self.file_stream = io.FileIO(self.name, mode="w")
super().__init__(
torch._C.PyTorchFileWriter( # pyrefly: ignore # no-matching-overload
self.file_stream, get_crc32_options(), _get_storage_alignment()
)
)
else:
super().__init__(
torch._C.PyTorchFileWriter(
self.name, get_crc32_options(), _get_storage_alignment()
)
)
def __exit__(self, *args) -> None:
self.file_like.write_end_of_file()
if self.file_stream is not None:
self.file_stream.close()
class _open_zipfile_writer_buffer(_opener[torch._C.PyTorchFileWriter]):
def __init__(self, buffer: IO[bytes]) -> None:
if not callable(getattr(buffer, "write", None)):
msg = f"Buffer of {str(type(buffer)).strip('<>')} has no callable attribute 'write'"
if not hasattr(buffer, "write"):
raise AttributeError(msg)
raise TypeError(msg)
self.buffer = buffer
super().__init__(
torch._C.PyTorchFileWriter(
buffer, get_crc32_options(), _get_storage_alignment()
)
)
def __exit__(self, *args) -> None:
self.file_like.write_end_of_file()
self.buffer.flush()
def _open_zipfile_writer(name_or_buffer: str | IO[bytes]) -> _opener:
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) # type: ignore[arg-type]
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 is not None and 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 >= {}, but you have dill {}."
" Please upgrade dill or switch to 'pickle'"
).format(
".".join([str(num) for num in required_dill_version]),
pickle_module.__version__,
)
)
def _check_save_filelike(f):
if not _is_path(f) and not hasattr(f, "write"):
raise AttributeError(
"expected 'f' to be string, path, or a file-like object with "
"a 'write' attribute"
)
def save(
obj: object,
f: FileLike,
pickle_module: Any = pickle,
pickle_protocol: int = DEFAULT_PROTOCOL,
_use_new_zipfile_serialization: bool = True,
_disable_byteorder_record: bool = False,
) -> None:
# Reference: https://github.com/pytorch/pytorch/issues/54354
# The first line of this docstring overrides the one Sphinx generates for the
# documentation. We need it so that Sphinx doesn't leak `pickle`s path from
# the build environment (e.g. `<module 'pickle' from '/leaked/path').
"""save(obj, f, pickle_module=pickle, pickle_protocol=2, _use_new_zipfile_serialization=True)
Saves an object to a disk file.
See also: :ref:`saving-loading-tensors`
See :ref:`layout-control` for more advanced tools to manipulate a checkpoint.
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
:ref:`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:
>>> # xdoctest: +SKIP("makes cwd dirty")
>>> # 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)
"""
torch._C._log_api_usage_once("torch.save")
_check_dill_version(pickle_module)
_check_save_filelike(f)
if isinstance(f, (str, os.PathLike)):
f = os.fspath(f)