espnet3.components.data.dataloader.DataLoaderBuilder

source

class espnet3.components.data.dataloader.DataLoaderBuilder(dataset, config, collate_fn, num_device: int, epoch: int)

Bases: object

Builder class for constructing training and validation DataLoaders in ESPnet3.

This class provides a unified interface for setting up PyTorch or ESPnet-specific DataLoaders based on the configuration. It supports advanced features such as:

• Custom collate functions (e.g., CommonCollateFn)
• Sequence-based batch sampling with batch_bins or batch_size
• Dynamic handling of DDP-compatible iteration strategies

Set dataloader.<mode>.batch_distribution to per_gpu (the default) to distribute complete batches, or global to split each supplied batch across ranks. Global mode requires unsharded datasets and at least one sample per rank in every batch. Factories that regenerate batches must preserve the distribution when doing so. For CategoryIterFactory, set sampler_args.distributed: true for multi-rank global batching.

Typically used by ESPnetLightningModule to instantiate training and validation DataLoaders with consistent behavior across single-GPU, multi-GPU, and distributed training.

• Parameters:
  • dataset (torch.utils.data.Dataset) – Dataset instance wrapped by DataLoader.
  • config (DictConfig) – Full training configuration (e.g., from OmegaConf).
  • collate_fn (Callable) – Function to collate individual samples into mini-batches.
  • num_device (int) – Number of devices in use (e.g., GPUs or nodes).
  • epoch (int) – Current epoch number. Used to reseed samplers deterministically.

Example

builder = DataLoaderBuilder(
    dataset=train_dataset,
    config=config,
    collate_fn=collate_fn,
    num_device=4,
    epoch=3,
)
train_loader = builder.build(mode="train")

Initialize DataLoaderBuilder object.

build(mode: str)

Build and return a DataLoader for the specified mode (“train” or “valid”).

This method supports two modes of operation depending on the configuration: (1) ESPnet-style iterator: Custom sampling and iteration logic. (2) Standard PyTorch DataLoader: Simpler use-case with fixed batch size.

The selection depends on whether config.dataloader.<mode>.iter_factory is defined.

• Parameters:mode (str) – One of {“train”, “valid”} indicating which dataloader to build.
• Returns: Configured DataLoader instance for training or : validation.
• Return type: torch.utils.data.DataLoader

Config-driven branching logic: : Case 1: ESPnet SequenceIterFactory is used Example config: <br/>

``

<br/> ` dataloader: train: <br/>

iter_factory: _target_: espnet2.iterators.sequence_iter_factory.SequenceIterFactory shuffle: true collate_fn: ${dataloader.collate_fn} batches: <br/>

_target_: espnet2.samplers.build_batch_sampler.build_batch_sampler type: numel shape_files: batch_bins: 4000000 min_batch_size: 2 <br/>

``

<br/>

`

<br/> Case 2: Standard PyTorch DataLoader is used Example config: <br/>

``

<br/> ` dataloader: train: <br/>

iter_factory: # Set iter_factory to None batch_size: 4 num_workers: 2 shuffle: true <br/>

``

<br/>

`

• Raises:
  • ValueError – If batch_distribution is not per_gpu or global.
  • RuntimeError – If global batching cannot preserve the shared batch plan.

Example

s

Set config.dataloader.train.batch_distribution = “global” and use an iter_factory to split the sampler’s batches across ranks. Omitting the setting preserves complete per-GPU batches.