update

examples/conv_tasnet/step_2_train_model.py

@@ -42,7 +42,7 @@ def get_args():
     parser.add_argument("--max_epochs", default=200, type=int)
 
     parser.add_argument("--batch_size", default=16, type=int)
-    parser.add_argument("--learning_rate", default=1e-4, type=float)
+    parser.add_argument("--learning_rate", default=1e-3, type=float)
     parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
     parser.add_argument("--patience", default=5, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
@@ -243,8 +243,8 @@ def main():
         progress_bar = tqdm(
             desc="Training; epoch-{}".format(idx_epoch),
         )
-        for batch in train_data_loader:
-            clean_audios, noisy_audios = batch
+        for train_batch in train_data_loader:
+            clean_audios, noisy_audios = train_batch
             clean_audios = clean_audios.to(device)
             noisy_audios = noisy_audios.to(device)
 
@@ -298,7 +298,7 @@ def main():
 
             # evaluation
             total_steps += 1
-            if total_steps % args.eval_steps:
+            if total_steps % args.eval_steps == 0:
                 model.eval()
                 torch.cuda.empty_cache()
 
@@ -313,8 +313,8 @@ def main():
                     desc="Evaluation; step-{}".format(total_steps),
                 )
                 with torch.no_grad():
-                    for batch in valid_data_loader:
-                        clean_audios, noisy_audios = batch
+                    for eval_batch in valid_data_loader:
+                        clean_audios, noisy_audios = eval_batch
                         clean_audios = clean_audios.to(device)
                         noisy_audios = noisy_audios.to(device)
 

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

@@ -49,15 +49,16 @@ class DenoiseJsonlDataset(IterableDataset):
                 item = next(iterable_source)
                 # 随机替换缓冲区元素
                 replace_idx = random.randint(0, len(self.buffer_samples) - 1)
-                yield self.buffer_samples[replace_idx]
+                sample = self.buffer_samples[replace_idx]
                 self.buffer_samples[replace_idx] = item
+                yield self.convert_sample(sample)
             except StopIteration:
                 break
 
         # 清空剩余元素
         random.shuffle(self.buffer_samples)
         for sample in self.buffer_samples:
-            yield sample
+            yield self.convert_sample(sample)
 
     def iterable_source(self):
         with open(self.jsonl_file, "r", encoding="utf-8") as f:
@@ -75,7 +76,7 @@ class DenoiseJsonlDataset(IterableDataset):
 
                 snr_db = row["snr_db"]
 
-                row = {
+                sample = {
                     "noise_filename": noise_filename,
                     "noise_raw_duration": noise_raw_duration,
                     "noise_offset": noise_offset,
@@ -88,7 +89,6 @@ class DenoiseJsonlDataset(IterableDataset):
 
                     "snr_db": snr_db,
                 }
-                sample = self.convert_sample(row)
                 yield sample
 
     def convert_sample(self, sample: dict):