def _persistent_id(self, obj):
if torch.is_storage(obj):
storage_type = normalize_storage_type(type(obj))
location = location_tag(obj)
# serialize storage if not already written
storage_present = self.storage_context.has_storage(obj)
storage_id = self.storage_context.get_or_add_storage(obj)
if not storage_present:
if obj.device.type != "cpu":
obj = obj.cpu()
num_bytes = obj.size() * obj.element_size()
self.zip_file.write_record(f".data/{storage_id}.storage",
obj.data_ptr(), num_bytes)
return ("storage", storage_type, storage_id, location, obj.size())
if hasattr(obj, "__reduce_package__"):
if _gate_torchscript_serialization and isinstance(
obj, torch.jit.RecursiveScriptModule):
raise Exception(
"Serializing ScriptModules directly into a package is a beta feature. "
"To use, set global "
"`torch.package.package_exporter._gate_torchscript_serialization` to `False`."
)
if self.serialized_reduces.get(id(obj)) is None:
self.serialized_reduces[id(obj)] = (
"reduce_package",
id(obj),
*obj.__reduce_package__(self),
)
return self.serialized_reduces[id(obj)]
return None
def _persistent_id(self, obj):
if torch.is_storage(obj):
storage_type = normalize_storage_type(type(obj))
obj_key = str(obj._cdata)
location = location_tag(obj)
name = f".data/{obj_key}.storage"
if name not in self.serialized_storages:
# check to see if storage was previously serialized
serialized_files = self.zip_file.get_all_written_records()
if name not in serialized_files:
if obj.device.type != "cpu":
obj = obj.cpu()
num_bytes = obj.size() * obj.element_size()
self.zip_file.write_record(name, obj.data_ptr(), num_bytes)
self.serialized_storages.add(name)
return ("storage", storage_type, obj_key, location, obj.size())
if hasattr(obj, "__reduce_package__"):
if _gate_torchscript_serialization and isinstance(
obj, torch.jit.RecursiveScriptModule):
raise Exception(
"Serializing ScriptModules directly into a package is a beta feature. "
"To use, set global "
"`torch.package.package_exporter._gate_torchscript_serialization` to `False`."
)
if self.serialized_reduces.get(id(obj)) is None:
self.serialized_reduces[id(obj)] = (
"reduce_package", id(obj), *obj.__reduce_package__(self))
return self.serialized_reduces[id(obj)]
return None
def _persistent_id(self, obj):
if torch.is_storage(obj):
storage_type = normalize_storage_type(type(obj))
obj_key = str(obj._cdata)
location = location_tag(obj)
name = f".data/{obj_key}.storage"
if name not in self.serialized_storages:
# check to see if storage was previously serialized
serialized_files = self.zip_file.get_all_written_records()
if name not in serialized_files:
if obj.device.type != "cpu":
obj = obj.cpu()
num_bytes = obj.size() * obj.element_size()
self.zip_file.write_record(name, obj.data_ptr(), num_bytes)
self.serialized_storages.add(name)
return ("storage", storage_type, obj_key, location, obj.size())
if hasattr(obj, "__reduce_package__"):
if self.serialized_reduces.get(id(obj)) is None:
self.serialized_reduces[id(obj)] = ("reduce_package", id(obj), *obj.__reduce_package__(self))
return self.serialized_reduces[id(obj)]
return None
def _persistent_id(self, obj):
if torch.is_storage(obj):
storage_type = normalize_storage_type(type(obj))
obj_key = str(obj._cdata)
location = location_tag(obj)
self.serialized_storages[obj_key] = obj
return ("storage", storage_type, obj_key, location, obj.size())
if hasattr(obj, "__reduce_package__"):
return ("reduce_package", *obj.__reduce_package__(self))
return None
def _persistent_id(self, obj):
if torch.is_storage(obj) or isinstance(obj,
torch.storage._TypedStorage):
if isinstance(obj, torch.storage._TypedStorage):
# TODO: Once we decide to break serialization FC, we can
# remove this case
untyped_storage = obj._storage
storage_type_str = obj.pickle_storage_type()
storage_type = getattr(torch, storage_type_str)
dtype = obj.dtype
storage_numel = obj.size()
elif isinstance(obj, torch._UntypedStorage):
untyped_storage = obj
storage_type = normalize_storage_type(type(storage))
dtype = torch.uint8
storage_numel = storage.nbytes()
else:
raise RuntimeError(f"storage type not recognized: {type(obj)}")
storage: Storage = cast(Storage, untyped_storage)
location = location_tag(storage)
# serialize storage if not already written
storage_present = self.storage_context.has_storage(storage)
storage_id = self.storage_context.get_or_add_storage(storage)
if not storage_present:
if storage.device.type != "cpu":
storage = storage.cpu()
num_bytes = storage.nbytes()
self.zip_file.write_record(f".data/{storage_id}.storage",
storage.data_ptr(), num_bytes)
return ("storage", storage_type, storage_id, location,
storage_numel)
if hasattr(obj, "__reduce_package__"):
if _gate_torchscript_serialization and isinstance(
obj, torch.jit.RecursiveScriptModule):
raise Exception(
"Serializing ScriptModules directly into a package is a beta feature. "
"To use, set global "
"`torch.package.package_exporter._gate_torchscript_serialization` to `False`."
)
if self.serialized_reduces.get(id(obj)) is None:
self.serialized_reduces[id(obj)] = (
"reduce_package",
id(obj),
*obj.__reduce_package__(self),
)
return self.serialized_reduces[id(obj)]
return None
def _persistent_id(self, 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)
self.serialized_storages[obj_key] = obj
return ('storage', storage_type, obj_key, location, obj.size())
return None
