Implements chunked memmap. - Designed to use with a dataset loader (no slicing support). - It handles saving/loading the meta data automatically, so can be used like a normal ndarray. - It supports chunking of the array to improve read/write speed and disk space efficiency. Usage:
| 24 | |
| 25 | |
| 26 | class ChunkedMemmap: |
| 27 | """Implements chunked memmap. |
| 28 | |
| 29 | - Designed to use with a dataset loader (no slicing support). |
| 30 | - It handles saving/loading the meta data automatically, so can be used like a normal ndarray. |
| 31 | - It supports chunking of the array to improve read/write speed and disk space efficiency. |
| 32 | |
| 33 | Usage: |
| 34 | |
| 35 | - add_all_samples(samples, chunk_idxs): |
| 36 | divide samples (one per row) into different numpy memmap according to chunk_idx. |
| 37 | - __getitem__: |
| 38 | only support reading ith sample (no slicing). The indexing is the same as if the samples are not chunked. |
| 39 | - __len__: |
| 40 | return total number of samples. |
| 41 | """ |
| 42 | |
| 43 | def __init__(self, working_dir: str, remove_exist=False): |
| 44 | """Create memmaps that store chunked of the samples. |
| 45 | |
| 46 | Args: |
| 47 | working_dir: |
| 48 | the directory where all the chunked memmaps are stored. |
| 49 | remove_exist: |
| 50 | whether or not to remove existing content in the working_dir. |
| 51 | Set to False if want to read existed chunked_memmap. |
| 52 | """ |
| 53 | self.working_dir = working_dir |
| 54 | self.info_filename = os.path.join(self.working_dir, "info.json") |
| 55 | |
| 56 | if remove_exist and os.path.exists(self.working_dir): |
| 57 | shutil.rmtree(self.working_dir) |
| 58 | |
| 59 | # read the memmaps if exist |
| 60 | self._read_info() |
| 61 | |
| 62 | def _read_info(self): |
| 63 | """Read the info file if existed, else set everything to None.""" |
| 64 | if not os.path.exists(self.info_filename): |
| 65 | self.dtype = None |
| 66 | self.itemsize = None # bytes per number |
| 67 | self.chunk_idxs: T.Sequence[int] = None |
| 68 | self.chunk_local_idxs: T.Sequence[int] = None # local index within each chunk memmap |
| 69 | self.chunk_filename_dict: T.Dict[int, str] = None # chunk_idx -> memmap filename |
| 70 | self.chunk_offsets_dict: T.Dict[ |
| 71 | int, T.Sequence[int] |
| 72 | ] = None # chunk_idx -> list of sample offset in the chunk |
| 73 | self.chunk_sample_shapes_dict: T.Dict[ |
| 74 | int, T.Sequence[T.Sequence[int]] |
| 75 | ] = None # chunk_idx -> list of sample_shapes in the chunk |
| 76 | self.chunk_shape_dict: T.Dict[int, T.Tuple[int]] = None # chunk_idx -> chunk's memmap shape |
| 77 | else: |
| 78 | with open(self.info_filename, "r") as f: |
| 79 | info_dict = json.load(f) |
| 80 | |
| 81 | for key in [ |
| 82 | "dtype", |
| 83 | "itemsize", |
no outgoing calls