| 778 | |
| 779 | |
| 780 | class LazyUnpickler(pickle.Unpickler): |
| 781 | def __init__(self, fp: IO[bytes], data_base_path: str, zip_file: zipfile.ZipFile): |
| 782 | super().__init__(fp) |
| 783 | self.data_base_path = data_base_path |
| 784 | self.zip_file = zip_file |
| 785 | |
| 786 | def persistent_load(self, pid: Any) -> Any: |
| 787 | assert pid[0] == 'storage' |
| 788 | assert isinstance(pid[1], LazyStorageKind) |
| 789 | data_type = pid[1].data_type |
| 790 | filename_stem = pid[2] |
| 791 | filename = f'{self.data_base_path}/{filename_stem}' |
| 792 | info = self.zip_file.getinfo(filename) |
| 793 | |
| 794 | def load(offset: int, elm_count: int) -> NDArray: |
| 795 | dtype = data_type.dtype |
| 796 | fp = self.zip_file.open(info) |
| 797 | fp.seek(offset * dtype.itemsize) |
| 798 | size = elm_count * dtype.itemsize |
| 799 | data = fp.read(size) |
| 800 | assert len(data) == size |
| 801 | return np.frombuffer(data, dtype) |
| 802 | description = f'storage data_type={data_type} path-in-zip={filename} path={self.zip_file.filename}' |
| 803 | return LazyStorage(load=load, kind=pid[1], description=description) |
| 804 | |
| 805 | @staticmethod |
| 806 | def lazy_rebuild_tensor_v2(storage: Any, storage_offset: Any, size: Any, stride: Any, |
| 807 | requires_grad: Any, backward_hooks: Any, metadata: Any = None) -> LazyTensor: |
| 808 | assert isinstance(storage, LazyStorage) |
| 809 | |
| 810 | def load() -> UnquantizedTensor: |
| 811 | elm_count = stride[0] * size[0] |
| 812 | return UnquantizedTensor(storage.load(storage_offset, elm_count).reshape(size)) |
| 813 | description = f'pickled storage_offset={storage_offset} in {storage.description}' |
| 814 | return LazyTensor(load, list(size), storage.kind.data_type, description) |
| 815 | |
| 816 | @staticmethod |
| 817 | def rebuild_from_type_v2(func, new_type, args, state): |
| 818 | return func(*args) |
| 819 | |
| 820 | CLASSES: dict[tuple[str, str], Any] = { |
| 821 | # getattr used here as a workaround for mypy not being smart enough to determine |
| 822 | # the staticmethods have a __func__ attribute. |
| 823 | ('torch._tensor', '_rebuild_from_type_v2'): getattr(rebuild_from_type_v2, '__func__'), |
| 824 | ('torch._utils', '_rebuild_tensor_v2'): getattr(lazy_rebuild_tensor_v2, '__func__'), |
| 825 | ('torch', 'BFloat16Storage'): LazyStorageKind(DT_BF16), |
| 826 | ('torch', 'HalfStorage'): LazyStorageKind(DT_F16), |
| 827 | ('torch', 'FloatStorage'): LazyStorageKind(DT_F32), |
| 828 | ('torch', 'IntStorage'): LazyStorageKind(DT_I32), |
| 829 | ('torch', 'Tensor'): LazyTensor, |
| 830 | } |
| 831 | |
| 832 | def find_class(self, module: str, name: str) -> Any: |
| 833 | if not module.startswith('torch'): |
| 834 | return super().find_class(module, name) |
| 835 | return self.CLASSES[(module, name)] |
| 836 | |
| 837 |
no test coverage detected