| 229 | return results |
| 230 | |
| 231 | class BlockTrainDataset(Dataset): |
| 232 | def __init__(self, size, seqlen, hidden_size, batch_size, dtype, cache_path='./cache/block_training_data', off_load_to_disk=False): |
| 233 | self.size = size |
| 234 | self.seqlen = seqlen |
| 235 | self.hidden_size = hidden_size |
| 236 | self.dtype = dtype |
| 237 | self.cache_path = cache_path |
| 238 | self.off_load_to_disk = off_load_to_disk |
| 239 | self.batch_size = batch_size |
| 240 | assert size%batch_size == 0 |
| 241 | |
| 242 | if self.off_load_to_disk: |
| 243 | if not os.path.exists(self.cache_path): |
| 244 | os.makedirs(self.cache_path) |
| 245 | self._initialize_data_on_disk() |
| 246 | else: |
| 247 | self.data = torch.zeros((self.size//self.batch_size, self.batch_size, self.seqlen, self.hidden_size), dtype=self.dtype) |
| 248 | |
| 249 | def _initialize_data_on_disk(self): |
| 250 | for idx in range(self.size//self.batch_size): |
| 251 | tensor = torch.zeros((self.batch_size, self.seqlen, self.hidden_size), dtype=self.dtype) |
| 252 | filepath = self._get_file_path(idx) |
| 253 | torch.save(tensor, filepath) |
| 254 | |
| 255 | def _get_file_path(self, idx): |
| 256 | return os.path.join(self.cache_path, f"data_{idx}.pt") |
| 257 | |
| 258 | def __len__(self): |
| 259 | return self.size//self.batch_size |
| 260 | |
| 261 | def __getitem__(self, idx): |
| 262 | if idx >= self.__len__(): |
| 263 | raise IndexError("Index out of range") |
| 264 | if self.off_load_to_disk: |
| 265 | filepath = self._get_file_path(idx) |
| 266 | tensor = torch.load(filepath) |
| 267 | else: |
| 268 | tensor = self.data[idx] |
| 269 | return tensor |
| 270 | |
| 271 | def update_data(self, idx, new_data): |
| 272 | if self.off_load_to_disk: |
| 273 | filepath = self._get_file_path(idx) |
| 274 | torch.save(new_data.to(self.dtype), filepath) |
| 275 | else: |
| 276 | self.data[idx] = new_data |