update

examples/clean_unet_aishell/run.sh

@@ -14,7 +14,8 @@ sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name fi
 
 sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
+--speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train" \
+--max_count 10000
 
 
 END
@@ -35,6 +36,8 @@ limit=10
 noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
 speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
 
+max_count=10000000
+
 nohup_name=nohup.out
 
 # model params
@@ -101,6 +104,7 @@ if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
   --speech_dir "${speech_dir}" \
   --train_dataset "${train_dataset}" \
   --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
 
 fi
 

examples/clean_unet_aishell/step_1_prepare_data.py

@@ -42,7 +42,7 @@ def get_args():
 
     parser.add_argument("--target_sample_rate", default=8000, type=int)
 
-    parser.add_argument("--scale", default=1, type=float)
+    parser.add_argument("--max_count", default=10000, type=int)
 
     args = parser.parse_args()
     return args
@@ -101,9 +101,8 @@ def get_dataset(args):
     count = 0
     process_bar = tqdm(desc="build dataset excel")
     for noise, speech in zip(noise_generator, speech_generator):
-        flag = random.random()
-        if flag > args.scale:
-            continue
+        if count > args.max_count:
+            break
 
         noise_filename = noise["filename"]
         noise_raw_duration = noise["raw_duration"]

examples/clean_unet_aishell/step_2_train_model.py

@@ -44,10 +44,10 @@ def get_args():
 
     parser.add_argument("--batch_size", default=64, type=int)
     parser.add_argument("--learning_rate", default=2e-4, 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)
+    parser.add_argument("--seed", default=0, type=int)
 
     parser.add_argument("--config_file", default="config.yaml", type=str)
 
@@ -119,10 +119,10 @@ def main():
 
     logger = logging_config(serialization_dir)
 
-    random.seed(config.seed)
-    np.random.seed(config.seed)
-    torch.manual_seed(config.seed)
-    logger.info(f"set seed: {config.seed}")
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    logger.info(f"set seed: {args.seed}")
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     n_gpu = torch.cuda.device_count()
@@ -141,7 +141,7 @@ def main():
     )
     train_data_loader = DataLoader(
         dataset=train_dataset,
-        batch_size=config.batch_size,
+        batch_size=args.batch_size,
         shuffle=True,
         sampler=None,
         # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
@@ -152,7 +152,7 @@ def main():
     )
     valid_data_loader = DataLoader(
         dataset=valid_dataset,
-        batch_size=config.batch_size,
+        batch_size=args.batch_size,
         shuffle=True,
         sampler=None,
         # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
@@ -168,7 +168,7 @@ def main():
 
     # optimizer
     logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
-    optimizer = torch.optim.AdamW(model.parameters(), config.learning_rate)
+    optimizer = torch.optim.AdamW(model.parameters(), args.learning_rate)
     lr_scheduler = LinearWarmupCosineDecay(
         optimizer,
         lr_max=args.learning_rate,

examples/clean_unet_aishell/yaml/config.yaml

@@ -11,3 +11,4 @@ tsfm_n_layers: 5
 tsfm_n_head: 8
 tsfm_d_model: 512
 tsfm_d_inner: 2048
+