Batching¶

Data loading and batching utilities for neural machine translation.

Overview¶

This module provides:

collate_fn: Pads variable-length sequences for batching
BucketBatchSampler: Groups similar-length sequences to minimize padding
create_dataloaders: Convenience function for creating train/val loaders

Quick Start¶

from torch.utils.data import DataLoader
from torchlingo.data_processing import NMTDataset, collate_fn, create_dataloaders

# Manual dataloader
dataset = NMTDataset("data/train.tsv")
loader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

# Or use the convenience function
train_loader, val_loader = create_dataloaders(
    train_file="data/train.tsv",
    val_file="data/val.tsv",
    batch_size=32,
)

API Reference¶

collate_fn¶

collate_fn ¶

collate_fn(batch: List[Tuple[Tensor, Tensor]], pad_idx: Optional[int] = None, config: Optional[Config] = None) -> Tuple[Tensor, Tensor]

Collate a batch of (source, target) tensor pairs with padding.

Pads all sequences in the batch to the same length (the length of the longest sequence) along the sequence dimension.

Parameters:

Name	Type	Description	Default
`batch`	`list[tuple[Tensor, Tensor]]`	List of (src, tgt) pairs where src and tgt are 1D long tensors of variable length.	required
`pad_idx`	`int`	Padding index value. Defaults to config.pad_idx.	`None`
`config`	`Config`	Configuration object. Defaults to get_default_config().	`None`

Returns:

Type	Description
`Tuple[Tensor, Tensor]`	tuple[torch.Tensor, torch.Tensor]: A tuple containing: - src_batch (torch.Tensor): Padded source batch with shape (batch_size, max_src_len). - tgt_batch (torch.Tensor): Padded target batch with shape (batch_size, max_tgt_len).

Notes

Used as the collate_fn for PyTorch DataLoaders. Padding is applied with batch_first=True so sequences are in (batch, seq_len) order.

Source code in src/torchlingo/data_processing/batching.py

def collate_fn(
    batch: List[Tuple[torch.Tensor, torch.Tensor]],
    pad_idx: Optional[int] = None,
    config: Optional[Config] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Collate a batch of (source, target) tensor pairs with padding.

    Pads all sequences in the batch to the same length (the length of the
    longest sequence) along the sequence dimension.

    Args:
        batch (list[tuple[torch.Tensor, torch.Tensor]]): List of (src, tgt)
            pairs where src and tgt are 1D long tensors of variable length.
        pad_idx (int, optional): Padding index value. Defaults to config.pad_idx.
        config (Config, optional): Configuration object. Defaults to
            get_default_config().

    Returns:
        tuple[torch.Tensor, torch.Tensor]: A tuple containing:
            - src_batch (torch.Tensor): Padded source batch with shape
              (batch_size, max_src_len).
            - tgt_batch (torch.Tensor): Padded target batch with shape
              (batch_size, max_tgt_len).

    Notes:
        Used as the collate_fn for PyTorch DataLoaders. Padding is applied
        with batch_first=True so sequences are in (batch, seq_len) order.
    """
    cfg = config if config is not None else get_default_config()
    pad_idx = pad_idx if pad_idx is not None else cfg.pad_idx

    src_batch, tgt_batch = zip(*batch)
    src_batch_t = pad_sequence(list(src_batch), batch_first=True, padding_value=pad_idx)
    tgt_batch_t = pad_sequence(list(tgt_batch), batch_first=True, padding_value=pad_idx)
    return src_batch_t, tgt_batch_t

BucketBatchSampler¶

BucketBatchSampler ¶

BucketBatchSampler(dataset: NMTDataset, batch_size: int, bucket_boundaries: Optional[List[int]] = None)

Bases: Sampler

Groups dataset samples into length buckets to minimize padding in batches.

Assigns each sample to a bucket based on its source sequence length. When iterated, yields batches where all samples come from the same bucket, reducing padding waste. Bucket boundaries are automatically computed from the data distribution if not provided.

Bucketing is particularly effective for datasets with large sequence length variance, e.g., 10-500 tokens. It reduces computational waste from padding longer samples with many pad tokens.

Parameters:

Name	Type	Description	Default
`dataset`	`NMTDataset`	The dataset to sample from.	required
`batch_size`	`int`	Number of samples per batch.	required
`bucket_boundaries`	`list[int]`	Bucket upper boundaries in ascending order. Defaults to None, in which case boundaries are computed automatically from sequence length percentiles.	`None`

Attributes:

Name	Type	Description
`buckets`	`list[list[int]]`	Per-bucket list of sample indices.
`num_batches`	`int`	Total number of batches that will be yielded.
`bucket_boundaries`	`list[int]`	Upper length boundary for each bucket.

Raises:

Type	Description
`ValueError`	If no bucket contains at least batch_size samples, so the sampler would yield zero batches (incomplete batches are dropped). This typically happens when batch_size is larger than the dataset; use a smaller batch_size or provide more data.

Parameters:

Name	Type	Description	Default
`dataset`	`NMTDataset`	The dataset to sample from.	required
`batch_size`	`int`	Number of samples per batch.	required
`bucket_boundaries`	`list[int]`	Pre-specified bucket boundaries. If None, computed automatically. Defaults to None.	`None`

Source code in src/torchlingo/data_processing/batching.py

def __init__(
    self,
    dataset: NMTDataset,
    batch_size: int,
    bucket_boundaries: Optional[List[int]] = None,
) -> None:
    """Initialize the sampler and assign samples to buckets.

    Args:
        dataset (NMTDataset): The dataset to sample from.
        batch_size (int): Number of samples per batch.
        bucket_boundaries (list[int], optional): Pre-specified bucket
            boundaries. If None, computed automatically. Defaults to None.
    """
    self.dataset = dataset
    self.batch_size = batch_size
    if bucket_boundaries is None:
        bucket_boundaries = self._calculate_bucket_boundaries()
    self.bucket_boundaries = sorted(bucket_boundaries)
    self.buckets: List[List[int]] = [[] for _ in range(len(bucket_boundaries) + 1)]
    for idx in range(len(dataset)):
        src_tensor, _ = dataset[idx]
        src_len = len(src_tensor)
        b = self._get_bucket(src_len)
        self.buckets[b].append(idx)
    self._compute_num_batches()
    if self.num_batches == 0:
        largest_bucket = max((len(b) for b in self.buckets), default=0)
        raise ValueError(
            f"BucketBatchSampler would yield 0 batches: the dataset has "
            f"{len(dataset)} example(s) and no bucket contains at least "
            f"batch_size={batch_size} of them (largest bucket has "
            f"{largest_bucket}). Incomplete batches are dropped, so "
            f"training would silently run 0 steps. Use a smaller "
            f"batch_size (at most {largest_bucket} for this data) or "
            f"provide more training data."
        )

iter ¶

__iter__() -> Iterator[List[int]]

Iterate over batches of sample indices.

Yields batches by: 1. Shuffling samples within each bucket independently. 2. Creating batches from each bucket (discarding incomplete batches). 3. Shuffling the list of batches. 4. Yielding batches of sample indices.

Yields:

Type	Description
`List[int]`	list[int]: Each batch is a list of sample indices, length = batch_size.

Notes

Incomplete batches (fewer than batch_size samples) are dropped.
Batch order is randomized for better training dynamics.
Within-bucket samples are shuffled for better gradient estimates.

Source code in src/torchlingo/data_processing/batching.py

def __iter__(self) -> Iterator[List[int]]:
    """Iterate over batches of sample indices.

    Yields batches by:
    1. Shuffling samples within each bucket independently.
    2. Creating batches from each bucket (discarding incomplete batches).
    3. Shuffling the list of batches.
    4. Yielding batches of sample indices.

    Yields:
        list[int]: Each batch is a list of sample indices, length = batch_size.

    Notes:
        - Incomplete batches (fewer than batch_size samples) are dropped.
        - Batch order is randomized for better training dynamics.
        - Within-bucket samples are shuffled for better gradient estimates.
    """
    for bucket in self.buckets:
        random.shuffle(bucket)
    all_batches: List[List[int]] = []
    for bucket in self.buckets:
        for i in range(
            0, len(bucket) - len(bucket) % self.batch_size, self.batch_size
        ):
            all_batches.append(bucket[i : i + self.batch_size])
    random.shuffle(all_batches)
    for batch in all_batches:
        yield batch

len ¶

__len__() -> int

Return the number of complete batches.

Returns:

Name	Type	Description
`int`	`int`	Total number of batches that will be yielded. Incomplete batches (fewer than batch_size samples) are not counted.

Source code in src/torchlingo/data_processing/batching.py

def __len__(self) -> int:
    """Return the number of complete batches.

    Returns:
        int: Total number of batches that will be yielded. Incomplete
            batches (fewer than batch_size samples) are not counted.
    """
    return self.num_batches

create_dataloaders¶

create_dataloaders ¶

create_dataloaders(train_file: Union[Path, NMTDataset], val_file: Optional[Union[Path, NMTDataset]] = None, batch_size: Optional[int] = None, num_workers: Optional[int] = None, use_sentencepiece: bool = False, sp_model_path: Optional[str] = None, sp_tgt_model_path: Optional[str] = None, use_bucketing: bool = False, bucket_boundaries: Optional[List[int]] = None, device: Optional[str] = None, pad_idx: Optional[int] = None, config: Optional[Config] = None) -> Tuple[DataLoader, Optional[DataLoader], BaseVocab, BaseVocab]

Create PyTorch DataLoaders for training and validation.

Constructs train and optional validation DataLoaders from structured data files. Supports two tokenization modes: simple token-based (SimpleVocab) or SentencePiece subword. Optionally uses bucketing by sequence length to minimize padding overhead.

Data must be in structured format (TSV/CSV/JSON/Parquet) with configurable source and target columns. Parallel .txt files should be merged into a single file first (see preprocessing.base.parallel_txt_to_dataframe).

Parameters:

Name	Type	Description	Default
`train_file`	`Path \| NMTDataset`	Path to training data file or a prebuilt NMTDataset instance.	required
`val_file`	`Path \| NMTDataset`	Path to validation data file or a prebuilt NMTDataset instance. If None, no validation loader is created. Defaults to None.	`None`
`batch_size`	`int`	Batch size for loaders. Defaults to config.batch_size.	`None`
`num_workers`	`int`	Number of worker processes for data loading. Defaults to config.num_workers.	`None`
`use_sentencepiece`	`bool`	If True, use SentencePiece models for tokenization. If False, use simple token-based vocabularies. Defaults to False.	`False`
`sp_model_path`	`str`	Path to SentencePiece source model file. Required if use_sentencepiece=True. Defaults to config.sentencepiece_src_model.	`None`
`sp_tgt_model_path`	`str`	Path to SentencePiece target model file. If None with use_sentencepiece=True, uses the configured target path or sp_model_path when identical. Defaults to None.	`None`
`use_bucketing`	`bool`	If True, use BucketBatchSampler to group sequences by length. Reduces padding overhead for datasets with variable sequence lengths. Defaults to False.	`False`
`bucket_boundaries`	`list[int]`	Pre-specified bucket boundaries for BucketBatchSampler. If None with use_bucketing=True, boundaries are computed automatically from the training data. Defaults to None.	`None`
`device`	`str`	Device string (e.g., 'cuda', 'cpu'). Used for pin_memory optimization. Defaults to config.device.	`None`
`pad_idx`	`int`	Padding index value. Defaults to config.pad_idx.	`None`
`config`	`Config`	Configuration object. Defaults to get_default_config().	`None`

Returns:

Name	Type	Description
`tuple`	`Tuple[DataLoader, Optional[DataLoader], BaseVocab, BaseVocab]`	A 4-tuple containing: - train_loader (DataLoader): Training data loader. - val_loader (DataLoader or None): Validation data loader, or None if val_file is not provided. - src_vocab (BaseVocab): Source vocabulary. - tgt_vocab (BaseVocab): Target vocabulary.

Raises:

Type	Description
`AssertionError`	If use_sentencepiece=True but sp_model_path is None.
`ValueError`	If data files have incorrect format or missing columns.
`ValueError`	If the training loader would yield 0 batches, e.g. when use_bucketing=True and batch_size is larger than every length bucket (incomplete batches are dropped), or when the training dataset is empty. Use a smaller batch_size or more data.

Examples:

>>> train_loader, val_loader, src_vocab, tgt_vocab = create_dataloaders(
...     train_file="data/train.tsv",
...     val_file="data/val.tsv",
...     batch_size=32,
...     use_bucketing=True
... )
>>> for src_batch, tgt_batch in train_loader:
...     print(src_batch.shape, tgt_batch.shape)
...     break
torch.Size([32, 45]) torch.Size([32, 48])

>>> train_ds = NMTDataset("data/train.tsv")
>>> val_ds = NMTDataset("data/val.tsv", src_vocab=train_ds.src_vocab, tgt_vocab=train_ds.tgt_vocab)
>>> train_loader, val_loader, src_vocab, tgt_vocab = create_dataloaders(
...     train_file=train_ds,
...     val_file=val_ds,
...     batch_size=32,
... )

Notes

DataLoader pin_memory is enabled for CUDA devices to speed up transfers.
Training loader always shuffles (or uses bucket-based shuffling).
Validation loader does not shuffle.
With use_bucketing=True, incomplete batches are dropped; without bucketing, the final (possibly smaller) batch is kept.

Source code in src/torchlingo/data_processing/batching.py

def create_dataloaders(
    train_file: Union[Path, NMTDataset],
    val_file: Optional[Union[Path, NMTDataset]] = None,
    batch_size: Optional[int] = None,
    num_workers: Optional[int] = None,
    use_sentencepiece: bool = False,
    sp_model_path: Optional[str] = None,
    sp_tgt_model_path: Optional[str] = None,
    use_bucketing: bool = False,
    bucket_boundaries: Optional[List[int]] = None,
    device: Optional[str] = None,
    pad_idx: Optional[int] = None,
    config: Optional[Config] = None,
) -> Tuple[DataLoader, Optional[DataLoader], BaseVocab, BaseVocab]:
    """Create PyTorch DataLoaders for training and validation.

    Constructs train and optional validation DataLoaders from structured data
    files. Supports two tokenization modes: simple token-based (SimpleVocab) or
    SentencePiece subword. Optionally uses bucketing by sequence length to
    minimize padding overhead.

    Data must be in structured format (TSV/CSV/JSON/Parquet) with configurable
    source and target columns. Parallel .txt files should be merged into a
    single file first (see preprocessing.base.parallel_txt_to_dataframe).

    Args:
        train_file (Path | NMTDataset): Path to training data file or a prebuilt
            NMTDataset instance.
        val_file (Path | NMTDataset, optional): Path to validation data file or
            a prebuilt NMTDataset instance. If None, no validation loader is
            created. Defaults to None.
        batch_size (int, optional): Batch size for loaders. Defaults to
            config.batch_size.
        num_workers (int, optional): Number of worker processes for data loading.
            Defaults to config.num_workers.
        use_sentencepiece (bool, optional): If True, use SentencePiece models for
            tokenization. If False, use simple token-based vocabularies. Defaults
            to False.
        sp_model_path (str, optional): Path to SentencePiece source model file.
            Required if use_sentencepiece=True. Defaults to
            config.sentencepiece_src_model.
        sp_tgt_model_path (str, optional): Path to SentencePiece target model file.
            If None with use_sentencepiece=True, uses the configured target path
            or sp_model_path when identical. Defaults to None.
        use_bucketing (bool, optional): If True, use BucketBatchSampler to group
            sequences by length. Reduces padding overhead for datasets with
            variable sequence lengths. Defaults to False.
        bucket_boundaries (list[int], optional): Pre-specified bucket boundaries
            for BucketBatchSampler. If None with use_bucketing=True, boundaries
            are computed automatically from the training data. Defaults to None.
        device (str, optional): Device string (e.g., 'cuda', 'cpu'). Used for
            pin_memory optimization. Defaults to config.device.
        pad_idx (int, optional): Padding index value. Defaults to config.pad_idx.
        config (Config, optional): Configuration object. Defaults to
            get_default_config().

    Returns:
        tuple: A 4-tuple containing:
            - train_loader (DataLoader): Training data loader.
            - val_loader (DataLoader or None): Validation data loader, or None
              if val_file is not provided.
            - src_vocab (BaseVocab): Source vocabulary.
            - tgt_vocab (BaseVocab): Target vocabulary.

    Raises:
        AssertionError: If use_sentencepiece=True but sp_model_path is None.
        ValueError: If data files have incorrect format or missing columns.
        ValueError: If the training loader would yield 0 batches, e.g. when
            use_bucketing=True and batch_size is larger than every length
            bucket (incomplete batches are dropped), or when the training
            dataset is empty. Use a smaller batch_size or more data.

    Examples:
        >>> train_loader, val_loader, src_vocab, tgt_vocab = create_dataloaders(
        ...     train_file="data/train.tsv",
        ...     val_file="data/val.tsv",
        ...     batch_size=32,
        ...     use_bucketing=True
        ... )
        >>> for src_batch, tgt_batch in train_loader:
        ...     print(src_batch.shape, tgt_batch.shape)
        ...     break
        torch.Size([32, 45]) torch.Size([32, 48])

        >>> train_ds = NMTDataset("data/train.tsv")
        >>> val_ds = NMTDataset("data/val.tsv", src_vocab=train_ds.src_vocab, tgt_vocab=train_ds.tgt_vocab)
        >>> train_loader, val_loader, src_vocab, tgt_vocab = create_dataloaders(
        ...     train_file=train_ds,
        ...     val_file=val_ds,
        ...     batch_size=32,
        ... )

    Notes:
        - DataLoader pin_memory is enabled for CUDA devices to speed up transfers.
        - Training loader always shuffles (or uses bucket-based shuffling).
        - Validation loader does not shuffle.
        - With use_bucketing=True, incomplete batches are dropped; without
          bucketing, the final (possibly smaller) batch is kept.
    """
    cfg = config if config is not None else get_default_config()
    batch_size = batch_size if batch_size is not None else cfg.batch_size
    num_workers = num_workers if num_workers is not None else cfg.num_workers
    device = device if device is not None else cfg.device
    pad_idx = pad_idx if pad_idx is not None else cfg.pad_idx

    # Create collate function with resolved pad_idx
    collate = partial(collate_fn, pad_idx=pad_idx)

    if isinstance(train_file, NMTDataset):
        train_dataset = train_file
        src_vocab = train_dataset.src_vocab
        tgt_vocab = train_dataset.tgt_vocab
        assert src_vocab is not None and tgt_vocab is not None, (
            "Provided train_dataset must have src_vocab and tgt_vocab initialized"
        )
    else:
        train_data = train_file
        if use_sentencepiece:
            sp_model_path = (
                sp_model_path
                if sp_model_path is not None
                else cfg.sentencepiece_src_model
            )
            sp_tgt_model_path = (
                sp_tgt_model_path
                if sp_tgt_model_path is not None
                else cfg.sentencepiece_tgt_model
            )
            assert sp_model_path is not None, (
                "Must provide sp_model_path when using SentencePiece"
            )
            src_vocab = SentencePieceVocab(sp_model_path)
            if sp_tgt_model_path and sp_tgt_model_path != sp_model_path:
                tgt_vocab = SentencePieceVocab(sp_tgt_model_path)
            else:
                tgt_vocab = src_vocab
        else:
            tmp = NMTDataset(train_data)
            src_vocab = tmp.src_vocab
            tgt_vocab = tmp.tgt_vocab
            assert src_vocab is not None and tgt_vocab is not None
        train_dataset = NMTDataset(train_data, src_vocab=src_vocab, tgt_vocab=tgt_vocab)

    val_dataset: Optional[NMTDataset]
    if val_file is not None:
        if isinstance(val_file, NMTDataset):
            val_dataset = val_file
            # Ensure vocab objects are shared with the training dataset
            val_dataset.src_vocab = src_vocab
            val_dataset.tgt_vocab = tgt_vocab
        else:
            val_dataset = NMTDataset(val_file, src_vocab=src_vocab, tgt_vocab=tgt_vocab)
    else:
        val_dataset = None
    if use_bucketing:
        train_sampler = BucketBatchSampler(
            train_dataset, batch_size=batch_size, bucket_boundaries=bucket_boundaries
        )
        train_loader = DataLoader(
            train_dataset,
            batch_sampler=train_sampler,
            num_workers=num_workers,
            collate_fn=collate,
            pin_memory=True if device == "cuda" else False,
        )
    else:
        train_loader = DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=num_workers,
            collate_fn=collate,
            pin_memory=True if device == "cuda" else False,
        )
    if len(train_loader) == 0:
        raise ValueError(
            f"Training loader would yield 0 batches (dataset has "
            f"{len(train_dataset)} example(s), batch_size={batch_size}). "
            f"Training would silently run 0 steps. Use a smaller batch_size "
            f"or provide more training data."
        )
    if val_dataset is not None:
        val_loader = DataLoader(
            val_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
            collate_fn=collate,
            pin_memory=True if device == "cuda" else False,
        )
    else:
        val_loader = None
    return train_loader, val_loader, src_vocab, tgt_vocab

Examples¶

Basic Batching¶

from torch.utils.data import DataLoader
from torchlingo.data_processing import NMTDataset, collate_fn

dataset = NMTDataset("data/train.tsv")

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    collate_fn=collate_fn,
    num_workers=4,  # Parallel data loading
)

for src_batch, tgt_batch in loader:
    print(f"Source: {src_batch.shape}")  # [32, max_src_len]
    print(f"Target: {tgt_batch.shape}")  # [32, max_tgt_len]
    break

Bucketed Batching¶

Bucketing groups sequences of similar length together, reducing wasted computation on padding:

from torch.utils.data import DataLoader
from torchlingo.data_processing import NMTDataset, collate_fn, BucketBatchSampler

dataset = NMTDataset("data/train.tsv")

sampler = BucketBatchSampler(dataset, batch_size=32)

loader = DataLoader(
    dataset,
    batch_sampler=sampler,  # Note: batch_sampler, not batch_size
    collate_fn=collate_fn,
)

Using create_dataloaders¶

The easiest way to create properly configured loaders:

from torchlingo.data_processing import create_dataloaders

train_loader, val_loader = create_dataloaders(
    train_file="data/train.tsv",
    val_file="data/val.tsv",
    batch_size=32,
    shuffle_train=True,
    use_bucketing=True,  # Optional: enable bucketing
)

With Shared Vocabularies¶

from torchlingo.data_processing import NMTDataset, create_dataloaders

# Build vocab on training data
train_dataset = NMTDataset("data/train.tsv")

# Pass to create_dataloaders
train_loader, val_loader = create_dataloaders(
    train_file="data/train.tsv",
    val_file="data/val.tsv",
    src_vocab=train_dataset.src_vocab,
    tgt_vocab=train_dataset.tgt_vocab,
    batch_size=32,
)

How Padding Works¶

The collate_fn pads sequences to the maximum length in each batch:

Before padding:
  Sample 1: [2, 5, 6, 3]         # length 4
  Sample 2: [2, 7, 8, 9, 10, 3]  # length 6
  Sample 3: [2, 11, 3]           # length 3

After padding (to length 6):
  Sample 1: [2, 5, 6, 3, 0, 0]   # padded with 0s
  Sample 2: [2, 7, 8, 9, 10, 3]  # no padding needed
  Sample 3: [2, 11, 3, 0, 0, 0]  # padded with 0s

How Bucketing Works¶

BucketBatchSampler assigns sequences to buckets based on length:

Bucket 0 (len ≤ 10):  [sample_3, sample_7, sample_12, ...]
Bucket 1 (len ≤ 20):  [sample_1, sample_5, sample_9, ...]
Bucket 2 (len ≤ 40):  [sample_2, sample_8, ...]
Bucket 3 (len > 40):  [sample_4, sample_6, ...]

Batches are created from samples within the same bucket, so sequences in a batch have similar lengths.

Benefits of Bucketing¶

Without Bucketing	With Bucketing
Random mixing of lengths	Similar lengths grouped
More padding tokens	Minimal padding
Wasted computation	Efficient computation

Custom Bucket Boundaries¶

sampler = BucketBatchSampler(
    dataset,
    batch_size=32,
    bucket_boundaries=[10, 25, 50, 100],  # Custom boundaries
)

Configuration¶

All batching functions respect the Config object:

from torchlingo.config import Config

config = Config(
    batch_size=64,
    pad_idx=0,
)

train_loader, val_loader = create_dataloaders(
    train_file="data/train.tsv",
    val_file="data/val.tsv",
    config=config,
)

Performance Tips¶

Use num_workers for parallel data loading:
```
DataLoader(..., num_workers=4)
```
Enable bucketing for datasets with variable lengths:
```
create_dataloaders(..., use_bucketing=True)
```
Pin memory for GPU training:
```
DataLoader(..., pin_memory=True)
```
Adjust batch size based on GPU memory:
Larger batches = faster training but more memory
Start with 32, increase until you run out of memory

Batching¶

Overview¶

Quick Start¶

API Reference¶

collate_fn¶

collate_fn ¶

BucketBatchSampler¶

BucketBatchSampler ¶

__iter__ ¶

__len__ ¶

create_dataloaders¶

create_dataloaders ¶

Examples¶

Basic Batching¶

Bucketed Batching¶

Using create_dataloaders¶

With Shared Vocabularies¶

How Padding Works¶

How Bucketing Works¶

Benefits of Bucketing¶

Custom Bucket Boundaries¶

Configuration¶

Performance Tips¶

iter ¶

len ¶