update

toolbox/torchaudio/models/dfnet2/modeling_dfnet2.py

@@ -135,7 +135,10 @@ class CausalConv2d(nn.Module):
         return x, new_cache
 
 
-class CausalConvTranspose2d(nn.Module):
+class CausalConvTranspose2dErrorCase(nn.Module):
+    """
+    错误的缓存方法。
+    """
     def __init__(self,
                  in_channels: int,
                  out_channels: int,
@@ -148,7 +151,7 @@ class CausalConvTranspose2d(nn.Module):
                  norm_layer: str = "batch_norm_2d",
                  activation_layer: str = "relu",
                  ):
-        super(CausalConvTranspose2d, self).__init__()
+        super(CausalConvTranspose2dErrorCase, self).__init__()
 
         kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
 
@@ -198,7 +201,7 @@ class CausalConvTranspose2d(nn.Module):
         else:
             self.activation = nn.Identity()
 
-    def forward(self, inputs: torch.Tensor, cache: Tuple[torch.Tensor, torch.Tensor] = None):
+    def forward(self, inputs: torch.Tensor, cache: torch.Tensor = None):
         """
         :param inputs: shape: [b, c, t, f]
         :param cache: shape: [b, c, lookback, f];
@@ -228,6 +231,101 @@ class CausalConvTranspose2d(nn.Module):
         return x, new_cache
 
 
+class CausalConvTranspose2d(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: Union[int, Iterable[int]],
+                 fstride: int = 1,
+                 dilation: int = 1,
+                 pad_f_dim: bool = True,
+                 bias: bool = True,
+                 separable: bool = False,
+                 norm_layer: str = "batch_norm_2d",
+                 activation_layer: str = "relu",
+                 ):
+        super(CausalConvTranspose2d, self).__init__()
+
+        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+
+        if pad_f_dim:
+            fpad = kernel_size[1] // 2
+        else:
+            fpad = 0
+
+        # for last 2 dim, pad (left, right, top, bottom).
+        self.lookback = kernel_size[0] - 1
+        if self.lookback > 0:
+            self.tpad = nn.ConstantPad2d(padding=(0, 0, self.lookback, 0), value=0.0)
+        else:
+            self.tpad = nn.Identity()
+
+        groups = math.gcd(in_channels, out_channels) if separable else 1
+        if groups == 1:
+            separable = False
+
+        self.convt = nn.ConvTranspose2d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            padding=(kernel_size[0] - 1, fpad + dilation - 1),
+            output_padding=(0, fpad),
+            stride=(1, fstride),  # stride over time is always 1
+            dilation=(1, dilation),  # dilation over time is always 1
+            groups=groups,
+            bias=bias,
+        )
+
+        if separable:
+            self.convp = nn.Conv2d(
+                out_channels,
+                out_channels,
+                kernel_size=1,
+                bias=False,
+            )
+        else:
+            self.convp = nn.Identity()
+
+        if norm_layer is not None:
+            norm_layer = norm_layer_dict[norm_layer]
+            self.norm = norm_layer(out_channels)
+        else:
+            self.norm = nn.Identity()
+
+        if activation_layer is not None:
+            activation_layer = activation_layer_dict[activation_layer]
+            self.activation = activation_layer()
+        else:
+            self.activation = nn.Identity()
+
+    def forward(self, inputs: torch.Tensor, cache: torch.Tensor = None):
+        """
+        :param inputs: shape: [b, c, t, f]
+        :param cache: shape: [b, c, lookback, f];
+        :return:
+        """
+        x = inputs
+
+        # x shape: [b, c, t, f]
+        x = self.convt(x)
+        # x shape: [b, c, t+lookback, f]
+
+        if cache is None:
+            x = self.tpad(x)
+        else:
+            x = torch.concat(tensors=[cache, x], dim=2)
+
+        new_cache = None
+        if self.lookback > 0:
+            new_cache = x[:, :, -self.lookback:, :]
+
+        x = self.convp(x)
+        x = self.norm(x)
+        x = self.activation(x)
+
+        return x, new_cache
+
+
 class GroupedLinear(nn.Module):
 
     def __init__(self, input_size: int, hidden_size: int, groups: int = 1):