update

examples/nx_clean_unet/yaml/config.yaml

@@ -16,6 +16,9 @@ tsfm_hidden_size: 512
 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_in_channel: 2

toolbox/torchaudio/models/nx_clean_unet/configuration_nx_clean_unet.py

@@ -24,6 +24,9 @@ class NXCleanUNetConfig(PretrainedConfig):
                  tsfm_attention_heads: int = 4,
                  tsfm_num_blocks: int = 6,
                  tsfm_dropout_rate: float = 0.1,
+                 tsfm_max_length: int = 1024,
+                 tsfm_chunk_size: int = 1,
+                 tsfm_num_left_chunks: int = 128,
 
                  discriminator_dim: int = 32,
                  discriminator_in_channel: int = 2,
@@ -56,6 +59,9 @@ class NXCleanUNetConfig(PretrainedConfig):
         self.tsfm_attention_heads = tsfm_attention_heads
         self.tsfm_num_blocks = tsfm_num_blocks
         self.tsfm_dropout_rate = tsfm_dropout_rate
+        self.tsfm_max_length = tsfm_max_length
+        self.tsfm_chunk_size = tsfm_chunk_size
+        self.tsfm_num_left_chunks = tsfm_num_left_chunks
 
         self.discriminator_dim = discriminator_dim
         self.discriminator_in_channel = discriminator_in_channel

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

@@ -410,7 +410,7 @@ class TransformerEncoder(nn.Module):
                  attention_heads: int = 4,
                  num_blocks: int = 6,
                  dropout_rate: float = 0.1,
-                 max_length: int = 512,
+                 max_length: int = 1024,
                  chunk_size: int = 1,
                  num_left_chunks: int = 128,
                  ):

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

@@ -12,10 +12,13 @@ down_sampling_hidden_channels: 64
 down_sampling_kernel_size: 4
 down_sampling_stride: 2
 
-tsfm_hidden_size: 256
+tsfm_hidden_size: 512
 tsfm_attention_heads: 4
 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_in_channel: 2