(self, size, seqlen, hidden_size, batch_size, dtype, cache_path='./cache/block_training_data', off_load_to_disk=False)
| 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): |
nothing calls this directly
no test coverage detected