update

examples/nx_clean_unet/step_2_train_model.py

@@ -143,7 +143,7 @@ def main():
         num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
         collate_fn=collate_fn,
         pin_memory=False,
-        # prefetch_factor=64,
+        prefetch_factor=16,
     )
     valid_data_loader = DataLoader(
         dataset=valid_dataset,
@@ -154,7 +154,7 @@ def main():
         num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
         collate_fn=collate_fn,
         pin_memory=False,
-        # prefetch_factor=64,
+        prefetch_factor=16,
     )
 
     # models

examples/nx_clean_unet/yaml/config.yaml

@@ -20,7 +20,7 @@ tsfm_max_length: 1024
 tsfm_chunk_size: 1
 tsfm_num_left_chunks: 128
 
-discriminator_dim: 32
+discriminator_dim: 16
 discriminator_in_channel: 2
 
 compress_factor: 0.3

toolbox/torchaudio/models/nx_clean_unet/configuration_nx_clean_unet.py

@@ -28,7 +28,7 @@ class NXCleanUNetConfig(PretrainedConfig):
                  tsfm_chunk_size: int = 1,
                  tsfm_num_left_chunks: int = 128,
 
-                 discriminator_dim: int = 32,
+                 discriminator_dim: int = 16,
                  discriminator_in_channel: int = 2,
 
                  compress_factor: float = 0.3,

toolbox/torchaudio/models/nx_clean_unet/transformer/transformer.py

@@ -562,6 +562,7 @@ def main():
         num_blocks=6,
         dropout_rate=0.1,
     )
+    print(encoder)
 
     x = torch.ones([4, 200, 64])
 

toolbox/torchaudio/models/nx_clean_unet/yaml/config.yaml

@@ -12,15 +12,15 @@ down_sampling_hidden_channels: 64
 down_sampling_kernel_size: 4
 down_sampling_stride: 2
 
-tsfm_hidden_size: 512
-tsfm_attention_heads: 4
+tsfm_hidden_size: 1024
+tsfm_attention_heads: 8
 tsfm_num_blocks: 6
 tsfm_dropout_rate: 0.1
 tsfm_max_length: 1024
 tsfm_chunk_size: 1
 tsfm_num_left_chunks: 128
 
-discriminator_dim: 32
+discriminator_dim: 16
 discriminator_in_channel: 2
 
 compress_factor: 0.3