Skip to content

mix

Dataset wrapper that mixes each sample with a random partner via a pairwise transform.

Classes:

  • MixDataset –

    Draws a second random sample and applies a pairwise mix transform.

MixDataset

MixDataset(
    dataset: Dataset,
    mix: Callable[
        [Dict[str, Any], Dict[str, Any]], Dict[str, Any]
    ],
    transform: Optional[
        Callable[[Dict[str, Any]], Dict[str, Any]]
    ] = None,
    generator: Optional[Generator] = None,
)

Bases: Dataset

Draws a second random sample and applies a pairwise mix transform.

Each source sample is produced by the wrapped dataset (its own transform runs first), then a partner sample is drawn at a random index and merged via mix(data, other). Pairwise mixes such as Mix3D, LaserMix, and PolarMix fit this contract; the mix's own p decides how often the merge actually happens.

Parameters:

  • dataset (Dataset) –

    The wrapped dataset; its own transform runs per source sample before mixing.

  • mix (Callable[[Dict[str, Any], Dict[str, Any]], Dict[str, Any]]) –

    Pairwise transform called as mix(data, other).

  • transform (Optional[Callable[[Dict[str, Any]], Dict[str, Any]]], default: None ) –

    Optional per-sample transform applied after mixing.

  • generator (Optional[Generator], default: None ) –

    Optional torch.Generator for the partner index. Every DataLoader worker receives a replica of this generator, so inside a worker the partner indices are drawn from a dedicated stream seeded with the generator's initial seed plus the worker id.