NTT123
a slow but working model
df1ad02
raw
history blame
7.04 kB
"""
WaveGRU model: melspectrogram => mu-law encoded waveform
"""
import jax
import jax.numpy as jnp
import pax
class ReLU(pax.Module):
def __call__(self, x):
return jax.nn.relu(x)
def dilated_residual_conv_block(dim, kernel, stride, dilation):
"""
Use dilated convs to enlarge the receptive field
"""
return pax.Sequential(
pax.Conv1D(dim, dim, kernel, stride, dilation, "VALID", with_bias=False),
pax.LayerNorm(dim, -1, True, True),
ReLU(),
pax.Conv1D(dim, dim, 1, 1, 1, "VALID", with_bias=False),
pax.LayerNorm(dim, -1, True, True),
ReLU(),
)
def tile_1d(x, factor):
"""
Tile tensor of shape N, L, D into N, L*factor, D
"""
N, L, D = x.shape
x = x[:, :, None, :]
x = jnp.tile(x, (1, 1, factor, 1))
x = jnp.reshape(x, (N, L * factor, D))
return x
def up_block(dim, factor):
"""
Tile >> Conv >> BatchNorm >> ReLU
"""
return 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(),
)
class Upsample(pax.Module):
"""
Upsample melspectrogram to match raw audio sample rate.
"""
def __init__(self, input_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),
)
self.upsample_factors = upsample_factors
self.dilated_convs = [
dilated_residual_conv_block(512, 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.final_tile = upsample_factors[-1]
def __call__(self, x):
x = self.input_conv(x)
for residual in self.dilated_convs:
y = residual(x)
pad = (x.shape[1] - y.shape[1]) // 2
x = x[:, pad:-pad, :] + y
for f in self.up_blocks:
x = f(x)
x = tile_1d(x, self.final_tile)
return x
class Pruner(pax.Module):
"""
Base class for pruners
"""
def __init__(self, update_freq=500):
super().__init__()
self.update_freq = update_freq
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)
return z
def prune(self, step, weights):
"""
Return a mask
"""
z = self.compute_sparsity(step)
x = weights
H, W = x.shape
x = x.reshape(H // 4, 4, W // 4, 4)
x = jnp.abs(x)
x = jnp.sum(x, axis=(1, 3), keepdims=True)
q = jnp.quantile(jnp.reshape(x, (-1,)), z)
x = x >= q
x = jnp.tile(x, (1, 4, 1, 4))
x = jnp.reshape(x, (H, W))
return x
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 __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)
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)
class LinearPruner(Pruner):
def __init__(self, linear, update_freq=500):
super().__init__(update_freq=update_freq)
self.mask = jnp.ones_like(linear.weight) == 1
def __call__(self, linear: pax.Linear):
"""
Apply mask after an optimization step
"""
return linear.replace(weight=jnp.where(self.mask, linear.weight, 0))
def update_mask(self, step, linear: pax.Linear):
"""
Update internal masks
"""
self.mask *= self.prune(step, linear.weight)
class WaveGRU(pax.Module):
"""
WaveGRU vocoder model
"""
def __init__(
self, mel_dim=80, embed_dim=32, rnn_dim=512, upsample_factors=(5, 4, 3, 5)
):
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.o1 = pax.Linear(rnn_dim, rnn_dim)
self.o2 = pax.Linear(rnn_dim, 256)
self.gru_pruner = GRUPruner(self.rnn)
self.o1_pruner = LinearPruner(self.o1)
self.o2_pruner = LinearPruner(self.o2)
def output(self, x):
x = self.o1(x)
x = jax.nn.relu(x)
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)
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)
output.append(x)
x = jnp.concatenate(output, axis=0)
return x
def __call__(self, mel, x):
x = self.embed(x)
y = self.upsample(mel)
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 = pax.scan(
self.rnn,
self.rnn.initial_state(x.shape[0]),
x,
time_major=False,
)
x = self.output(x)
return x