import torch import torch.nn as nn import torch.nn.functional as F from utils.display import * from utils.dsp import * class WaveRNN(nn.Module) : def __init__(self, hidden_size=896, quantisation=256) : super(WaveRNN, self).__init__() self.hidden_size = hidden_size self.split_size = hidden_size // 2 # The main matmul self.R = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=False) # Output fc layers self.O1 = nn.Linear(self.split_size, self.split_size) self.O2 = nn.Linear(self.split_size, quantisation) self.O3 = nn.Linear(self.split_size, self.split_size) self.O4 = nn.Linear(self.split_size, quantisation) # Input fc layers self.I_coarse = nn.Linear(2, 3 * self.split_size, bias=False) self.I_fine = nn.Linear(3, 3 * self.split_size, bias=False) # biases for the gates self.bias_u = nn.Parameter(torch.zeros(self.hidden_size)) self.bias_r = nn.Parameter(torch.zeros(self.hidden_size)) self.bias_e = nn.Parameter(torch.zeros(self.hidden_size)) # display num params self.num_params() def forward(self, prev_y, prev_hidden, current_coarse) : # Main matmul - the projection is split 3 ways R_hidden = self.R(prev_hidden) R_u, R_r, R_e, = torch.split(R_hidden, self.hidden_size, dim=1) # Project the prev input coarse_input_proj = self.I_coarse(prev_y) I_coarse_u, I_coarse_r, I_coarse_e = \ torch.split(coarse_input_proj, self.split_size, dim=1) # Project the prev input and current coarse sample fine_input = torch.cat([prev_y, current_coarse], dim=1) fine_input_proj = self.I_fine(fine_input) I_fine_u, I_fine_r, I_fine_e = \ torch.split(fine_input_proj, self.split_size, dim=1) # concatenate for the gates I_u = torch.cat([I_coarse_u, I_fine_u], dim=1) I_r = torch.cat([I_coarse_r, I_fine_r], dim=1) I_e = torch.cat([I_coarse_e, I_fine_e], dim=1) # Compute all gates for coarse and fine u = F.sigmoid(R_u + I_u + self.bias_u) r = F.sigmoid(R_r + I_r + self.bias_r) e = torch.tanh(r * R_e + I_e + self.bias_e) hidden = u * prev_hidden + (1. - u) * e # Split the hidden state hidden_coarse, hidden_fine = torch.split(hidden, self.split_size, dim=1) # Compute outputs out_coarse = self.O2(F.relu(self.O1(hidden_coarse))) out_fine = self.O4(F.relu(self.O3(hidden_fine))) return out_coarse, out_fine, hidden def generate(self, seq_len): with torch.no_grad(): # First split up the biases for the gates b_coarse_u, b_fine_u = torch.split(self.bias_u, self.split_size) b_coarse_r, b_fine_r = torch.split(self.bias_r, self.split_size) b_coarse_e, b_fine_e = torch.split(self.bias_e, self.split_size) # Lists for the two output seqs c_outputs, f_outputs = [], [] # Some initial inputs out_coarse = torch.LongTensor([0]).cuda() out_fine = torch.LongTensor([0]).cuda() # We'll meed a hidden state hidden = self.init_hidden() # Need a clock for display start = time.time() # Loop for generation for i in range(seq_len) : # Split into two hidden states hidden_coarse, hidden_fine = \ torch.split(hidden, self.split_size, dim=1) # Scale and concat previous predictions out_coarse = out_coarse.unsqueeze(0).float() / 127.5 - 1. out_fine = out_fine.unsqueeze(0).float() / 127.5 - 1. prev_outputs = torch.cat([out_coarse, out_fine], dim=1) # Project input coarse_input_proj = self.I_coarse(prev_outputs) I_coarse_u, I_coarse_r, I_coarse_e = \ torch.split(coarse_input_proj, self.split_size, dim=1) # Project hidden state and split 6 ways R_hidden = self.R(hidden) R_coarse_u , R_fine_u, \ R_coarse_r, R_fine_r, \ R_coarse_e, R_fine_e = torch.split(R_hidden, self.split_size, dim=1) # Compute the coarse gates u = F.sigmoid(R_coarse_u + I_coarse_u + b_coarse_u) r = F.sigmoid(R_coarse_r + I_coarse_r + b_coarse_r) e = torch.tanh(r * R_coarse_e + I_coarse_e + b_coarse_e) hidden_coarse = u * hidden_coarse + (1. - u) * e # Compute the coarse output out_coarse = self.O2(F.relu(self.O1(hidden_coarse))) posterior = F.softmax(out_coarse, dim=1) distrib = torch.distributions.Categorical(posterior) out_coarse = distrib.sample() c_outputs.append(out_coarse) # Project the [prev outputs and predicted coarse sample] coarse_pred = out_coarse.float() / 127.5 - 1. fine_input = torch.cat([prev_outputs, coarse_pred.unsqueeze(0)], dim=1) fine_input_proj = self.I_fine(fine_input) I_fine_u, I_fine_r, I_fine_e = \ torch.split(fine_input_proj, self.split_size, dim=1) # Compute the fine gates u = F.sigmoid(R_fine_u + I_fine_u + b_fine_u) r = F.sigmoid(R_fine_r + I_fine_r + b_fine_r) e = torch.tanh(r * R_fine_e + I_fine_e + b_fine_e) hidden_fine = u * hidden_fine + (1. - u) * e # Compute the fine output out_fine = self.O4(F.relu(self.O3(hidden_fine))) posterior = F.softmax(out_fine, dim=1) distrib = torch.distributions.Categorical(posterior) out_fine = distrib.sample() f_outputs.append(out_fine) # Put the hidden state back together hidden = torch.cat([hidden_coarse, hidden_fine], dim=1) # Display progress speed = (i + 1) / (time.time() - start) stream('Gen: %i/%i -- Speed: %i', (i + 1, seq_len, speed)) coarse = torch.stack(c_outputs).squeeze(1).cpu().data.numpy() fine = torch.stack(f_outputs).squeeze(1).cpu().data.numpy() output = combine_signal(coarse, fine) return output, coarse, fine def init_hidden(self, batch_size=1) : return torch.zeros(batch_size, self.hidden_size).cuda() def num_params(self) : parameters = filter(lambda p: p.requires_grad, self.parameters()) parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000 print('Trainable Parameters: %.3f million' % parameters)