40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
|
"""Dataset bundle returned by data backends."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Any, Callable, Iterable, Optional
|
||
|
|
|
||
|
|
|
||
|
|
def _noop_set_epoch(_epoch: int) -> None:
|
||
|
|
"""Default ``set_epoch`` for backends with epoch-agnostic shuffling."""
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class DataBundle:
|
||
|
|
"""Container returned by every backend ``build()``.
|
||
|
|
|
||
|
|
Attributes:
|
||
|
|
dataset: backend-private; trainer code should treat it as opaque.
|
||
|
|
train_loader: anything iterable that yields training batches.
|
||
|
|
val_loader: val iterable or None if the backend does not split val.
|
||
|
|
train_iters: one-epoch step count for the train loader. Backends
|
||
|
|
with dynamic resizing should set this to a stable snapshot
|
||
|
|
and expose the live value separately on ``dataset``.
|
||
|
|
val_iters: one-epoch step count for the val loader; 0 if no val.
|
||
|
|
set_epoch: per-epoch seed hook. Called before each epoch by
|
||
|
|
the trainer. Backends that don't need per-epoch reshuffle
|
||
|
|
should use ``_noop_set_epoch``.
|
||
|
|
"""
|
||
|
|
|
||
|
|
dataset: Any
|
||
|
|
train_loader: Iterable
|
||
|
|
val_loader: Optional[Iterable] = None
|
||
|
|
train_iters: int = 0
|
||
|
|
val_iters: int = 0
|
||
|
|
set_epoch: Callable[[int], None] = field(default=_noop_set_epoch)
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = ["DataBundle"]
|