| 274 | |
| 275 | # Caches data loaded by a DataLoader for use across multiple runners. |
| 276 | class DataLoaderCache: |
| 277 | def __init__(self, data_loader, save_inputs_path=None): |
| 278 | self.data_loader = data_loader |
| 279 | self.cache = [] # List[OrderedDict[str, numpy.ndarray]] |
| 280 | self.save_inputs_path = save_inputs_path |
| 281 | |
| 282 | @func.constantmethod |
| 283 | def __getitem__(self, iteration): |
| 284 | """ |
| 285 | Load the specified iteration from the cache if present, or load it from the data loader. |
| 286 | |
| 287 | Args: |
| 288 | iteration (int): The iteration whose data to retrieve. |
| 289 | """ |
| 290 | if iteration >= len(self.cache): |
| 291 | raise IndexError() |
| 292 | |
| 293 | # Attempts to match existing input buffers to the requested input_metadata |
| 294 | def coerce_cached_input(index, name, dtype, shape): |
| 295 | cached_feed_dict = self.cache[iteration] |
| 296 | cached_name = util.find_str_in_iterable(name, cached_feed_dict.keys(), index) |
| 297 | if cached_name is None: |
| 298 | G_LOGGER.critical(f"Input tensor: {name} | Does not exist in the data loader cache.") |
| 299 | |
| 300 | if cached_name != name: |
| 301 | G_LOGGER.warning( |
| 302 | f"Input tensor: {name} | Buffer name ({cached_name}) does not match expected input name ({name})." |
| 303 | ) |
| 304 | |
| 305 | buffer = cached_feed_dict[cached_name] |
| 306 | |
| 307 | if dtype != buffer.dtype: |
| 308 | G_LOGGER.warning( |
| 309 | f"Input tensor: {name} | Buffer dtype ({buffer.dtype}) does not match expected input dtype ({np.dtype(dtype).name}), attempting to cast. " |
| 310 | ) |
| 311 | |
| 312 | type_info = None |
| 313 | if np.issubdtype(dtype, np.integer): |
| 314 | type_info = np.iinfo(np.dtype(dtype)) |
| 315 | elif np.issubdtype(dtype, np.floating): |
| 316 | type_info = np.finfo(np.dtype(dtype)) |
| 317 | |
| 318 | if type_info is not None and np.any((buffer < type_info.min) | (buffer > type_info.max)): |
| 319 | G_LOGGER.warning( |
| 320 | f"Some values in this input are out of range of {dtype}. Unexpected behavior may ensue!" |
| 321 | ) |
| 322 | buffer = buffer.astype(dtype) |
| 323 | |
| 324 | if not util.is_valid_shape_override(buffer.shape, shape): |
| 325 | G_LOGGER.warning( |
| 326 | f"Input tensor: {name} | Buffer shape ({buffer.shape}) does not match expected input shape ({shape}), attempting to transpose/reshape. " |
| 327 | ) |
| 328 | buffer = util.try_match_shape(buffer, shape) |
| 329 | |
| 330 | if buffer.dtype != dtype or not util.is_valid_shape_override(buffer.shape, shape): |
| 331 | G_LOGGER.critical( |
| 332 | f"Input tensor: {name} | Cannot reuse input data due to mismatch in shape or data type.\nNote: Cached input: [dtype={buffer.dtype}, shape={buffer.shape}], Requested input: [dtype={dtype}, shape={shape}]" |
| 333 | ) |
no outgoing calls