fovi.training.loader

class fovi.training.loader.FlashLoader(*args: Any, **kwargs: Any)[source]

Bases: Loader

Extended FFCV Loader with support for after-batch transforms.

Extends the base FFCV Loader to apply additional PyTorch transformations to batches after they have been processed by the FFCV pipeline.

Parameters:

  • path (str) – Path to the FFCV dataset file (.beton).

  • batch_size (int) – Number of samples per batch.

  • order (ORDER_TYPE) – Sample ordering strategy.

  • num_workers (int, optional) – Number of data loading workers. Defaults to -1.

  • os_cache (bool, optional) – Whether to use OS page cache. Defaults to DEFAULT_OS_CACHE.

  • distributed (bool, optional) – Whether to use distributed sampling. Defaults to False.

  • seed (int, optional) – Random seed for sample ordering. Defaults to None.

  • indices (Sequence[int], optional) – Subset of indices to use. Defaults to None.

  • pipelines (Mapping, optional) – FFCV processing pipelines per field. Defaults to {}.

  • custom_fields (Mapping, optional) – Custom field type mappings. Defaults to {}.

  • drop_last (bool, optional) – Whether to drop the last incomplete batch. Defaults to True.

  • batches_ahead (int, optional) – Number of batches to prefetch. Defaults to 3.

  • recompile (bool, optional) – Whether to recompile pipelines each epoch. Defaults to False.

  • custom_field_mapper (int, optional) – Custom field mapper. Defaults to None.

  • after_batch_pipelines (dict, optional) – Dictionary mapping field names to transforms applied after batch loading. Can be a single transform or list of transforms. Defaults to None.

__init__(path: str, batch_size: int, order: ffcv.loader.loader.ORDER_TYPE, num_workers: int = -1, os_cache: bool = DEFAULT_OS_CACHE, distributed: bool = False, seed: int = None, indices: Sequence[int] = None, pipelines: Mapping[str, Sequence[ffcv.pipeline.operation.Operation | Module]] = {}, custom_fields: Mapping[str, Type[ffcv.fields.base.Field]] = {}, drop_last: bool = True, batches_ahead: int = 3, recompile: bool = False, custom_field_mapper: int = None, after_batch_pipelines: Dict[str, Any] | None = None)[source]