| 307 | |
| 308 | |
| 309 | class IncrementalPyTorchPickler(pickle.Pickler): |
| 310 | def __init__(self, saver, *args, **kwargs): |
| 311 | super().__init__(*args, **kwargs) |
| 312 | self.storage_dtypes = {} |
| 313 | self.saver = saver |
| 314 | self.id_map = {} |
| 315 | |
| 316 | # this logic is taken from PyTorch 2.0+ torch/serialization.py |
| 317 | def persistent_id(self, obj): |
| 318 | # FIXME: the docs say that persistent_id should only return a string |
| 319 | # but torch store returns tuples. This works only in the binary protocol |
| 320 | # see |
| 321 | # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects |
| 322 | # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537 |
| 323 | if isinstance(obj, SavingProxyForStorage): |
| 324 | return obj.storage_info |
| 325 | |
| 326 | if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj): |
| 327 | if isinstance(obj, torch.storage.TypedStorage): |
| 328 | # TODO: Once we decide to break serialization FC, this case |
| 329 | # can be deleted |
| 330 | storage = obj._untyped_storage |
| 331 | storage_dtype = obj.dtype |
| 332 | storage_type_str = obj._pickle_storage_type() |
| 333 | storage_type = getattr(torch, storage_type_str) |
| 334 | storage_numel = obj._size() |
| 335 | |
| 336 | else: |
| 337 | storage = obj |
| 338 | storage_dtype = torch.uint8 |
| 339 | storage_type = normalize_storage_type(type(obj)) |
| 340 | storage_numel = storage.nbytes() |
| 341 | |
| 342 | # If storage is allocated, ensure that any other saved storages |
| 343 | # pointing to the same data all have the same dtype. If storage is |
| 344 | # not allocated, don't perform this check |
| 345 | if storage.data_ptr() != 0: |
| 346 | if storage.data_ptr() in self.storage_dtypes: |
| 347 | if storage_dtype != self.storage_dtypes[storage.data_ptr()]: |
| 348 | raise RuntimeError( |
| 349 | "Cannot save multiple tensors or storages that view the same data as different types" |
| 350 | ) |
| 351 | else: |
| 352 | self.storage_dtypes[storage.data_ptr()] = storage_dtype |
| 353 | |
| 354 | storage_key = self.id_map.get(storage._cdata) |
| 355 | if storage_key is None: |
| 356 | storage_key = self.saver._write_storage_and_return_key(storage) |
| 357 | self.id_map[storage._cdata] = storage_key |
| 358 | location = torch.serialization.location_tag(storage) |
| 359 | |
| 360 | return ("storage", storage_type, storage_key, location, storage_numel) |
| 361 | |
| 362 | return None |
| 363 | |
| 364 | |
| 365 | class incremental_save: |