optimize

Author: yxlllc

modules/fast_D/discriminator.py

@@ -3,15 +3,6 @@
 import torch.nn.functional as F
 import torch.nn as nn
 
-
-def combine_frames(x, n):
-    B, L, C = x.shape
-    num_groups = L // n
-    if num_groups == 0:
-        return torch.empty(B, 0, n * C, device=x.device, dtype=x.dtype)
-    x = x[:, :num_groups * n, :].reshape(B, num_groups, n * C)
-    return x
-  
 
 class Transpose(nn.Module):
     def __init__(self, dims):
@@ -116,30 +107,26 @@ def __init__(self, period, init_channel=8, strides=[4, 4, 4], kernel_size=11):
     def forward(self, x):
         fmap = []
 
-        # 1d to 2d
         b, _, t = x.shape
-        if t % self.period != 0: # pad first
-            n_pad = self.period - (t % self.period)
-            x = F.pad(x, (0, n_pad), "reflect")
-            t = t + n_pad
-        x = x.view(b, 1, t // self.period, self.period)
-        x = x.permute(0, 3, 2, 1).reshape(b * self.period, t // self.period, 1)
+        n = self.period * np.prod(self.strides)
+        x = x[:, :, : (t // n) * n].view(b, -1, self.period)
+        x = x.transpose(1, 2).reshape(b * self.period, -1, 1)
 
         x = self.pre(x)
         for i, layer in enumerate(self.residual_layers):
             if self.strides[i] > 1:
-                x = combine_frames(x, self.strides[i])
+                x = x.view(b, -1, x.size(2) * self.strides[i])
             x, norm_x = layer(x)
             if i > 0:
-                fmap.append(norm_x.reshape(b, -1))
+                fmap.append(norm_x.view(b, -1))
         x = self.post(F.rms_norm(x, (x.size(-1), )))
-        x = x.reshape(b, -1)
+        x = x.view(b, -1)
 
         return x, fmap
 
 
 class FastMPD(torch.nn.Module):
-    def __init__(self,periods=None, init_channel=8, strides=[1, 2, 4, 4, 2], kernel_size=31):
+    def __init__(self,periods=None, init_channel=8, strides=[4, 4, 4], kernel_size=11):
         super(FastMPD, self).__init__()
         self.periods = periods if periods is not None else [2, 3, 5, 7, 11]
         self.discriminators = nn.ModuleList()