update

examples/nx_clean_unet/yaml/config.yaml

@@ -16,9 +16,10 @@ tsfm_hidden_size: 256
 tsfm_attention_heads: 8
 tsfm_num_blocks: 6
 tsfm_dropout_rate: 0.1
-tsfm_max_length: 5120
+tsfm_max_length: 512
 tsfm_chunk_size: 4
 tsfm_num_left_chunks: 64
+tsfm_num_right_chunks: 2
 
 discriminator_dim: 32
 discriminator_in_channel: 2

toolbox/torchaudio/models/nx_clean_unet/configuration_nx_clean_unet.py

@@ -25,8 +25,9 @@ class NXCleanUNetConfig(PretrainedConfig):
                  tsfm_num_blocks: int = 6,
                  tsfm_dropout_rate: float = 0.1,
                  tsfm_max_length: int = 1024,
-                 tsfm_chunk_size: int = 1,
+                 tsfm_chunk_size: int = 4,
                  tsfm_num_left_chunks: int = 128,
+                 tsfm_num_right_chunks: int = 2,
 
                  discriminator_dim: int = 16,
                  discriminator_in_channel: int = 2,
@@ -62,6 +63,7 @@ class NXCleanUNetConfig(PretrainedConfig):
         self.tsfm_max_length = tsfm_max_length
         self.tsfm_chunk_size = tsfm_chunk_size
         self.tsfm_num_left_chunks = tsfm_num_left_chunks
+        self.tsfm_num_right_chunks = tsfm_num_right_chunks
 
         self.discriminator_dim = discriminator_dim
         self.discriminator_in_channel = discriminator_in_channel

toolbox/torchaudio/models/nx_clean_unet/modeling_nx_clean_unet.py

@@ -172,6 +172,9 @@ class NXCleanUNet(nn.Module):
             attention_heads=config.tsfm_attention_heads,
             num_blocks=config.tsfm_num_blocks,
             dropout_rate=config.tsfm_dropout_rate,
+            chunk_size=config.chunk_size,
+            num_left_chunks=config.num_left_chunks,
+            num_right_chunks=config.num_right_chunks,
         )
         self.up_sampling = UpSampling(
             num_layers=config.down_sampling_num_layers,

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

@@ -7,7 +7,7 @@ import torch
 import torch.nn as nn
 
 
-class MultiHeadAttention(nn.Module):
+class MultiHeadSelfAttention(nn.Module):
     def __init__(self, n_head: int, n_feat: int, dropout_rate: float):
         """
         :param n_head: int. the number of heads.
@@ -86,14 +86,12 @@ class MultiHeadAttention(nn.Module):
         return self.linear_out(x)  # (batch, time1, n_feat)
 
     def forward(self,
-                query: torch.Tensor,
-                key: torch.Tensor,
-                value: torch.Tensor,
+                x: torch.Tensor,
                 mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
                 cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
                 ) -> Tuple[torch.Tensor, torch.Tensor]:
 
-        q, k, v = self.forward_qkv(query, key, value)
+        q, k, v = self.forward_qkv(x, x, x)
 
         if cache.size(0) > 0:
             key_cache, value_cache = torch.split(
@@ -157,32 +155,40 @@ class RelativeMultiHeadSelfAttention(nn.Module):
     def forward_attention(self,
                           value: torch.Tensor,
                           scores: torch.Tensor,
-                          mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+                          mask: torch.Tensor = None
                           ) -> torch.Tensor:
         """
         compute attention context vector.
-        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, time2, d_k).
-        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, time1, time2).
-        :param mask: torch.Tensor. mask. shape=(batch_size, 1, time2) or
-                (batch_size, time1, time2), (0, 0, 0) means fake mask.
-        :return: torch.Tensor. transformed value. (batch_size, time1, d_model).
-                weighted by the attention score (batch_size, time1, time2).
+        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, key_time_steps, d_k).
+        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, query_time_steps, key_time_steps).
+        :param mask: torch.Tensor. mask. shape=(batch_size, 1, key_time_steps) or (batch_size, query_time_steps, key_time_steps).
+        :return: torch.Tensor. transformed value. (batch_size, query_time_steps, d_model).
+                weighted by the attention score (batch_size, query_time_steps, key_time_steps).
         """
         n_batch = value.size(0)
-        if mask.size(2) > 0:  # time2 > 0
-            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
-            # For last chunk, time2 might be larger than scores.size(-1)
-            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
+        if mask is not None:
+            mask = mask.unsqueeze(1).eq(0)
+            # mask shape: [batch_size, 1, query_time_steps, key_time_steps]
             scores = scores.masked_fill(mask, -float('inf'))
-            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)
         else:
-            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+            attn = torch.softmax(scores, dim=-1)
+        # attn shape: [batch_size, n_head, query_time_steps, key_time_steps]
 
         p_attn = self.dropout(attn)
-        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
-        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
 
-        return self.linear_out(x)  # (batch, time1, n_feat)
+        x = torch.matmul(p_attn, value)
+        # x shape: [batch_size, n_head, query_time_steps, d_k]
+        x = x.transpose(1, 2)
+        # x shape: [batch_size, query_time_steps, n_head, d_k]
+
+        x = x.contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
+        # x shape: [batch_size, query_time_steps, n_head * d_k]
+        # x shape: [batch_size, query_time_steps, n_feat]
+
+        x = self.linear_out(x)
+        # x shape: [batch_size, query_time_steps, n_feat]
+        return x
 
     def relative_position_encoding(self, length: int) -> torch.Tensor:
         """
@@ -197,18 +203,16 @@ class RelativeMultiHeadSelfAttention(nn.Module):
         return final_mat
 
     def forward(self,
-                query: torch.Tensor,
-                key: torch.Tensor,
-                value: torch.Tensor,
-                mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
-                cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+                x: torch.Tensor,
+                mask: torch.Tensor = None,
+                cache: torch.Tensor = None
                 ) -> Tuple[torch.Tensor, torch.Tensor]:
         # attention! self attention.
 
-        q, k, v = self.forward_qkv(query, key, value)
-        # q shape: [batch_size, self.h, time_steps, self.d_k]
+        q, k, v = self.forward_qkv(x, x, x)
+        # q k v shape: [batch_size, self.h, query_time_steps, self.d_k]
 
-        if cache.size(0) > 0:
+        if cache is not None:
             key_cache, value_cache = torch.split(
                 cache, cache.size(-1) // 2, dim=-1)
             k = torch.cat([key_cache, k], dim=2)
@@ -217,11 +221,13 @@ class RelativeMultiHeadSelfAttention(nn.Module):
         # new_cache shape: [batch_size, self.h, time_steps, self.d_k * 2]
         new_cache = torch.cat((k, v), dim=-1)
 
+        # native_scores shape: [batch_size, self.h, q_time_steps, k_time_steps]
+        native_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+
         # Compute relative position encoding
         q_length, k_length = q.size(2), k.size(2)
         relative_position = self.relative_position_encoding(k_length)
 
-        # 流式推理时 q_length 与 k_length 不同。
         relative_position = relative_position[-q_length:]
 
         relative_position_k = self.relative_position_k[relative_position.view(-1)].view(q_length, k_length, -1)
@@ -229,11 +235,10 @@ class RelativeMultiHeadSelfAttention(nn.Module):
         relative_position_k = relative_position_k.unsqueeze(0).unsqueeze(0)  # (1, 1, q_length, k_length, d_k)
         relative_position_k = relative_position_k.expand(q.size(0), q.size(1), -1, -1, -1)  # (batch, head, q_length, k_length, d_k)
 
-        native_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
-        # native_scores shape: [batch_size, self.h, q_time_steps, k_time_steps]
-
         relative_position_scores = torch.matmul(q.unsqueeze(3), relative_position_k.transpose(-2, -1)).squeeze(3) / math.sqrt(self.d_k)
         # relative_position_scores shape: [batch_size, self.h, q_time_steps, k_time_steps]
+
+        # score
         scores = native_scores + relative_position_scores
 
         return self.forward_attention(v, scores, mask), new_cache

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

@@ -25,6 +25,7 @@ def subsequent_chunk_mask(
         size: int,
         chunk_size: int,
         num_left_chunks: int = -1,
+        num_right_chunks: int = 0,
         device: torch.device = torch.device("cpu"),
 ) -> torch.Tensor:
     """
@@ -41,6 +42,7 @@ def subsequent_chunk_mask(
     :param size: int. size of mask.
     :param chunk_size: int. size of chunk.
     :param num_left_chunks: int. number of left chunks. <0: use full chunk. >=0 use num_left_chunks.
+    :param num_right_chunks: int. number of right chunks.
     :param device: torch.device. "cpu" or "cuda" or torch.Tensor.device.
     :return: torch.Tensor. mask
     """
@@ -51,7 +53,7 @@ def subsequent_chunk_mask(
             start = 0
         else:
             start = max((i // chunk_size - num_left_chunks) * chunk_size, 0)
-        ending = min((i // chunk_size + 1) * chunk_size, size)
+        ending = min((i // chunk_size + 1 + num_right_chunks) * chunk_size, size)
         ret[i, start:ending] = True
     return ret
 
@@ -59,6 +61,12 @@ def subsequent_chunk_mask(
 def main():
     chunk_mask = subsequent_chunk_mask(size=8, chunk_size=2, num_left_chunks=2)
     print(chunk_mask)
+
+    chunk_mask = subsequent_chunk_mask(size=8, chunk_size=2, num_left_chunks=2, num_right_chunks=1)
+    print(chunk_mask)
+
+    chunk_mask = subsequent_chunk_mask(size=9, chunk_size=2, num_left_chunks=2, num_right_chunks=1)
+    print(chunk_mask)
     return
 
 

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

@@ -7,7 +7,7 @@ import torch
 import torch.nn as nn
 
 from toolbox.torchaudio.models.nx_clean_unet.transformer.mask import subsequent_chunk_mask
-from toolbox.torchaudio.models.nx_clean_unet.transformer.attention import MultiHeadAttention, RelativeMultiHeadSelfAttention
+from toolbox.torchaudio.models.nx_clean_unet.transformer.attention import MultiHeadSelfAttention, RelativeMultiHeadSelfAttention
 
 
 class PositionwiseFeedForward(nn.Module):
@@ -87,7 +87,7 @@ class TransformerEncoderLayer(nn.Module):
         xt = self.norm1(x)
 
         x_att, new_att_cache = self.attention.forward(
-            xt, xt, xt, mask=mask, cache=attention_cache
+            xt, mask=mask, cache=attention_cache
         )
         x = x + self.dropout1(xt)
         xt = self.norm2(x)
@@ -112,6 +112,7 @@ class TransformerEncoder(nn.Module):
                  max_relative_position: int = 1024,
                  chunk_size: int = 1,
                  num_left_chunks: int = 128,
+                 num_right_chunks: int = 2,
                  ):
         super().__init__()
         self.input_size = input_size
@@ -120,6 +121,7 @@ class TransformerEncoder(nn.Module):
         self.max_relative_position = max_relative_position
         self.chunk_size = chunk_size
         self.num_left_chunks = num_left_chunks
+        self.num_right_chunks = num_right_chunks
 
         self.input_linear = nn.Linear(
             in_features=self.input_size,
@@ -155,7 +157,8 @@ class TransformerEncoder(nn.Module):
         chunk_masks = subsequent_chunk_mask(
             size=time_steps,
             chunk_size=self.chunk_size,
-            num_left_chunks=self.num_left_chunks
+            num_left_chunks=self.num_left_chunks,
+            num_right_chunks=self.num_right_chunks,
         )
         chunk_masks = chunk_masks.to(xs.device)
         # chunk_masks shape: [1, time_steps, time_steps]

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

@@ -10,23 +10,24 @@ hop_size: 80
 # 例如 2**5=32 就意味着 32个值在降采样之后是一个时间步，
 # 则一步是 32/sample_rate = 0.004秒。
 # 那么 tsfm_chunk_size=4 则为16ms，tsfm_chunk_size=8 则为32ms
-# 假设每次向左看1秒，则：
-# tsfm_chunk_size=1，tsfm_num_left_chunks: 256
-# tsfm_chunk_size=4，tsfm_num_left_chunks: 64
-# tsfm_chunk_size=8，tsfm_num_left_chunks: 32
+# 假设每次向左看1秒，向右看30ms，则：
+# tsfm_chunk_size=1，tsfm_num_left_chunks=256，tsfm_num_right_chunks=8
+# tsfm_chunk_size=4，tsfm_num_left_chunks=64，tsfm_num_right_chunks=2
+# tsfm_chunk_size=8，tsfm_num_left_chunks=32，tsfm_num_right_chunks=1
 down_sampling_num_layers: 5
 down_sampling_in_channels: 1
 down_sampling_hidden_channels: 64
 down_sampling_kernel_size: 4
 down_sampling_stride: 2
 
-tsfm_hidden_size: 64
-tsfm_attention_heads: 4
+tsfm_hidden_size: 256
+tsfm_attention_heads: 8
 tsfm_num_blocks: 6
 tsfm_dropout_rate: 0.1
-tsfm_max_length: 5120
+tsfm_max_length: 512
 tsfm_chunk_size: 4
 tsfm_num_left_chunks: 64
+tsfm_num_right_chunks: 2
 
 discriminator_dim: 32
 discriminator_in_channel: 2