Load data.
(self)
| 309 | pass |
| 310 | |
| 311 | def load_data(self): |
| 312 | """ |
| 313 | Load data. |
| 314 | """ |
| 315 | if self.prefetch_data: |
| 316 | if self._data_prefetched is None: |
| 317 | self._data_prefetched = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True) |
| 318 | data = self._data_prefetched |
| 319 | self._data_prefetched = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True) |
| 320 | else: |
| 321 | data = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True) |
| 322 | |
| 323 | # if the data is a dict, we need to split it into multiple dicts with batch_size_per_gpu |
| 324 | if isinstance(data, dict): |
| 325 | if self.batch_split == 1: |
| 326 | data_list = [data] |
| 327 | else: |
| 328 | batch_size = list(data.values())[0].shape[0] |
| 329 | data_list = [ |
| 330 | {k: v[i * batch_size // self.batch_split:(i + 1) * batch_size // self.batch_split] for k, v in data.items()} |
| 331 | for i in range(self.batch_split) |
| 332 | ] |
| 333 | elif isinstance(data, list): |
| 334 | data_list = data |
| 335 | else: |
| 336 | raise ValueError('Data must be a dict or a list of dicts.') |
| 337 | |
| 338 | return data_list |
| 339 | |
| 340 | @abstractmethod |
| 341 | def run_step(self, data_list): |
no test coverage detected