PupuM2D/dataloader_webdataset.py at main · sizigi/PupuM2D · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import random
import torch
import torchaudio
import webdataset as wds
from glob import glob
from torch.utils.data import get_worker_info


class StreamingAudioWebDataset:
    def __init__(self, root_dir, cfg):
        self.root_dir = root_dir
        self.cfg = cfg

        self.target_sr = self.cfg.preprocess.sample_rate
        self.frame_len = (
            self.cfg.preprocess.cut_mel_frame * self.cfg.preprocess.hop_size
        )

        self.file_list = glob(f"{self.root_dir}/*/*.flac")
        random.shuffle(self.file_list)

        self.pipeline = wds.DataPipeline(
            self._get_file_list,
            wds.shuffle(self.cfg.train.dataloader.shuffle_small),
            self._load_and_chunk_audio,
            wds.shuffle(self.cfg.train.dataloader.shuffle_large),
            wds.batched(
                self.cfg.train.batch_size,
                collation_fn=self._collate_fn,
                partial=not self.cfg.train.dataloader.drop_last,
            ),
        )

    def _get_file_list(self):
        worker_info = get_worker_info()
        if worker_info is not None:
            worker_files = self.file_list[worker_info.id :: worker_info.num_workers]
        else:
            worker_files = self.file_list

        random.shuffle(worker_files)
        for path in worker_files:
            yield {"path": path}

    def _load_and_chunk_audio(self, sample_iter):
        target_sr = self.target_sr
        frame_len = self.frame_len
        resamplers = {}

        for sample in sample_iter:
            path = sample["path"]
            try:
                audio, sr = torchaudio.load(path, normalize=True)

                if audio.shape[0] > 1:
                    audio = audio.mean(dim=0, keepdim=True)

                if sr != target_sr:
                    if sr not in resamplers:
                        resamplers[sr] = torchaudio.transforms.Resample(sr, target_sr)
                    audio = resamplers[sr](audio)

                audio = audio.squeeze(0)
                audio_len = audio.shape[-1]

                if audio_len < frame_len:
                    pad_amt = frame_len - audio_len
                    chunk = torch.nn.functional.pad(audio, (0, pad_amt))
                    yield {"audio": chunk}
                else:
                    chunks = audio.unfold(0, frame_len, frame_len)
                    remainder_start = chunks.shape[0] * frame_len
                    if remainder_start < audio_len:
                        remainder = audio[remainder_start:]
                        pad_amt = frame_len - remainder.shape[0]
                        last_chunk = torch.nn.functional.pad(remainder, (0, pad_amt))
                        for i in range(chunks.shape[0]):
                            yield {"audio": chunks[i].clone()}
                        yield {"audio": last_chunk}
                    else:
                        for i in range(chunks.shape[0]):
                            yield {"audio": chunks[i].clone()}
            except Exception as e:
                print(f"[WARN] Failed to load {path}: {e}")
                continue

    def _collate_fn(self, batch):
        audios = [b["audio"] for b in batch]
        audios = torch.stack(audios)
        return audios

    def get_dataloader(self):
        loader = wds.WebLoader(
            self.pipeline,
            num_workers=self.cfg.train.dataloader.num_workers,
            pin_memory=self.cfg.train.dataloader.pin_memory,
            prefetch_factor=self.cfg.train.dataloader.prefetch_factor,
            persistent_workers=True,
        )
        return loader