| 122 | self.init_cache() |
| 123 | |
| 124 | def init_cache(self): |
| 125 | assert self.cache_mode in ["part", "full"] |
| 126 | n_sample = len(self.samples) |
| 127 | global_rank = dist.get_rank() |
| 128 | world_size = dist.get_world_size() |
| 129 | |
| 130 | samples_bytes = [None for _ in range(n_sample)] |
| 131 | start_time = time.time() |
| 132 | for index in range(n_sample): |
| 133 | if index % (n_sample // 10) == 0: |
| 134 | t = time.time() - start_time |
| 135 | print(f'global_rank {dist.get_rank()} cached {index}/{n_sample} takes {t:.2f}s per block') |
| 136 | start_time = time.time() |
| 137 | path, target = self.samples[index] |
| 138 | if self.cache_mode == "full": |
| 139 | samples_bytes[index] = (ZipReader.read(path), target) |
| 140 | elif self.cache_mode == "part" and index % world_size == global_rank: |
| 141 | samples_bytes[index] = (ZipReader.read(path), target) |
| 142 | else: |
| 143 | samples_bytes[index] = (path, target) |
| 144 | self.samples = samples_bytes |
| 145 | |
| 146 | def __getitem__(self, index): |
| 147 | """ |