It would be nice to be able to add arguments for a given TypedDict, in addition to dataclasses.
class FfcvLoaderConfig(TypedDict, total=False):
os_cache: bool
""" Leverages the operating system for caching purposes. This is beneficial when there is
enough memory to cache the dataset and/or when multiple processes on the same machine training
using the same dataset. See https://docs.ffcv.io/performance_guide.html for more information.
"""
order: TraversalOrder
"""Traversal order, one of: SEQUENTIAL, RANDOM, QUASI_RANDOM
QUASI_RANDOM is a random order that tries to be as uniform as possible while minimizing the
amount of data read from the disk. Note that it is mostly useful when `os_cache=False`.
Currently unavailable in distributed mode.
"""
distributed: bool
"""For distributed training (multiple GPUs). Emulates the behavior of DistributedSampler from
PyTorch.
"""
seed: int
"""Random seed for batch ordering."""
indices: Sequence[int]
"""Subset of dataset by filtering only some indices. """
custom_fields: Mapping[str, type[Field]]
"""Dictonary informing the loader of the types associated to fields that are using a custom
type.
"""
drop_last: bool
"""Drop non-full batch in each iteration."""
batches_ahead: int
"""Number of batches prepared in advance; balances latency and memory. """
recompile: bool
"""Recompile every iteration. This is necessary if the implementation of some augmentations
are expected to change during training.
"""
if __name__ == "__main__":
parser = simple_parsing.ArgumentParser()
parser.add_arguments(FfcvLoaderConfig, "loader")
args = parser.parse_args()
loader_config: FfcvLoaderConfig = args.loader
This could also extend to TypedDict fields on dataclasses, or on TypedDicts themselves.
It would be nice to be able to add arguments for a given TypedDict, in addition to dataclasses.
This could also extend to TypedDict fields on dataclasses, or on TypedDicts themselves.