Class DataIteratorFactory

espnet2/speechlm/dataloader/iterator.py:21–345

Factory for creating data iterators for SpeechLM training. This class manages batching, data sharding across GPUs, and provides DataLoader instances for training with support for endless epochs. Features: - Supports multiple tasks and datasets with resampling factors -

class DataIteratorFactory:
    """Factory for creating data iterators for SpeechLM training.

    This class manages batching, data sharding across GPUs, and provides
    DataLoader instances for training with support for endless epochs.

    Features:
        - Supports multiple tasks and datasets with resampling factors
        - Bucket or pack batching strategies
        - Distributed training with automatic batch synchronization
        - Deterministic shuffling with configurable seeds
        - State saving/loading for training resumption

    Args:
        unregistered_specifier: Space-separated unregistered data specs.
            Format: "task:name:data_json[:factor]"
            Example: "asr:librispeech:train.json:2.0"
        registered_specifier: Space-separated registered data specs.
            Format: "task:name[:factor]"
            Example: "tts:ljspeech:1.5"
        stats_dir: Directory containing statistics files (str or Path).
            Each file should be named "stats_{task}_{data_name}.jsonl"
        collate_fn: Optional collate function for DataLoader.
        loader_state: Optional saved state dict to restore from.
        batchfy_method: Batching method ("bucket" or "pack").
        batch_size: Maximum tokens per batch.
        num_workers: Number of DataLoader workers.
        rank: GPU rank for distributed training (0-indexed).
        world_size: Total number of GPUs in distributed training.
        shuffle: Whether to shuffle batches.
        seed: Random seed for reproducibility.

    Example:
        >>> factory = DataIteratorFactory(
        ...     unregistered_specifier="asr:libri:train.json:2.0",
        ...     registered_specifier="tts:lj:1.0",
        ...     stats_dir="/path/to/stats",
        ...     batch_size=10000,
        ...     shuffle=True,
        ... )
        >>> loader = factory.get_iterator(global_step=0, length=100)
        >>> for batch in loader:
        ...     # Training loop
        ...     pass
    """

    def __init__(
        self,
        unregistered_specifier: str = "",
        registered_specifier: str = "",
        stats_dir: Union[str, Path] = None,
        loader_state: Optional[Path] = None,
        collate_fn: Optional[Callable] = None,
        batchfy_method: str = "bucket",
        batch_size: int = 1000,
        num_workers: int = 4,
        rank: int = 0,
        world_size: int = 1,