update

examples/conv_tasnet/run.sh

@@ -71,9 +71,8 @@ file_dir="${work_dir}/${file_folder_name}"
 final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
 evaluation_audio_dir="${file_dir}/evaluation_audio"
 
-dataset="${file_dir}/dataset.xlsx"
-train_dataset="${file_dir}/train.xlsx"
-valid_dataset="${file_dir}/valid.xlsx"
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
 
 $verbose && echo "system_version: ${system_version}"
 $verbose && echo "file_folder_name: ${file_folder_name}"

examples/conv_tasnet/step_2_train_model.py

@@ -25,7 +25,7 @@ from torch.nn import functional as F
 from torch.utils.data.dataloader import DataLoader
 from tqdm import tqdm
 
-from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
+from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
 from toolbox.torchaudio.models.conv_tasnet.configuration_conv_tasnet import ConvTasNetConfig
 from toolbox.torchaudio.models.conv_tasnet.modeling_conv_tasnet import ConvTasNet, ConvTasNetPretrainedModel
 from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
@@ -125,37 +125,37 @@ def main():
     logger.info(f"GPU available count: {n_gpu}; device: {device}")
 
     # datasets
-    train_dataset = DenoiseExcelDataset(
-        excel_file=args.train_dataset,
+    train_dataset = DenoiseJsonlDataset(
+        jsonl_file=args.train_dataset,
         expected_sample_rate=8000,
         max_wave_value=32768.0,
     )
-    valid_dataset = DenoiseExcelDataset(
-        excel_file=args.valid_dataset,
+    valid_dataset = DenoiseJsonlDataset(
+        jsonl_file=args.valid_dataset,
         expected_sample_rate=8000,
         max_wave_value=32768.0,
     )
     train_data_loader = DataLoader(
         dataset=train_dataset,
         batch_size=config.batch_size,
-        shuffle=True,
+        # shuffle=True,
         sampler=None,
         # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
         num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
         collate_fn=collate_fn,
         pin_memory=False,
-        prefetch_factor=16,
+        prefetch_factor=2,
     )
     valid_data_loader = DataLoader(
         dataset=valid_dataset,
         batch_size=config.batch_size,
-        shuffle=True,
+        # shuffle=True,
         sampler=None,
         # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
         num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
         collate_fn=collate_fn,
         pin_memory=False,
-        prefetch_factor=16,
+        prefetch_factor=2,
     )
 
     # models

toolbox/torch/utils/data/dataset/denoise_jsonl_dataset.py

@@ -2,6 +2,8 @@
 # -*- coding: utf-8 -*-
 import json
 import os
+import random
+from typing import List
 
 import librosa
 import numpy as np
@@ -9,28 +11,54 @@ import pandas as pd
 from scipy.io import wavfile
 import torch
 import torchaudio
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, IterableDataset
 from tqdm import tqdm
 
 
-class DenoiseJsonlDataset(Dataset):
+class DenoiseJsonlDataset(IterableDataset):
     def __init__(self,
                  jsonl_file: str,
                  expected_sample_rate: int,
                  resample: bool = False,
                  max_wave_value: float = 1.0,
+                 buffer_size: int = 1000,
                  ):
         self.jsonl_file = jsonl_file
         self.expected_sample_rate = expected_sample_rate
         self.resample = resample
         self.max_wave_value = max_wave_value
 
-        self.samples = self.load_samples(jsonl_file)
+        self.buffer_size = buffer_size
+        self.buffer_samples: List[dict] = list()
 
-    @staticmethod
-    def load_samples(filename: str):
-        samples = list()
-        with open(filename, "r", encoding="utf-8") as f:
+    def __iter__(self):
+        iterable_source = self.iterable_source()
+
+        # 初始填充缓冲区
+        try:
+            for _ in range(self.buffer_size):
+                self.buffer_samples.append(next(iterable_source))
+        except StopIteration:
+            pass
+
+        # 动态替换逻辑
+        while True:
+            try:
+                item = next(iterable_source)
+                # 随机替换缓冲区元素
+                replace_idx = random.randint(0, len(self.buffer_samples) - 1)
+                yield self.buffer_samples[replace_idx]
+                self.buffer_samples[replace_idx] = item
+            except StopIteration:
+                break
+
+        # 清空剩余元素
+        random.shuffle(self.buffer_samples)
+        for sample in self.buffer_samples:
+            yield sample
+
+    def iterable_source(self):
+        with open(self.jsonl_file, "r", encoding="utf-8") as f:
             for row in f:
                 row = json.loads(row)
                 noise_filename = row["noise_filename"]
@@ -58,11 +86,10 @@ class DenoiseJsonlDataset(Dataset):
 
                     "snr_db": snr_db,
                 }
-                samples.append(row)
-        return samples
+                sample = self.convert_sample(row)
+                yield sample
 
-    def __getitem__(self, index):
-        sample = self.samples[index]
+    def convert_sample(self, sample: dict):
         noise_filename = sample["noise_filename"]
         noise_offset = sample["noise_offset"]
         noise_duration = sample["noise_duration"]
@@ -92,9 +119,6 @@ class DenoiseJsonlDataset(Dataset):
         }
         return result
 
-    def __len__(self):
-        return len(self.samples)
-
     def filename_to_waveform(self, filename: str, offset: float, duration: float):
         try:
             waveform, sample_rate = librosa.load(
@@ -129,5 +153,5 @@ class DenoiseJsonlDataset(Dataset):
         return noisy_signal, noise_adjusted
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     pass