espnet2.layers.batch_wav_augment.BatchWavAugment

source

class espnet2.layers.batch_wav_augment.BatchWavAugment(noise_scp: str | None = None, noise_apply_prob: float = 1.0, noise_db_range: str = '13_15', rir_scp: str | None = None, rir_apply_prob: float = 1.0, rir_shift: bool = True, snr_reference: str = 'dry', noise_rir_apply_prob: float = 0.0, context: int = 0, fs: int = 16000, silence_threshold: float = 0.01, frame_length: int = 1024, frame_shift: int = 512, keep_on_device: bool = False)

Bases: Module

Add reverberation and noise to a batch of single-channel waveforms.

• Parameters:
  • noise_scp – Kaldi-style list <id> <path> (or <path>) of noise recordings, or None for no noise. Each channel of a multi-channel recording is used as a separate noise.
  • noise_apply_prob – Probability of adding noise to each utterance.
  • noise_db_range – SNR range in dB, "low_high" or a single value. Quote a range in YAML ('13_15'): unquoted, it is the integer 1315.
  • rir_scp – Kaldi-style list <id> <path> (or <path>) of room impulse responses, or None for no reverberation. Each channel of a multi-channel file is used as a separate impulse response.
  • rir_apply_prob – Probability of convolving each utterance with an impulse response.
  • rir_shift – Shift the reverberated speech back by the position of the impulse response’s peak (argmax |rir|, the largest absolute value, not an estimate of the direct path; each channel of a multi-channel file is its own impulse response with its own peak), so that it stays aligned with the dry speech and its labels, as Kaldi’s wav-reverberate --shift-output (default), Lhotse and SpeechBrain do. False keeps the first samples of the convolution, as CommonPreprocessor does, which delays the speech by the peak position (e.g. the propagation delay of a measured or simulated impulse response).
  • snr_reference – The speech power the SNR refers to. “dry”: the non-silent power of the speech before reverberation (the reverberated speech is scaled back to it), as CommonPreprocessor does. “early”: the non-silent power of the early reverberation, the speech convolved with the impulse response from 1 ms before its peak to 50 ms after it, as Kaldi’s wav-reverberate does; the late reverberation then counts as distortion, not as speech.
  • noise_rir_apply_prob – Probability of convolving the noise of each utterance with an impulse response of its own (drawn independently of the speech’s), so that a point-source noise is reverberated as Kaldi’s reverberate_data_dir.py does. Needs rir_scp.
  • context – Number of leading samples of each row that only give the reverberation its history: the row is augmented as a whole, the powers and the peak are measured after the context, and the first context samples are removed (also in eval mode, without augmentation), so the labels cover the rest. Give the data loader windows that start context samples early (zeros before the start of a recording); without it the first samples of a window are reverberated as if nothing preceded them. 0 for none.
  • fs – Sampling rate of the speech; noise and impulse responses at another rate are resampled when they are loaded.
  • silence_threshold – A frame is non-silent when its power exceeds this fraction of the mean frame power (detect_non_silence).
  • frame_length – Frame length (samples) of the silence detection.
  • frame_shift – Frame shift (samples) of the silence detection.
  • keep_on_device – Copy the whole banks to the device of the batch once, instead of copying what each batch needs from host memory. Faster, but uses 2 bytes of device memory per noise sample (11.4 GB for 99 h at 16 kHz) and 4 bytes per impulse-response sample.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

apply(speech: Tensor, valid: Tensor, draws: Dict) → Tensor

Apply the choices of draws to speech (B, T); valid marks the samples.

This is the deterministic part of forward(): given the same choices it gives the same result as CommonPreprocessor.

bank(kind: str, device: device) → Tuple[Tensor, ...]

Return (bank, starts, lengths) of “noise” or “rir” on device.

The first call loads the list; each device gets one copy.

draw(lengths: Tensor) → Dict

Draw the random choices of a batch: what to apply, which and where.

lengths (B,) holds the number of samples of each utterance. Returns a dict of (B,) tensors: rir_apply, rir_which, noise_apply, noise_which and snr (dB) on the CPU, and noise_offset on the device of lengths. As in CommonPreprocessor, the noise is cut for each utterance by its own length: a noise at least as long as the utterance starts at any offset where the whole utterance fits (n_len - length + 1 offsets); a shorter noise starts at any offset and repeats.

extra_repr() → str

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(speech: Tensor, speech_lengths: Tensor | None = None) → Tuple[Tensor, Tensor | None]

Add reverberation and noise to speech (B, T) in training mode.

Returns the augmented speech and speech_lengths. With context the first context samples of each row are removed afterwards, in training and in eval mode.

remove_context(speech: Tensor, speech_lengths: Tensor | None = None) → Tuple[Tensor, Tensor | None]

Remove the first context samples of each row (nothing if 0).

Raises ValueError when the batch has no sample after the context; a row whose length does not exceed the context fails an assertion on the device (checked without waiting for it).