use a customized gru.

app.py

@@ -1,6 +1,5 @@
-import os
-
 ## build wavegru-cpp
+# import os
 # os.system("./bazelisk-linux-amd64 clean --expunge")
 # os.system("./bazelisk-linux-amd64 build wavegru_mod -c opt --copt=-march=native")
 
@@ -18,14 +17,12 @@ alphabet, tacotron_net, tacotron_config = load_tacotron_model(
 wavegru_config, wavegru_net = load_wavegru_net("./wavegru.yaml", "./wavegru.ckpt")
 
 wave_cpp_weight_mask = extract_weight_mask(wavegru_net)
-wavecpp = load_wavegru_cpp(wave_cpp_weight_mask)
+wavecpp = load_wavegru_cpp(wave_cpp_weight_mask, wavegru_config["upsample_factors"][-1])
 
 
 def speak(text):
     mel = text_to_mel(tacotron_net, text, alphabet, tacotron_config)
-    print(mel.shape)
     y = mel_to_wav(wavegru_net, wavecpp, mel, wavegru_config)
-    print(y.shape)
     return 24_000, y
 
 
@@ -38,7 +35,7 @@ gr.Interface(
     examples=[
         "this is a test!",
         "October arrived, spreading a damp chill over the grounds and into the castle. Madam Pomfrey, the nurse, was kept busy by a sudden spate of colds among the staff and students.",
-        "Artificial intelligence is intelligence demonstrated by machines, as opposed to natural intelligence displayed by animals including humans",
+        "Artificial intelligence is intelligence demonstrated by machines, as opposed to natural intelligence displayed by animals including humans.",
     ],
     outputs="audio",
     title=title,

extract_model.py

@@ -0,0 +1,5 @@
+import pickle
+
+dic = pickle.load(open("./wavegru.ckpt", "rb"))
+del dic["optim_state_dict"]
+pickle.dump(dic, open("./wavegru.ckpt", "wb"))

inference.py

@@ -49,7 +49,6 @@ def load_wavegru_net(config_file, model_file):
     config = load_wavegru_config(config_file)
     net = WaveGRU(
         mel_dim=config["mel_dim"],
-        embed_dim=config["embed_dim"],
         rnn_dim=config["rnn_dim"],
         upsample_factors=config["upsample_factors"],
     )
@@ -74,7 +73,7 @@ def mel_to_wav(net, netcpp, mel, config):
     )
     ft = wavegru_inference(net, mel)
     ft = jax.device_get(ft[0])
-    wav = netcpp.inference(ft, 1.0)
+    wav = netcpp.inference(ft, 0.9)
     wav = np.array(wav)
     wav = librosa.mu_expand(wav - 127, mu=255)
     wav = librosa.effects.deemphasis(wav, coef=0.86)

wavegru.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1439b9d8bfd62848cd53d1da08bd8f311004eab0f6bcc682144d4dc2b3c1e6fd
-size 56479601
+oid sha256:c06310d989fd524359d5f3ecf8ea1dc146980bb594b7b90553d0d42a64c512d8
+size 58039876

wavegru.py

@@ -2,9 +2,13 @@
 WaveGRU model: melspectrogram => mu-law encoded waveform
 """
 
+from typing import Tuple
+
 import jax
 import jax.numpy as jnp
 import pax
+from pax import GRUState
+from tqdm.cli import tqdm
 
 
 class ReLU(pax.Module):
@@ -37,16 +41,20 @@ def tile_1d(x, factor):
     return x
 
 
-def up_block(dim, factor):
+def up_block(in_dim, out_dim, factor, relu=True):
     """
     Tile >> Conv >> BatchNorm >> ReLU
     """
-    return pax.Sequential(
+    f = pax.Sequential(
         lambda x: tile_1d(x, factor),
-        pax.Conv1D(dim, dim, 2 * factor, stride=1, padding="VALID", with_bias=False),
-        pax.LayerNorm(dim, -1, True, True),
-        ReLU(),
+        pax.Conv1D(
+            in_dim, out_dim, 2 * factor, stride=1, padding="VALID", with_bias=False
+        ),
+        pax.LayerNorm(out_dim, -1, True, True),
     )
+    if relu:
+        f >>= ReLU()
+    return f
 
 
 class Upsample(pax.Module):
@@ -54,21 +62,24 @@ class Upsample(pax.Module):
     Upsample melspectrogram to match raw audio sample rate.
     """
 
-    def __init__(self, input_dim, upsample_factors):
+    def __init__(self, input_dim, rnn_dim, upsample_factors):
         super().__init__()
         self.input_conv = pax.Sequential(
-            pax.Conv1D(input_dim, 512, 1, with_bias=False),
-            pax.LayerNorm(512, -1, True, True),
+            pax.Conv1D(input_dim, rnn_dim, 1, with_bias=False),
+            pax.LayerNorm(rnn_dim, -1, True, True),
         )
         self.upsample_factors = upsample_factors
         self.dilated_convs = [
-            dilated_residual_conv_block(512, 3, 1, 2**i) for i in range(5)
+            dilated_residual_conv_block(rnn_dim, 3, 1, 2**i) for i in range(5)
         ]
         self.up_factors = upsample_factors[:-1]
-        self.up_blocks = [up_block(512, x) for x in self.up_factors]
+        self.up_blocks = [up_block(rnn_dim, rnn_dim, x) for x in self.up_factors[:-1]]
+        self.up_blocks.append(
+            up_block(rnn_dim, 3 * rnn_dim, self.up_factors[-1], relu=False)
+        )
         self.final_tile = upsample_factors[-1]
 
-    def __call__(self, x):
+    def __call__(self, x, no_repeat=False):
         x = self.input_conv(x)
         for residual in self.dilated_convs:
             y = residual(x)
@@ -78,28 +89,55 @@ class Upsample(pax.Module):
         for f in self.up_blocks:
             x = f(x)
 
+        if no_repeat:
+            return x
         x = tile_1d(x, self.final_tile)
         return x
 
 
-class Pruner(pax.Module):
+class GRU(pax.Module):
     """
-    Base class for pruners
+    A customized GRU module.
     """
 
-    def __init__(self, update_freq=500):
+    input_dim: int
+    hidden_dim: int
+
+    def __init__(self, hidden_dim: int):
         super().__init__()
-        self.update_freq = update_freq
+        self.hidden_dim = hidden_dim
+        self.h_zrh_fc = pax.Linear(hidden_dim, hidden_dim * 3)
+
+    def initial_state(self, batch_size: int) -> GRUState:
+        """Create an all zeros initial state."""
+        return GRUState(jnp.zeros((batch_size, self.hidden_dim), dtype=jnp.float32))
+
+    def __call__(self, state: GRUState, x) -> Tuple[GRUState, jnp.ndarray]:
+        hidden = state.hidden
+        x_zrh = x
+        h_zrh = self.h_zrh_fc(hidden)
+        x_zr, x_h = jnp.split(x_zrh, [2 * self.hidden_dim], axis=-1)
+        h_zr, h_h = jnp.split(h_zrh, [2 * self.hidden_dim], axis=-1)
+
+        zr = x_zr + h_zr
+        zr = jax.nn.sigmoid(zr)
+        z, r = jnp.split(zr, 2, axis=-1)
+
+        h_hat = x_h + r * h_h
+        h_hat = jnp.tanh(h_hat)
+
+        h = (1 - z) * hidden + z * h_hat
+        return GRUState(h), h
+
+
+class Pruner(pax.Module):
+    """
+    Base class for pruners
+    """
 
     def compute_sparsity(self, step):
-        """
-        Two-stages pruning
-        """
-        t = jnp.power(1 - (step * 1.0 - 1_000) / 300_000, 3)
-        z = 0.5 * jnp.clip(1.0 - t, a_min=0, a_max=1)
-        for i in range(4):
-            t = jnp.power(1 - (step * 1.0 - 1_000 - 400_000 - i * 200_000) / 100_000, 3)
-            z = z + 0.1 * jnp.clip(1 - t, a_min=0, a_max=1)
+        t = jnp.power(1 - (step * 1.0 - 1_000) / 200_000, 3)
+        z = 0.9 * jnp.clip(1.0 - t, a_min=0, a_max=1)
         return z
 
     def prune(self, step, weights):
@@ -120,35 +158,32 @@ class Pruner(pax.Module):
 
 
 class GRUPruner(Pruner):
-    def __init__(self, gru, update_freq=500):
-        super().__init__(update_freq=update_freq)
-        self.xh_zr_fc_mask = jnp.ones_like(gru.xh_zr_fc.weight) == 1
-        self.xh_h_fc_mask = jnp.ones_like(gru.xh_h_fc.weight) == 1
+    def __init__(self, gru):
+        super().__init__()
+        self.h_zrh_fc_mask = jnp.ones_like(gru.h_zrh_fc.weight) == 1
 
     def __call__(self, gru: pax.GRU):
         """
         Apply mask after an optimization step
         """
-        zr_masked_weights = jnp.where(self.xh_zr_fc_mask, gru.xh_zr_fc.weight, 0)
-        gru = gru.replace_node(gru.xh_zr_fc.weight, zr_masked_weights)
-        h_masked_weights = jnp.where(self.xh_h_fc_mask, gru.xh_h_fc.weight, 0)
-        gru = gru.replace_node(gru.xh_h_fc.weight, h_masked_weights)
+        zrh_masked_weights = jnp.where(self.h_zrh_fc_mask, gru.h_zrh_fc.weight, 0)
+        gru = gru.replace_node(gru.h_zrh_fc.weight, zrh_masked_weights)
         return gru
 
     def update_mask(self, step, gru: pax.GRU):
         """
         Update internal masks
         """
-        xh_z_weight, xh_r_weight = jnp.split(gru.xh_zr_fc.weight, 2, axis=1)
-        xh_z_weight = self.prune(step, xh_z_weight)
-        xh_r_weight = self.prune(step, xh_r_weight)
-        self.xh_zr_fc_mask *= jnp.concatenate((xh_z_weight, xh_r_weight), axis=1)
-        self.xh_h_fc_mask *= self.prune(step, gru.xh_h_fc.weight)
+        z_weight, r_weight, h_weight = jnp.split(gru.h_zrh_fc.weight, 3, axis=1)
+        z_mask = self.prune(step, z_weight)
+        r_mask = self.prune(step, r_weight)
+        h_mask = self.prune(step, h_weight)
+        self.h_zrh_fc_mask *= jnp.concatenate((z_mask, r_mask, h_mask), axis=1)
 
 
 class LinearPruner(Pruner):
-    def __init__(self, linear, update_freq=500):
-        super().__init__(update_freq=update_freq)
+    def __init__(self, linear):
+        super().__init__()
         self.mask = jnp.ones_like(linear.weight) == 1
 
     def __call__(self, linear: pax.Linear):
@@ -166,16 +201,16 @@ class LinearPruner(Pruner):
 
 class WaveGRU(pax.Module):
     """
-    WaveGRU vocoder model
+    WaveGRU vocoder model.
     """
 
-    def __init__(
-        self, mel_dim=80, embed_dim=32, rnn_dim=512, upsample_factors=(5, 4, 3, 5)
-    ):
+    def __init__(self, mel_dim=80, rnn_dim=512, upsample_factors=(5, 3, 20)):
         super().__init__()
-        self.embed = pax.Embed(256, embed_dim)
-        self.upsample = Upsample(input_dim=mel_dim, upsample_factors=upsample_factors)
-        self.rnn = pax.GRU(embed_dim + rnn_dim, rnn_dim)
+        self.embed = pax.Embed(256, 3 * rnn_dim)
+        self.upsample = Upsample(
+            input_dim=mel_dim, rnn_dim=rnn_dim, upsample_factors=upsample_factors
+        )
+        self.rnn = GRU(rnn_dim)
         self.o1 = pax.Linear(rnn_dim, rnn_dim)
         self.o2 = pax.Linear(rnn_dim, 256)
         self.gru_pruner = GRUPruner(self.rnn)
@@ -188,31 +223,30 @@ class WaveGRU(pax.Module):
         x = self.o2(x)
         return x
 
-    @jax.jit
-    def inference_step(self, rnn_state, mel, rng_key, x):
-        """one inference step"""
-        x = self.embed(x)
-        x = jnp.concatenate((x, mel), axis=-1)
-        rnn_state, x = self.rnn(rnn_state, x)
-        x = self.output(x)
-        rng_key, next_rng_key = jax.random.split(rng_key, 2)
-        x = jax.random.categorical(rng_key, x, axis=-1)
-        return rnn_state, next_rng_key, x
-
     def inference(self, mel, no_gru=False, seed=42):
         """
         generate waveform form melspectrogram
         """
 
-        y = self.upsample(mel)
+        @jax.jit
+        def step(rnn_state, mel, rng_key, x):
+            x = self.embed(x)
+            x = x + mel
+            rnn_state, x = self.rnn(rnn_state, x)
+            x = self.output(x)
+            rng_key, next_rng_key = jax.random.split(rng_key, 2)
+            x = jax.random.categorical(rng_key, x, axis=-1)
+            return rnn_state, next_rng_key, x
+
+        y = self.upsample(mel, no_repeat=no_gru)
         if no_gru:
             return y
         x = jnp.array([127], dtype=jnp.int32)
         rnn_state = self.rnn.initial_state(1)
         output = []
         rng_key = jax.random.PRNGKey(seed)
-        for i in range(y.shape[1]):
-            rnn_state, rng_key, x = self.inference_step(rnn_state, y[:, i], rng_key, x)
+        for i in tqdm(range(y.shape[1])):
+            rnn_state, rng_key, x = step(rnn_state, y[:, i], rng_key, x)
             output.append(x)
         x = jnp.concatenate(output, axis=0)
         return x
@@ -223,7 +257,7 @@ class WaveGRU(pax.Module):
         pad_left = (x.shape[1] - y.shape[1]) // 2
         pad_right = x.shape[1] - y.shape[1] - pad_left
         x = x[:, pad_left:-pad_right]
-        x = jnp.concatenate((x, y), axis=-1)
+        x = x + y
         _, x = pax.scan(
             self.rnn,
             self.rnn.initial_state(x.shape[0]),

wavegru.yaml

@@ -11,4 +11,4 @@ embed_dim: 32
 rnn_dim: 512
 frames_per_sequence: 67
 num_pad_frames: 62
-upsample_factors: [5, 4, 3, 5]
+upsample_factors: [5, 3, 20]

wavegru_cpp.py

@@ -1,18 +1,13 @@
 import numpy as np
-import sys
 from wavegru_mod import WaveGRU
 
 
 def extract_weight_mask(net):
     data = {}
     data["embed_weight"] = net.embed.weight
-    data["gru_xh_zr_weight"] = net.rnn.xh_zr_fc.weight
-    data["gru_xh_zr_mask"] = net.gru_pruner.xh_zr_fc_mask
-    data["gru_xh_zr_bias"] = net.rnn.xh_zr_fc.bias
-
-    data["gru_xh_h_weight"] = net.rnn.xh_h_fc.weight
-    data["gru_xh_h_mask"] = net.gru_pruner.xh_h_fc_mask
-    data["gru_xh_h_bias"] = net.rnn.xh_h_fc.bias
+    data["gru_h_zrh_weight"] = net.rnn.h_zrh_fc.weight
+    data["gru_h_zrh_mask"] = net.gru_pruner.h_zrh_fc_mask
+    data["gru_h_zrh_bias"] = net.rnn.h_zrh_fc.bias
 
     data["o1_weight"] = net.o1.weight
     data["o1_mask"] = net.o1_pruner.mask
@@ -23,31 +18,16 @@ def extract_weight_mask(net):
     return data
 
 
-def load_wavegru_cpp(data):
+def load_wavegru_cpp(data, repeat_factor):
+    """load wavegru weight to cpp object"""
     embed = data["embed_weight"]
-    embed_dim = embed.shape[1]
-    rnn_dim = data["gru_xh_h_bias"].shape[0]
-    input_dim = data["gru_xh_zr_weight"].shape[1] - rnn_dim
-    net = WaveGRU(input_dim, embed_dim, rnn_dim)
+    rnn_dim = data["gru_h_zrh_bias"].shape[0] // 3
+    net = WaveGRU(rnn_dim, repeat_factor)
     net.load_embed(embed)
-    dim = embed_dim + input_dim + rnn_dim
-    z, r = np.split(data["gru_xh_zr_weight"].T, 2, axis=0)
-    h = data["gru_xh_h_weight"].T
-    z = np.ascontiguousarray(z)
-    r = np.ascontiguousarray(r)
-    h = np.ascontiguousarray(h)
-
-    b1, b2 = np.split(data["gru_xh_zr_bias"], 2)
-    b3 = data["gru_xh_h_bias"]
-    m1, m2, m3 = z, r, h
-
-    mask_z, mask_r = np.split(data["gru_xh_zr_mask"].T, 2, axis=0)
-    mask_h = data["gru_xh_h_mask"].T
-    mask_z = np.ascontiguousarray(mask_z)
-    mask_r = np.ascontiguousarray(mask_r)
-    mask_h = np.ascontiguousarray(mask_h)
 
-    mask1, mask2, mask3 = mask_z, mask_r, mask_h
+    m = np.ascontiguousarray(data["gru_h_zrh_weight"].T)
+    mask = np.ascontiguousarray(data["gru_h_zrh_mask"].T)
+    b = data["gru_h_zrh_bias"]
 
     o1 = np.ascontiguousarray(data["o1_weight"].T)
     masko1 = np.ascontiguousarray(data["o1_mask"].T)
@@ -57,22 +37,6 @@ def load_wavegru_cpp(data):
     masko2 = np.ascontiguousarray(data["o2_mask"].T)
     o2b = data["o2_bias"]
 
-    net.load_weights(
-        m1,
-        mask1,
-        b1,
-        m2,
-        mask2,
-        b2,
-        m3,
-        mask3,
-        b3,
-        o1,
-        masko1,
-        o1b,
-        o2,
-        masko2,
-        o2b,
-    )
+    net.load_weights(m, mask, b, o1, masko1, o1b, o2, masko2, o2b)
 
     return net

wavegru_mod.cc

@@ -30,12 +30,11 @@ mat create_mat(int h, int w) {
 }
 
 struct WaveGRU {
-  int input_dim;
-  int embed_dim;
   int hidden_dim;
-  mat m1, m2, m3;
-  vec b1, b2, b3;
-  vec z, r, hh;
+  int repeat_factor;
+  mat m;
+  vec b;
+  vec z, r, hh, zrh;
   vec fco1, fco2;
   vec o1b, o2b;
   vec t;
@@ -43,30 +42,26 @@ struct WaveGRU {
   mat o1, o2;
   std::vector<vec> embed;
 
-  WaveGRU(int input_dim, int embed_dim, int hidden_dim)
-      : input_dim(input_dim),
-        embed_dim(embed_dim),
-        hidden_dim(hidden_dim),
-        b1(hidden_dim),
-        b2(hidden_dim),
-        b3(hidden_dim),
+  WaveGRU(int hidden_dim, int repeat_factor)
+      : hidden_dim(hidden_dim),
+        repeat_factor(repeat_factor),
+        b(3*hidden_dim),
+        t(3*hidden_dim),
+        zrh(3*hidden_dim),
         z(hidden_dim),
         r(hidden_dim),
         hh(hidden_dim),
         fco1(hidden_dim),
         fco2(256),
-        t(hidden_dim + input_dim + embed_dim),
         h(hidden_dim),
         o1b(hidden_dim),
         o2b(256) {
-    m1 = create_mat(input_dim + hidden_dim + embed_dim, hidden_dim);
-    m2 = create_mat(input_dim + hidden_dim + embed_dim, hidden_dim);
-    m3 = create_mat(input_dim + hidden_dim + embed_dim, hidden_dim);
+    m = create_mat(hidden_dim, 3*hidden_dim);
     o1 = create_mat(hidden_dim, hidden_dim);
     o2 = create_mat(hidden_dim, 256);
     embed = std::vector<vec>();
     for (int i = 0; i < 256; i++) {
-      embed.emplace_back(embed_dim);
+      embed.emplace_back(hidden_dim * 3);
       embed[i].FillRandom();
     }
   }
@@ -74,7 +69,7 @@ struct WaveGRU {
   void load_embed(fndarray embed_weights) {
     auto a_embed = embed_weights.unchecked<2>();
     for (int i = 0; i < 256; i++) {
-      for (int j = 0; j < embed_dim; j++) embed[i][j] = a_embed(i, j);
+      for (int j = 0; j < hidden_dim * 3; j++) embed[i][j] = a_embed(i, j);
     }
   }
 
@@ -90,43 +85,35 @@ struct WaveGRU {
     return mmm;
   }
 
-  void load_weights(fndarray m1, indarray m1_mask, fndarray b1, fndarray m2,
-                    indarray m2_mask, fndarray b2, fndarray m3,
-                    indarray m3_mask, fndarray b3, fndarray o1,
-                    indarray o1_mask, fndarray o1b, fndarray o2,
+  void load_weights(fndarray m, indarray m_mask, fndarray b,
+                    fndarray o1, indarray o1_mask, 
+                    fndarray o1b, fndarray o2,
                     indarray o2_mask, fndarray o2b) {
-    this->m1 = load_linear(this->b1, m1, m1_mask, b1);
-    this->m2 = load_linear(this->b2, m2, m2_mask, b2);
-    this->m3 = load_linear(this->b3, m3, m3_mask, b3);
+    this->m = load_linear(this->b, m, m_mask, b);
     this->o1 = load_linear(this->o1b, o1, o1_mask, o1b);
     this->o2 = load_linear(this->o2b, o2, o2_mask, o2b);
   }
 
   std::vector<int> inference(fndarray ft, float temperature) {
     auto rft = ft.unchecked<2>();
-    std::vector<vec> xs;
-    for (int i = 0; i < rft.shape(0); i++) {
-      xs.emplace_back(input_dim);
-      for (int j = 0; j < input_dim; j++) xs[i][j] = rft(i, j);
-    }
-
     int value = 127;
-    std::vector<int> signal(xs.size());
+    std::vector<int> signal(rft.shape(0) * repeat_factor);
     h.FillZero();
-    for (int index = 0; index < xs.size(); index++) {
-      for (int i = 0; i < embed_dim; i++) t[i] = embed[value][i];
-      for (int i = 0; i < input_dim; i++) t[embed_dim + i] = xs[index][i];
-      for (int i = 0; i < hidden_dim; i++) t[embed_dim + input_dim + i] = h[i];
-      m1.SpMM_bias(t, b1, &z, false);
-      m2.SpMM_bias(t, b2, &r, false);
+    for (int index = 0; index < signal.size(); index++) {
+      m.SpMM_bias(h, b, &zrh, false);
+
+      for (int i = 0; i < 3 * hidden_dim; i++) t[i] = embed[value][i] + rft(index / repeat_factor, i);
+      for (int i = 0; i < hidden_dim; i++) {
+        z[i] = zrh[i] + t[i];
+        r[i] = zrh[hidden_dim + i] + t[hidden_dim + i];
+      }
+
       z.Sigmoid();
       r.Sigmoid();
 
       for (int i = 0; i < hidden_dim; i++) {
-        t[embed_dim + input_dim + i] = h[i] * r[i];
+        hh[i] = zrh[hidden_dim * 2 + i]  * r[i] + t[hidden_dim * 2 + i];
       }
-
-      m3.SpMM_bias(t, b3, &hh, false);
       hh.Tanh();
       for (int i = 0; i < hidden_dim; i++) {
         h[i] = (1. - z[i]) * h[i] + z[i] * hh[i];
@@ -142,7 +129,7 @@ struct WaveGRU {
 
 PYBIND11_MODULE(wavegru_mod, m) {
   py::class_<WaveGRU>(m, "WaveGRU")
-      .def(py::init<int, int, int>())
+      .def(py::init<int, int>())
       .def("load_embed", &WaveGRU::load_embed)
       .def("load_weights", &WaveGRU::load_weights)
       .def("inference", &WaveGRU::inference);

wavegru_mod.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d5fd47ca8d441ba204b3fcff514341918d5795618ad2740704d09f62c5b6c47a
-size 527976
+oid sha256:12c27f0ea07f8da3a3ab48bc01bb0f68971ce7d57b19ada87669eab138623a9c
+size 525536