[dataset] new io for code reuse for many speech tasks

[Mddct]

#2152

• TODO

  • chain call
  • support raw dataset source
  • support tar shard dataset source
  • prefetch in synchronize way
  • mapper Ignore error
  • sort data pipes
  • check shuffle deterministic
  • batching
    • static batch
    • dynamic batch
  • padding
  • unit test
    • raw dataset
    • tar dataset
    • unit test for all related function in processor.py
  • reproduce exp
    • raw
    • shard

• next pr todo

  • bucket batch [dataset] support bucket by seq length #2333
  • multiprocessing prefetch for CPU bound operation
  • cv support distributed and join
  • 修复断点处的种子

[Mddct]

chain works:

from torch.utils.data import DataLoader
from wenet.dataset.datapipes import (WenetRawDatasetSource,
                                     WenetTarShardDatasetSource)
from wenet.dataset.processor_v2 import (compute_fbank,
                                        compute_log_mel_spectrogram,
                                        decode_wav, parse_json)

if __name__ == '__main__':

    dataset = WenetTarShardDatasetSource(
        'test/resources/dataset/data.shards.list')
    # dataset = dataset.map(decode_wav)
    # dataset = dataset.map(compute_fbank)
    # dataset = dataset.map(compute_log_mel_spectrogram)
    dataloader = DataLoader(dataset,
                            num_workers=1,
                            persistent_workers=True,
                            batch_size=None)
    for d in dataloader:
        print(d['file_name'], d['tar_file_name'], d['txt'], len(d['wav']))
    for d in dataloader:
        print(d['file_name'], d['tar_file_name'], d['txt'], len(d['wav']))
    for d in dataloader:
        print(d['file_name'], d['tar_file_name'], d['txt'], len(d['wav']))

    dataset = dataset.map(decode_wav).map(compute_fbank)
    dataloader = DataLoader(dataset,
                            num_workers=1,
                            persistent_workers=True,
                            batch_size=None)
    for d in dataloader:
        print(d['file_name'], d['tar_file_name'], d['txt'], len(d['wav']),
              d['feat'].size())

    dataset = WenetRawDatasetSource("test/resources/dataset/data.list",
                                    prefetch=10)
    dataset = dataset.map(parse_json)
    dataset = dataset.map(decode_wav)
    dataset = dataset.map(compute_log_mel_spectrogram)
    dataloader = DataLoader(dataset,
                            num_workers=2,
                            persistent_workers=True,
                            batch_size=None)

    for d in dataloader:
        print(d.keys())

    for d in dataloader:
        print(d.keys())

[Mddct]

batch work

    dataset = WenetRawDatasetSource("test/resources/dataset/data.list",
                                    prefetch=10)
    dataset = dataset.map(parse_json)
    dataset = dataset.map(decode_wav)
    dataset = dataset.map(compute_log_mel_spectrogram)

    def fake_labels(sample):
        sample['label'] = [1, 2, 3, 4]
        return sample

    dataset = dataset.map(fake_labels)
    dataset = dataset.batch(batch_size=2, wrapper_class=padding)

    for d in dataset:
        print(d)

[Mddct]

dynamic batch work

[Mddct]

shard 模式下 new io 和old io loss下降曲线几乎一致

[Mddct]

shard 模式结果复现

[xingchensong]

great job！

[robin1001]

Big job!

[kobenaxie]

对比了 #2152 (comment) 中提到的lhotse和new-IO的速度：

• IO	分钟/epoch
  new-IO-dynamic	18
  new-IO-bucket	21
  Lhotse-IO	13

相关配置

• dynamic: max_frames_in_batch = 18000
• bucket_boundaries: [500, 750, 1000]
• bucket_batch_sizes: [36, 24, 18, 9]
• num_workers = 8
• prefectch = 20
• 数据集：好未来中英混数据集
• GPU: 8卡A100
• torch2.2.1-cu118