|
|
|
|
|
|
|
|
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
import numpy as np |
|
|
|
|
|
class ResBlock(nn.Module): |
|
def __init__(self, dims): |
|
super().__init__() |
|
self.conv1 = nn.Conv1d(dims, dims, kernel_size=1, bias=False) |
|
self.conv2 = nn.Conv1d(dims, dims, kernel_size=1, bias=False) |
|
self.batch_norm1 = nn.BatchNorm1d(dims) |
|
self.batch_norm2 = nn.BatchNorm1d(dims) |
|
|
|
def forward(self, x): |
|
residual = x |
|
x = self.conv1(x) |
|
x = self.batch_norm1(x) |
|
x = F.relu(x) |
|
x = self.conv2(x) |
|
x = self.batch_norm2(x) |
|
x = x + residual |
|
return x |
|
|
|
|
|
class MelResNet(nn.Module): |
|
def __init__(self, res_blocks, in_dims, compute_dims, res_out_dims, pad): |
|
super().__init__() |
|
kernel_size = pad * 2 + 1 |
|
self.conv_in = nn.Conv1d( |
|
in_dims, compute_dims, kernel_size=kernel_size, bias=False |
|
) |
|
self.batch_norm = nn.BatchNorm1d(compute_dims) |
|
self.layers = nn.ModuleList() |
|
for i in range(res_blocks): |
|
self.layers.append(ResBlock(compute_dims)) |
|
self.conv_out = nn.Conv1d(compute_dims, res_out_dims, kernel_size=1) |
|
|
|
def forward(self, x): |
|
x = self.conv_in(x) |
|
x = self.batch_norm(x) |
|
x = F.relu(x) |
|
for f in self.layers: |
|
x = f(x) |
|
x = self.conv_out(x) |
|
return x |
|
|
|
|
|
class Stretch2d(nn.Module): |
|
def __init__(self, x_scale, y_scale): |
|
super().__init__() |
|
self.x_scale = x_scale |
|
self.y_scale = y_scale |
|
|
|
def forward(self, x): |
|
b, c, h, w = x.size() |
|
x = x.unsqueeze(-1).unsqueeze(3) |
|
x = x.repeat(1, 1, 1, self.y_scale, 1, self.x_scale) |
|
return x.view(b, c, h * self.y_scale, w * self.x_scale) |
|
|
|
|
|
class UpsampleNetwork(nn.Module): |
|
def __init__( |
|
self, feat_dims, upsample_scales, compute_dims, res_blocks, res_out_dims, pad |
|
): |
|
super().__init__() |
|
total_scale = np.cumproduct(upsample_scales)[-1] |
|
self.indent = pad * total_scale |
|
self.resnet = MelResNet(res_blocks, feat_dims, compute_dims, res_out_dims, pad) |
|
self.resnet_stretch = Stretch2d(total_scale, 1) |
|
self.up_layers = nn.ModuleList() |
|
for scale in upsample_scales: |
|
kernel_size = (1, scale * 2 + 1) |
|
padding = (0, scale) |
|
stretch = Stretch2d(scale, 1) |
|
conv = nn.Conv2d(1, 1, kernel_size=kernel_size, padding=padding, bias=False) |
|
conv.weight.data.fill_(1.0 / kernel_size[1]) |
|
self.up_layers.append(stretch) |
|
self.up_layers.append(conv) |
|
|
|
def forward(self, m): |
|
aux = self.resnet(m).unsqueeze(1) |
|
aux = self.resnet_stretch(aux) |
|
aux = aux.squeeze(1) |
|
m = m.unsqueeze(1) |
|
for f in self.up_layers: |
|
m = f(m) |
|
m = m.squeeze(1)[:, :, self.indent : -self.indent] |
|
return m.transpose(1, 2), aux.transpose(1, 2) |
|
|
|
|
|
class WaveRNN(nn.Module): |
|
def __init__(self, cfg): |
|
super().__init__() |
|
|
|
self.cfg = cfg |
|
self.pad = self.cfg.VOCODER.MEL_FRAME_PAD |
|
|
|
if self.cfg.VOCODER.MODE == "mu_law_quantize": |
|
self.n_classes = 2**self.cfg.VOCODER.BITS |
|
elif self.cfg.VOCODER.MODE == "mu_law" or self.cfg.VOCODER: |
|
self.n_classes = 30 |
|
|
|
self._to_flatten = [] |
|
|
|
self.rnn_dims = self.cfg.VOCODER.RNN_DIMS |
|
self.aux_dims = self.cfg.VOCODER.RES_OUT_DIMS // 4 |
|
self.hop_length = self.cfg.VOCODER.HOP_LENGTH |
|
self.fc_dims = self.cfg.VOCODER.FC_DIMS |
|
self.upsample_factors = self.cfg.VOCODER.UPSAMPLE_FACTORS |
|
self.feat_dims = self.cfg.VOCODER.INPUT_DIM |
|
self.compute_dims = self.cfg.VOCODER.COMPUTE_DIMS |
|
self.res_out_dims = self.cfg.VOCODER.RES_OUT_DIMS |
|
self.res_blocks = self.cfg.VOCODER.RES_BLOCKS |
|
|
|
self.upsample = UpsampleNetwork( |
|
self.feat_dims, |
|
self.upsample_factors, |
|
self.compute_dims, |
|
self.res_blocks, |
|
self.res_out_dims, |
|
self.pad, |
|
) |
|
self.I = nn.Linear(self.feat_dims + self.aux_dims + 1, self.rnn_dims) |
|
|
|
self.rnn1 = nn.GRU(self.rnn_dims, self.rnn_dims, batch_first=True) |
|
self.rnn2 = nn.GRU( |
|
self.rnn_dims + self.aux_dims, self.rnn_dims, batch_first=True |
|
) |
|
self._to_flatten += [self.rnn1, self.rnn2] |
|
|
|
self.fc1 = nn.Linear(self.rnn_dims + self.aux_dims, self.fc_dims) |
|
self.fc2 = nn.Linear(self.fc_dims + self.aux_dims, self.fc_dims) |
|
self.fc3 = nn.Linear(self.fc_dims, self.n_classes) |
|
|
|
self.num_params() |
|
|
|
self._flatten_parameters() |
|
|
|
def forward(self, x, mels): |
|
device = next(self.parameters()).device |
|
|
|
self._flatten_parameters() |
|
|
|
batch_size = x.size(0) |
|
h1 = torch.zeros(1, batch_size, self.rnn_dims, device=device) |
|
h2 = torch.zeros(1, batch_size, self.rnn_dims, device=device) |
|
mels, aux = self.upsample(mels) |
|
|
|
aux_idx = [self.aux_dims * i for i in range(5)] |
|
a1 = aux[:, :, aux_idx[0] : aux_idx[1]] |
|
a2 = aux[:, :, aux_idx[1] : aux_idx[2]] |
|
a3 = aux[:, :, aux_idx[2] : aux_idx[3]] |
|
a4 = aux[:, :, aux_idx[3] : aux_idx[4]] |
|
|
|
x = torch.cat([x.unsqueeze(-1), mels, a1], dim=2) |
|
x = self.I(x) |
|
res = x |
|
x, _ = self.rnn1(x, h1) |
|
|
|
x = x + res |
|
res = x |
|
x = torch.cat([x, a2], dim=2) |
|
x, _ = self.rnn2(x, h2) |
|
|
|
x = x + res |
|
x = torch.cat([x, a3], dim=2) |
|
x = F.relu(self.fc1(x)) |
|
|
|
x = torch.cat([x, a4], dim=2) |
|
x = F.relu(self.fc2(x)) |
|
return self.fc3(x) |
|
|
|
def num_params(self, print_out=True): |
|
parameters = filter(lambda p: p.requires_grad, self.parameters()) |
|
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000 |
|
if print_out: |
|
print("Trainable Parameters: %.3fM" % parameters) |
|
return parameters |
|
|
|
def _flatten_parameters(self): |
|
[m.flatten_parameters() for m in self._to_flatten] |
|
|
