| 170 | |
| 171 | |
| 172 | class IncrementalPyTorchPickler(pickle.Pickler): |
| 173 | def __init__(self, saver, *args, **kwargs): |
| 174 | super().__init__(*args, **kwargs) |
| 175 | self.storage_dtypes = {} |
| 176 | self.saver = saver |
| 177 | self.id_map = {} |
| 178 | |
| 179 | # this logic is taken from PyTorch 2.0+ torch/serialization.py |
| 180 | def persistent_id(self, obj): |
| 181 | # FIXME: the docs say that persistent_id should only return a string |
| 182 | # but torch store returns tuples. This works only in the binary protocol |
| 183 | # see |
| 184 | # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects |
| 185 | # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537 |
| 186 | if isinstance(obj, SavingProxyForStorage): |
| 187 | return obj.storage_info |
| 188 | |
| 189 | if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj): |
| 190 | if isinstance(obj, torch.storage.TypedStorage): |
| 191 | # TODO: Once we decide to break serialization FC, this case |
| 192 | # can be deleted |
| 193 | storage = obj._untyped_storage |
| 194 | storage_dtype = obj.dtype |
| 195 | storage_type_str = obj._pickle_storage_type() |
| 196 | storage_type = getattr(torch, storage_type_str) |
| 197 | storage_numel = obj._size() |
| 198 | |
| 199 | else: |
| 200 | storage = obj |
| 201 | storage_dtype = torch.uint8 |
| 202 | storage_type = normalize_storage_type(type(obj)) |
| 203 | storage_numel = storage.nbytes() |
| 204 | |
| 205 | # If storage is allocated, ensure that any other saved storages |
| 206 | # pointing to the same data all have the same dtype. If storage is |
| 207 | # not allocated, don't perform this check |
| 208 | if storage.data_ptr() != 0: |
| 209 | if storage.data_ptr() in self.storage_dtypes: |
| 210 | if storage_dtype != self.storage_dtypes[storage.data_ptr()]: |
| 211 | raise RuntimeError( |
| 212 | 'Cannot save multiple tensors or storages that view the same data as different types' |
| 213 | ) |
| 214 | else: |
| 215 | self.storage_dtypes[storage.data_ptr()] = storage_dtype |
| 216 | |
| 217 | storage_key = self.id_map.get(storage._cdata) |
| 218 | if storage_key is None: |
| 219 | storage_key = self.saver._write_storage_and_return_key(storage) |
| 220 | self.id_map[storage._cdata] = storage_key |
| 221 | location = torch.serialization.location_tag(storage) |
| 222 | |
| 223 | return ('storage', storage_type, storage_key, location, storage_numel) |
| 224 | |
| 225 | return None |
| 226 | |
| 227 | |
| 228 | class incremental_save: |