Source code for lhotse.dataset.sampling.stateless

import logging
import random
from pathlib import Path
from typing import Callable, Dict, Generator, Iterable, Optional, Sequence, Tuple, Union

import torch
from cytoolz import compose_left

from lhotse import CutSet, Seconds
from lhotse.cut.set import deserialize_cut
from lhotse.dataset.dataloading import get_rank, get_world_size
from lhotse.dataset.sampling.base import SamplingDiagnostics
from lhotse.lazy import Dillable
from lhotse.serialization import decode_json_line
from lhotse.utils import Pathlike

PathlikeAndScale = Tuple[Pathlike, float]

[docs] class StatelessSampler(, Dillable): """ An infinite and stateless cut sampler that selects data at random from one or more cut manifests. The main idea is to make training resumption easy while guaranteeing the data seen each time by the model is shuffled differently. It discards the notion of an "epoch" and it never finishes iteration. It makes no strong guarantees about avoiding data duplication, but in practice you would rarely see duplicated data. The recommended way to use this sampler is by placing it into a dataloader worker subprocess with Lhotse's :class:``~lhotse.dataset.iterable_dataset.IterableDatasetWrapper``, so that each worker has its own sampler replica that uses a slightly different random seed:: >>> import torch >>> import lhotse >>> dloader = ... lhotse.dataset.iterable_dataset.IterableDatasetWrapper( ... dataset=lhotse.dataset.K2SpeechRecognitionDataset(...), ... sampler=StatelessSampler(...), ... ), ... batch_size=None, ... num_workers=4, ... ) This sampler's design was originally proposed by Dan Povey. For details see: Example 1: Get a non-bucketing :class:``.StatelessSampler``:: >>> sampler = StatelessSampler( ... cuts_paths=["data/cuts_a.jsonl", "data/cuts_b.jsonl"], ... index_path="data/files.idx", ... max_duration=600.0, ... ) Example 2: Get a bucketing :class:``.StatelessSampler``:: >>> sampler = StatelessSampler( ... cuts_paths=["data/cuts_a.jsonl", "data/cuts_b.jsonl"], ... index_path="data/files.idx", ... max_duration=600.0, ... num_buckets=50, ... quadratic_duration=30.0, ... ) Example 3: Get a bucketing :class:``.StatelessSampler`` with scaled weights for each cutset:: >>> sampler = StatelessSampler( ... cuts_paths=[ ... ("data/cuts_a.jsonl", 2.0), ... ("data/cuts_b.jsonl", 1.0), ... ], ... index_path="data/files.idx", ... max_duration=600.0, ... num_buckets=50, ... quadratic_duration=30.0, ... ) .. note:: This sampler works only with uncompressed jsonl manifests, as it creates extra index files with line byte offsets to quickly find and sample JSON lines. This means this sampler will not work with Webdataset and Lhotse Shar data format. :param cuts_paths: Path, or list of paths, or list of tuples of (path, scale) to cutset files. :param index_path: Path to a file that contains the index of all cutsets and their line count (will be auto-created the first time this object is initialized). :param base_seed: Int, user-provided part of the seed used to initialize the RNG for sampling (each node and worker are still going to produce different results). When continuing the training it should be a function of the number of training steps to ensure the model doesn't see identical mini-batches again. :param max_duration: Maximum total number of audio seconds in a mini-batch (dynamic batch size). :param max_cuts: Maximum number of examples in a mini-batch (static batch size). :param num_buckets: If set, enables bucketing (each mini-batch has examples of a similar duration). :param quadratic_duration: If set, adds a penalty term for longer duration cuts. Works well with models that have quadratic time complexity to keep GPU utilization similar when using bucketing. Suggested values are between 30 and 45. """
[docs] def __init__( self, cuts_paths: Union[Pathlike, Iterable[Pathlike], Iterable[PathlikeAndScale]], index_path: Pathlike, base_seed: int, max_duration: Optional[Seconds] = None, max_cuts: Optional[int] = None, num_buckets: Optional[int] = None, quadratic_duration: Optional[Seconds] = None, ) -> None: super().__init__(data_source=None) # Parse file paths and scales self.paths = [] self.scales = [] if isinstance(cuts_paths, (Path, str)): # Single path input self.paths.append(Path(cuts_paths)) self.scales.append(1.0) else: # Multiple path input cuts_paths = list(cuts_paths) self.paths = [] if isinstance(cuts_paths[0], (Path, str)): # Without scales for p in cuts_paths: assert isinstance( p, (Path, str) ), "Mixing paths with and without scales is not allowed." self.paths.append(Path(p)) self.scales.append(1.0) else: for tpl in cuts_paths: # With scales assert len(tpl) == 2, ( f"Expected (path, scale) but got: {tpl} " f"[note: mixing paths with and without scales is not allowed]" ) p, scale = tpl assert isinstance( p, (Path, str) ), f"Path must be a string or Path, got: {p}" assert isinstance( scale, (int, float) ), f"Scale must be an int or float, got: {scale}" self.paths.append(Path(p)) self.scales.append(scale) self.index_path = index_path self.max_duration = max_duration self.max_cuts = max_cuts self.num_buckets = num_buckets self.quadratic_duration = quadratic_duration self.base_seed = base_seed assert any( v is not None for v in (self.max_duration, self.max_cuts) ), "At least one of max_duration or max_cuts has to be set." self.diagnostics = SamplingDiagnostics() self.index = ManifestIndex(self.paths, self.index_path) self.scaled_line_counts = [ lc * scale for lc, scale in zip(self.index.line_counts.values(), self.scales) ] self._transforms = [] # DDP related info self.rank = get_rank() self.world_size = get_world_size()
[docs] def map(self, fn: Callable[[CutSet], CutSet]) -> "StatelessSampler": """Apply ``fn`` to each mini-batch of ``CutSet`` before yielding it.""" self._transforms.append(fn) return self
[docs] def state_dict(self) -> Dict: """Stub state_dict method that returns nothing - this sampler is stateless.""" return {}
[docs] def load_state_dict(self, state_dict: Dict) -> None: """Stub load_state_dict method that does nothing - this sampler is stateless.""" return
def __iter__(self) -> Generator[CutSet, None, None]: from lhotse.dataset import DynamicBucketingSampler, DynamicCutSampler worker_info = worker_id = 0 if worker_info is None else my_id = worker_id + 1000 * self.rank seed = self.base_seed + my_id rng = random.Random(seed) f"[{type(self).__name__}] Initialized sampler RNG with seed {seed} (== base_seed={self.base_seed} + my_id={my_id}) " f"[ddp_rank={self.rank} worker_id={worker_id}]" ) def _inner(): """ Infinite generator of cuts. Each cut is samples in two steps: - first we select a cutset file, weighted by line count (num cuts) - then we randomly select a line from that file using uniform distribution """ n = 0 while True: # Choose a file path = rng.choices(self.paths, self.scaled_line_counts)[0] # Choose a line line_offsets = self.index.line_offsets[path] begin_idx = rng.randrange(len(line_offsets) - 1) begin, end = line_offsets[begin_idx], line_offsets[begin_idx + 1] # Read JSON line and initialize a Cut object with as f: line = - begin) data = decode_json_line(line) cut = deserialize_cut(data) # Update cut ID since the same item may land in a single mini-batch # (note: CutSet prohibits two items sharing the same ID) = f"{}_it{n}" yield cut n += 1 if self.num_buckets is not None and self.num_buckets > 1: inner_sampler = DynamicBucketingSampler( _inner(), max_duration=self.max_duration, max_cuts=self.max_cuts, num_buckets=self.num_buckets, shuffle=False, drop_last=False, quadratic_duration=self.quadratic_duration, world_size=1, rank=0, ) else: inner_sampler = DynamicCutSampler( _inner(), max_duration=self.max_duration, max_cuts=self.max_cuts, shuffle=False, drop_last=False, world_size=1, rank=0, )*self._transforms)) self.diagnostics = inner_sampler.diagnostics yield from inner_sampler
[docs] def get_report(self) -> str: """Returns a string describing the statistics of the sampling process so far.""" return self.diagnostics.get_report()
class ManifestIndex: """ An index of line count and line offset for each cutset manifest. When created for the first time, it writes a .jsonl.idx file for each .jsonl file that contains byte offsets for each line. It also writes a file at ``index_path`` that has the line count and path for each manifest. When this object is instantiated again (e.g. when resuming training), it will just load the contents of existing files from disk. Objects of this class expose two members: ``line_counts: Dict[Path, int]`` and ``line_offsets: Dict[Path, List[int]]`` to simplify working with manifests. :param manifest_paths: A list of paths to cut sets. :param index_path: A path where we should write the line count index (if it doesn't exist). :param force: When true, we'll ignore existing files and reindex the cutsets. """ def __init__( self, manifest_paths: Sequence[Pathlike], index_path: Pathlike, force: bool = False, ) -> None: self.line_counts: Dict[Path, int] = {} self.line_offsets: Dict[Path, Tuple[int]] = {} for p in map(Path, manifest_paths): assert ( p.suffix == ".jsonl" ), f"We only support uncompressed .jsonl files in this sampler, but received: {p}" offset_path = p.with_suffix(".jsonl.idx") if offset_path.is_file() and not force: offsets = self._load(offset_path) else: offsets = self._process(p, offset_path) self.line_counts[p] = len(offsets) self.line_offsets[p] = offsets # Write a single cutset index in format: # <number-of-lines> <cutset-path> # Example: # 10015 data/cuts-part-0001.jsonl # 376101 data/cuts-part-0002.jsonl # 572 data/cuts-part-0003.jsonl if not index_path.is_file() or force: with"w") as index_f: for p, lc in self.line_counts.items(): print(f"{lc} {p}", file=index_f) def _load(self, file_index: Path) -> Tuple[int]: with as f: offsets = tuple(map(int, f)) return offsets def _process(self, manifest: Path, file_index: Path) -> Tuple[int]: # Write line index for each cutset in format <begin-byte> per line, e.g.: # 0 # 214 # 357 # ... offsets = [0] with as cuts_f,"w") as index_f: print(0, file=index_f) line = cuts_f.readline() while line: offsets.append(cuts_f.tell()) print(offsets[-1], file=index_f) line = cuts_f.readline() return tuple(offsets)