reverted in-place optimizations for generators and discriminators, simplified checkpointing code, refinegan reworked

Author: AznamirWoW

rvc/lib/algorithm/discriminators.py

@@ -21,36 +21,20 @@ class MultiPeriodDiscriminator(torch.nn.Module):
     """
 
     def __init__(self, use_spectral_norm: bool = False, checkpointing: bool = False):
-        super(MultiPeriodDiscriminator, self).__init__()
+        super().__init__()
         periods = [2, 3, 5, 7, 11, 17, 23, 37]
         self.checkpointing = checkpointing
         self.discriminators = torch.nn.ModuleList(
-            [
-                DiscriminatorS(
-                    use_spectral_norm=use_spectral_norm, checkpointing=checkpointing
-                )
-            ]
-            + [
-                DiscriminatorP(
-                    p, use_spectral_norm=use_spectral_norm, checkpointing=checkpointing
-                )
-                for p in periods
-            ]
+            [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+            + [DiscriminatorP(p, use_spectral_norm=use_spectral_norm) for p in periods]
         )
 
     def forward(self, y, y_hat):
         y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
         for d in self.discriminators:
             if self.training and self.checkpointing:
-
-                def forward_discriminator(d, y, y_hat):
-                    y_d_r, fmap_r = d(y)
-                    y_d_g, fmap_g = d(y_hat)
-                    return y_d_r, fmap_r, y_d_g, fmap_g
-
-                y_d_r, fmap_r, y_d_g, fmap_g = checkpoint(
-                    forward_discriminator, d, y, y_hat, use_reentrant=False
-                )
+                y_d_r, fmap_r = checkpoint(d, y, use_reentrant=False)
+                y_d_g, fmap_g = checkpoint(d, y_hat, use_reentrant=False)
             else:
                 y_d_r, fmap_r = d(y)
                 y_d_g, fmap_g = d(y_hat)
@@ -71,9 +55,9 @@ class DiscriminatorS(torch.nn.Module):
     convolutional layers that are applied to the input signal.
     """
 
-    def __init__(self, use_spectral_norm: bool = False, checkpointing: bool = False):
-        super(DiscriminatorS, self).__init__()
-        self.checkpointing = checkpointing
+    def __init__(self, use_spectral_norm: bool = False):
+        super().__init__()
+
         norm_f = spectral_norm if use_spectral_norm else weight_norm
         self.convs = torch.nn.ModuleList(
             [
@@ -86,16 +70,12 @@ def __init__(self, use_spectral_norm: bool = False, checkpointing: bool = False)
             ]
         )
         self.conv_post = norm_f(torch.nn.Conv1d(1024, 1, 3, 1, padding=1))
-        self.lrelu = torch.nn.LeakyReLU(LRELU_SLOPE, inplace=True)
+        self.lrelu = torch.nn.LeakyReLU(LRELU_SLOPE)
 
     def forward(self, x):
         fmap = []
         for conv in self.convs:
-            if self.training and self.checkpointing:
-                x = checkpoint(conv, x, use_reentrant=False)
-                x = checkpoint(self.lrelu, x, use_reentrant=False)
-            else:
-                x = self.lrelu(conv(x))
+            x = self.lrelu(conv(x))
             fmap.append(x)
         x = self.conv_post(x)
         fmap.append(x)
@@ -125,10 +105,8 @@ def __init__(
         kernel_size: int = 5,
         stride: int = 3,
         use_spectral_norm: bool = False,
-        checkpointing: bool = False,
     ):
-        super(DiscriminatorP, self).__init__()
-        self.checkpointing = checkpointing
+        super().__init__()
         self.period = period
         norm_f = spectral_norm if use_spectral_norm else weight_norm
 
@@ -151,7 +129,7 @@ def __init__(
         )
 
         self.conv_post = norm_f(torch.nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
-        self.lrelu = torch.nn.LeakyReLU(LRELU_SLOPE, inplace=True)
+        self.lrelu = torch.nn.LeakyReLU(LRELU_SLOPE)
 
     def forward(self, x):
         fmap = []
@@ -162,13 +140,8 @@ def forward(self, x):
         x = x.view(b, c, -1, self.period)
 
         for conv in self.convs:
-            if self.training and self.checkpointing:
-                x = checkpoint(conv, x, use_reentrant=False)
-                x = checkpoint(self.lrelu, x, use_reentrant=False)
-            else:
-                x = self.lrelu(conv(x))
+            x = self.lrelu(conv(x))
             fmap.append(x)
-
         x = self.conv_post(x)
         fmap.append(x)
         x = torch.flatten(x, 1, -1)

rvc/lib/algorithm/generators/hifigan.py

@@ -75,12 +75,10 @@ def forward(self, x: torch.Tensor, g: Optional[torch.Tensor] = None):
         x = self.conv_pre(x)
 
         if g is not None:
-            # in-place call
-            x += self.cond(g)
+            x = x + self.cond(g)
 
         for i in range(self.num_upsamples):
-            # in-place call
-            x = torch.nn.functional.leaky_relu_(x, LRELU_SLOPE)
+            x = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
             x = self.ups[i](x)
             xs = None
             for j in range(self.num_kernels):
@@ -90,10 +88,10 @@ def forward(self, x: torch.Tensor, g: Optional[torch.Tensor] = None):
                     xs += self.resblocks[i * self.num_kernels + j](x)
             x = xs / self.num_kernels
         # in-place call
-        x = torch.nn.functional.leaky_relu_(x)
+        x = torch.nn.functional.leaky_relu(x)
         x = self.conv_post(x)
         # in-place call
-        x = torch.tanh_(x)
+        x = torch.tanh(x)
 
         return x
 

rvc/lib/algorithm/generators/hifigan_mrf.py

@@ -43,11 +43,9 @@ def __init__(self, channels, kernel_size, dilation):
         )
 
     def forward(self, x: torch.Tensor):
-        # new tensor
         y = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
         y = self.conv1(y)
-        # in-place call
-        y = torch.nn.functional.leaky_relu_(y, LRELU_SLOPE)
+        y = torch.nn.functional.leaky_relu(y, LRELU_SLOPE)
         y = self.conv2(y)
         return x + y
 
@@ -344,36 +342,31 @@ def forward(
         f0 = self.f0_upsample(f0[:, None, :]).transpose(-1, -2)
         har_source, _, _ = self.m_source(f0)
         har_source = har_source.transpose(-1, -2)
-        # new tensor
         x = self.conv_pre(x)
 
         if g is not None:
-            # in-place call
-            x += self.cond(g)
+            x = x + self.cond(g)
 
         for ups, mrf, noise_conv in zip(self.upsamples, self.mrfs, self.noise_convs):
-            # in-place call
-            x = torch.nn.functional.leaky_relu_(x, LRELU_SLOPE)
+            x = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
 
             if self.training and self.checkpointing:
                 x = checkpoint(ups, x, use_reentrant=False)
+                x = x + noise_conv(har_source)
+                xs = sum([
+                    checkpoint(layer, x, use_reentrant=False)
+                    for layer in mrf])
             else:
                 x = ups(x)
+                x = x + noise_conv(har_source)
+                xs = sum([
+                    layer(x)
+                    for layer in mrf])
+            x = xs / self.num_kernels
 
-            x += noise_conv(har_source)
+        x = torch.nn.functional.leaky_relu(x)
+        x = torch.tanh(self.conv_post(x))
 
-            def mrf_sum(x, layers):
-                return sum(layer(x) for layer in layers) / self.num_kernels
-
-            if self.training and self.checkpointing:
-                x = checkpoint(mrf_sum, x, mrf, use_reentrant=False)
-            else:
-                x = mrf_sum(x, mrf)
-        # in-place call
-        x = torch.nn.functional.leaky_relu_(x)
-        x = self.conv_post(x)
-        # in-place call
-        x = torch.tanh_(x)
         return x
 
     def remove_weight_norm(self):

rvc/lib/algorithm/generators/hifigan_nsf.py

@@ -179,37 +179,29 @@ def forward(
         x = self.conv_pre(x)
 
         if g is not None:
-            # in-place call
-            x += self.cond(g)
+            x = x + self.cond(g)
 
         for i, (ups, noise_convs) in enumerate(zip(self.ups, self.noise_convs)):
-            # in-place call
-            x = torch.nn.functional.leaky_relu_(x, self.lrelu_slope)
-
+            x = torch.nn.functional.leaky_relu(x, self.lrelu_slope)
             # Apply upsampling layer
             if self.training and self.checkpointing:
                 x = checkpoint(ups, x, use_reentrant=False)
+                x = x + noise_convs(har_source)
+                xs = sum([
+                    checkpoint(resblock, x, use_reentrant=False)
+                    for j, resblock in enumerate(self.resblocks)
+                    if j in range(i * self.num_kernels, (i + 1) * self.num_kernels)])
             else:
                 x = ups(x)
-
-            # Add noise excitation
-            x += noise_convs(har_source)
-
-            # Apply residual blocks
-            def resblock_forward(x, blocks):
-                return sum(block(x) for block in blocks) / len(blocks)
-
-            blocks = self.resblocks[i * self.num_kernels : (i + 1) * self.num_kernels]
-
-            # Checkpoint or regular computation for ResBlocks
-            if self.training and self.checkpointing:
-                x = checkpoint(resblock_forward, x, blocks, use_reentrant=False)
-            else:
-                x = resblock_forward(x, blocks)
-        # in-place call
-        x = torch.nn.functional.leaky_relu_(x)
-        # in-place call
-        x = torch.tanh_(self.conv_post(x))
+                x = x + noise_convs(har_source)
+                xs = sum([
+                    resblock(x)
+                    for j, resblock in enumerate(self.resblocks)
+                    if j in range(i * self.num_kernels, (i + 1) * self.num_kernels)])
+            x = xs / self.num_kernels
+
+        x = torch.nn.functional.leaky_relu(x)
+        x = torch.tanh(self.conv_post(x))
 
         return x