| import pickle |
| from copy import deepcopy |
|
|
| import numpy as np |
|
|
|
|
| class IndexedDataset: |
| def __init__(self, path, num_cache=1): |
| super().__init__() |
| self.path = path |
| self.data_file = None |
| self.data_offsets = np.load(f"{path}.idx", allow_pickle=True).item()['offsets'] |
| self.data_file = open(f"{path}.data", 'rb', buffering=-1) |
| self.cache = [] |
| self.num_cache = num_cache |
|
|
| def check_index(self, i): |
| if i < 0 or i >= len(self.data_offsets) - 1: |
| raise IndexError('index out of range') |
|
|
| def __del__(self): |
| if self.data_file: |
| self.data_file.close() |
|
|
| def __getitem__(self, i): |
| self.check_index(i) |
| if self.num_cache > 0: |
| for c in self.cache: |
| if c[0] == i: |
| return c[1] |
| self.data_file.seek(self.data_offsets[i]) |
| b = self.data_file.read(self.data_offsets[i + 1] - self.data_offsets[i]) |
| item = pickle.loads(b) |
| if self.num_cache > 0: |
| self.cache = [(i, deepcopy(item))] + self.cache[:-1] |
| return item |
|
|
| def __len__(self): |
| return len(self.data_offsets) - 1 |
|
|
| class IndexedDatasetBuilder: |
| def __init__(self, path): |
| self.path = path |
| self.out_file = open(f"{path}.data", 'wb') |
| self.byte_offsets = [0] |
|
|
| def add_item(self, item): |
| s = pickle.dumps(item) |
| bytes = self.out_file.write(s) |
| self.byte_offsets.append(self.byte_offsets[-1] + bytes) |
|
|
| def finalize(self): |
| self.out_file.close() |
| np.save(open(f"{self.path}.idx", 'wb'), {'offsets': self.byte_offsets}) |
|
|
|
|
| if __name__ == "__main__": |
| import random |
| from tqdm import tqdm |
| ds_path = '/tmp/indexed_ds_example' |
| size = 100 |
| items = [{"a": np.random.normal(size=[10000, 10]), |
| "b": np.random.normal(size=[10000, 10])} for i in range(size)] |
| builder = IndexedDatasetBuilder(ds_path) |
| for i in tqdm(range(size)): |
| builder.add_item(items[i]) |
| builder.finalize() |
| ds = IndexedDataset(ds_path) |
| for i in tqdm(range(10000)): |
| idx = random.randint(0, size - 1) |
| assert (ds[idx]['a'] == items[idx]['a']).all() |
|
|
