Spaces:

rjattt
/

echofocus

Paused

echofocus / vanta /data /dynamic_dataset.py

Upload folder using huggingface_hub

d67c04d verified 3 days ago

3.3 kB

"""On-the-fly mixing dataset — the fix for the memorization problem.

The old pipeline froze 20k mixtures to disk (build_dataset.py) and trained on
that same fixed set every epoch, so the model memorized it (train SI-SDR climbed
while val stayed pinned at ~+1 dB). This dataset instead calls the synthesis
engine *live*: every __getitem__ produces a brand-new random mixture, so the
model never sees the same example twice and has to learn to separate voices
rather than memorize files.

Only the *training* set should be dynamic. Keep validation on a fixed manifest
(VantaDataset) so val numbers are comparable across runs.

Freshness across epochs:
    We reseed the per-process mixer from torch.initial_seed(), which PyTorch
    varies per (epoch, worker). With num_workers=0 the mixer is created once and
    its RNG simply keeps advancing across epochs — either way, no epoch ever
    repeats the previous epoch's mixtures.
"""

from __future__ import annotations

import torch
from torch.utils.data import Dataset, get_worker_info

from vanta.data.indexer import NoiseIndex, RirIndex, SpeakerIndex
from vanta.data.synthesize import MixConfig, Mixer


class DynamicMixDataset(Dataset):
    def __init__(
        self,
        mix_cfg: MixConfig,
        speakers: SpeakerIndex,
        noise: NoiseIndex | None,
        rirs: RirIndex | None,
        epoch_size: int,
        base_seed: int = 0,
    ):
        self.mix_cfg = mix_cfg
        self.speakers = speakers
        self.noise = noise
        self.rirs = rirs
        self.epoch_size = int(epoch_size)
        self.base_seed = int(base_seed)
        # Lazily built per process so each DataLoader worker gets its own RNG.
        self._mixer: Mixer | None = None

    def __len__(self) -> int:
        # Nominal length = how many fresh mixtures make up one "epoch". Purely a
        # cadence knob (when to validate / checkpoint); the data is unbounded.
        return self.epoch_size

    def _mixer_for_process(self) -> Mixer:
        if self._mixer is None:
            info = get_worker_info()
            if info is None:
                # Main process (num_workers=0): seed once; the RNG advances
                # across epochs because this object persists, so data stays fresh.
                seed = self.base_seed
            else:
                # Worker process: torch.initial_seed() is set by PyTorch per
                # (epoch, worker), so reseeding from it gives fresh data every
                # epoch even when workers are torn down and recreated.
                seed = torch.initial_seed() % (2**31)
            self._mixer = Mixer(
                self.mix_cfg,
                speakers=self.speakers,
                noise=self.noise,
                rirs=self.rirs,
                seed=seed,
            )
        return self._mixer

    def __getitem__(self, idx: int) -> dict:
        # idx is ignored — every call is a new random draw.
        r = self._mixer_for_process().mix()
        return {
            "mixture": torch.from_numpy(r.mixture),
            "target": torch.from_numpy(r.target),
            "enrollment": torch.from_numpy(r.enrollment),
            "target_speaker": r.meta["target_speaker"],
        }

	"""On-the-fly mixing dataset — the fix for the memorization problem.

	The old pipeline froze 20k mixtures to disk (build_dataset.py) and trained on
	that same fixed set every epoch, so the model memorized it (train SI-SDR climbed
	while val stayed pinned at ~+1 dB). This dataset instead calls the synthesis
	engine live: every __getitem__ produces a brand-new random mixture, so the
	model never sees the same example twice and has to learn to separate voices
	rather than memorize files.

	Only the training set should be dynamic. Keep validation on a fixed manifest
	(VantaDataset) so val numbers are comparable across runs.

	Freshness across epochs:
	We reseed the per-process mixer from torch.initial_seed(), which PyTorch
	varies per (epoch, worker). With num_workers=0 the mixer is created once and
	its RNG simply keeps advancing across epochs — either way, no epoch ever
	repeats the previous epoch's mixtures.
	"""

	from __future__ import annotations

	import torch
	from torch.utils.data import Dataset, get_worker_info

	from vanta.data.indexer import NoiseIndex, RirIndex, SpeakerIndex
	from vanta.data.synthesize import MixConfig, Mixer


	class DynamicMixDataset(Dataset):
	def __init__(
	self,
	mix_cfg: MixConfig,
	speakers: SpeakerIndex,
	noise: NoiseIndex \| None,
	rirs: RirIndex \| None,
	epoch_size: int,
	base_seed: int = 0,
	):
	self.mix_cfg = mix_cfg
	self.speakers = speakers
	self.noise = noise
	self.rirs = rirs
	self.epoch_size = int(epoch_size)
	self.base_seed = int(base_seed)
	# Lazily built per process so each DataLoader worker gets its own RNG.
	self._mixer: Mixer \| None = None

	def __len__(self) -> int:
	# Nominal length = how many fresh mixtures make up one "epoch". Purely a
	# cadence knob (when to validate / checkpoint); the data is unbounded.
	return self.epoch_size

	def _mixer_for_process(self) -> Mixer:
	if self._mixer is None:
	info = get_worker_info()
	if info is None:
	# Main process (num_workers=0): seed once; the RNG advances
	# across epochs because this object persists, so data stays fresh.
	seed = self.base_seed
	else:
	# Worker process: torch.initial_seed() is set by PyTorch per
	# (epoch, worker), so reseeding from it gives fresh data every
	# epoch even when workers are torn down and recreated.
	seed = torch.initial_seed() % (2**31)
	self._mixer = Mixer(
	self.mix_cfg,
	speakers=self.speakers,
	noise=self.noise,
	rirs=self.rirs,
	seed=seed,
	)
	return self._mixer

	def __getitem__(self, idx: int) -> dict:
	# idx is ignored — every call is a new random draw.
	r = self._mixer_for_process().mix()
	return {
	"mixture": torch.from_numpy(r.mixture),
	"target": torch.from_numpy(r.target),
	"enrollment": torch.from_numpy(r.enrollment),
	"target_speaker": r.meta["target_speaker"],
	}