optimize

Author: yxlllc

modules/fast_D/discriminator.py

@@ -34,23 +34,22 @@ def forward(ctx, x, min_val, max_val, leak_slope):
         ctx.max_val = max_val
         ctx.leak_slope = leak_slope
         below_mask = x < min_val
-        any_below = torch.any(below_mask)
-        if any_below:
-            x[below_mask] = leak_slope * x[below_mask] + (1 - leak_slope) * min_val
         above_mask = x > max_val
-        any_above = torch.any(above_mask)
-        if any_above:
-            x[above_mask] = leak_slope * x[above_mask] + (1 - leak_slope) * max_val
-        if any_below or any_above:
-            ctx.save_for_backward(below_mask | above_mask)
+        mask = below_mask | above_mask
+        if mask.any().item():
+            ctx.save_for_backward(mask)
+            below_val = leak_slope * x + (1.0 - leak_slope) * min_val
+            above_val = leak_slope * x + (1.0 - leak_slope) * max_val
+            x = torch.where(below_mask, below_val, x)
+            x = torch.where(above_mask, above_val, x)
         return x
 
     @staticmethod
-    def backward(ctx, grad_output):
+    def backward(ctx, grad):
         if len(ctx.saved_tensors) > 0:
             mask, = ctx.saved_tensors
-            grad_output[mask] *= ctx.leak_slope
-        return grad_output, None, None, None
+            grad = torch.where(mask, grad * ctx.leak_slope, grad)
+        return grad, None, None, None
 
 
 class ATanGLU(nn.Module):
@@ -82,7 +81,7 @@ def __init__(self, dim, expansion_factor, kernel_size=31, dropout=0.):
             nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim),
             Transpose((1, 2)),
             nn.Linear(dim, inner_dim * 2),
-            ATanGLU(),
+            ATanGLU(hard_limit=True),
             nn.Linear(inner_dim, inner_dim * 2),
             ATanGLU(hard_limit=True),
             nn.Linear(inner_dim, dim),