class Model(nn.Module):
def __init__(self,
quantization_channels=256,
gru_channels=896,
fc_channels=896,
lc_channels=80,
upsample_factor=(5, 5, 8),
use_gru_in_upsample=True):
super().__init__()
self.upsample = ConvInUpsampleNetwork(upsample_scales=upsample_factor,
upsample_activation="none",
upsample_activation_params={},
mode="nearest",
cin_channels=lc_channels,
use_gru=use_gru_in_upsample)
self.wavernn = WaveRNN(quantization_channels, gru_channels,
fc_channels, lc_channels)
def forward(self, inputs, conditions):
conditions = self.upsample(conditions.transpose(1, 2))
return self.wavernn(inputs, conditions[:, 1:, :])
def after_update(self):
self.wavernn.after_update()
def generate(self, conditions):
self.eval()
with torch.no_grad():
conditions = self.upsample(conditions.transpose(1, 2))
output = self.wavernn.generate(conditions)
self.train()
return output
class Model(nn.Module):
def __init__(self,
quantization_channels=256,
gru_channels=896,
fc_channels=896,
lc_channels=80,
lc_out_channles=80,
upsample_factor=(5, 5, 8),
use_lstm=True,
lstm_layer=2,
upsample_method='duplicate'):
super().__init__()
self.frame_net = FrameRateNet(lc_channels, lc_out_channles)
self.upsample = UpsampleNet(input_size=lc_out_channles,
output_size=lc_out_channles,
upsample_factor=upsample_factor,
use_lstm=use_lstm,
lstm_layer=lstm_layer,
upsample_method=upsample_method)
self.wavernn = WaveRNN(quantization_channels, gru_channels,
fc_channels, lc_channels)
self.num_params()
def forward(self, inputs, conditions):
conditions = self.frame_net(conditions.transpose(1, 2))
conditions = self.upsample(conditions.transpose(1, 2))
return self.wavernn(inputs, conditions[:, 1:, :])
def after_update(self):
self.wavernn.after_update()
def generate(self, conditions):
self.eval()
with torch.no_grad():
conditions = self.frame_net(conditions.transpose(1, 2))
conditions = self.upsample(conditions.transpose(1, 2))
output = self.wavernn.generate(conditions)
self.train()
return output
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)
class Model(nn.Module):
def __init__(self, rnn_dims, fc_dims, pad, upsample_factors, feat_dims):
super().__init__()
self.n_classes = 256
self.upsample = UpsampleNetwork(feat_dims, upsample_factors)
self.wavernn = WaveRNN(rnn_dims, fc_dims, feat_dims, 0)
self.num_params()
def forward(self, x, mels):
#logger.log(f'x: {x.size()} mels: {mels.size()}')
cond = self.upsample(mels)
#logger.log(f'cond: {cond.size()}')
return self.wavernn(x, cond.transpose(1, 2), None, None, None)
def after_update(self):
self.wavernn.after_update()
def preview_upsampling(self, mels):
return self.upsample(mels)
def forward_generate(self,
mels,
deterministic=False,
use_half=False,
verbose=False):
n = mels.size(0)
if use_half:
mels = mels.half()
self.eval()
with torch.no_grad():
cond = self.upsample(mels)
output = self.wavernn.generate(cond.transpose(1, 2),
None,
None,
None,
use_half=use_half,
verbose=verbose)
self.train()
return output
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
logger.log('Trainable Parameters: %.3f million' % parameters)
def load_state_dict(self, dict):
return super().load_state_dict(upgrade_state_dict(dict))
def do_train(self,
paths,
dataset,
optimiser,
epochs,
batch_size,
step,
lr=1e-4,
valid_index=[],
use_half=False):
if use_half:
import apex
optimiser = apex.fp16_utils.FP16_Optimizer(optimiser,
dynamic_loss_scale=True)
for p in optimiser.param_groups:
p['lr'] = lr
criterion = nn.NLLLoss().cuda()
k = 0
saved_k = 0
print(win_length, hop_length, win_length / hop_length)
for e in range(epochs):
# trn_loader = DataLoader(dataset, collate_fn=lambda batch: env.collate(0, int( win_length/hop_length), 0, batch), batch_size=batch_size,
# num_workers=2, shuffle=True, pin_memory=True)
trn_loader = DataLoader(
dataset,
collate_fn=lambda batch: env.collate(0, 16, 0, batch),
batch_size=batch_size,
num_workers=2,
shuffle=True,
pin_memory=True)
start = time.time()
running_loss_c = 0.
running_loss_f = 0.
iters = len(trn_loader)
for i, (mels, coarse, fine, coarse_f,
fine_f) in enumerate(trn_loader):
mels, coarse, fine, coarse_f, fine_f = mels.cuda(
), coarse.cuda(), fine.cuda(), coarse_f.cuda(), fine_f.cuda()
coarse, fine, coarse_f, fine_f = [
t[:, hop_length:1 - hop_length]
for t in [coarse, fine, coarse_f, fine_f]
]
if use_half:
mels = mels.half()
coarse_f = coarse_f.half()
fine_f = fine_f.half()
x = torch.cat([
coarse_f[:, :-1].unsqueeze(-1),
fine_f[:, :-1].unsqueeze(-1), coarse_f[:, 1:].unsqueeze(-1)
],
dim=2)
p_c, p_f, _h_n = self(x, mels)
loss_c = criterion(p_c.transpose(1, 2).float(), coarse[:, 1:])
loss_f = criterion(p_f.transpose(1, 2).float(), fine[:, 1:])
loss = loss_c + loss_f
optimiser.zero_grad()
if use_half:
optimiser.backward(loss)
else:
loss.backward()
optimiser.step()
running_loss_c += loss_c.item()
running_loss_f += loss_f.item()
self.after_update()
speed = (i + 1) / (time.time() - start)
avg_loss_c = running_loss_c / (i + 1)
avg_loss_f = running_loss_f / (i + 1)
step += 1
k = step // 1000
logger.status(
f'Epoch: {e+1}/{epochs} -- Batch: {i+1}/{iters} -- Loss: c={avg_loss_c:#.4} f={avg_loss_f:#.4} -- Speed: {speed:#.4} steps/sec -- Step: {k}k '
)
os.makedirs(paths.checkpoint_dir, exist_ok=True)
torch.save(self.state_dict(), paths.model_path())
np.save(paths.step_path(), step)
logger.log_current_status()
logger.log(
f' <saved>; w[0][0] = {self.wavernn.gru.weight_ih_l0[0][0]}')
if k > saved_k + 50:
torch.save(self.state_dict(), paths.model_hist_path(step))
saved_k = k
self.do_generate(paths,
step,
dataset.path,
valid_index,
use_half=use_half)
def do_generate(self,
paths,
step,
data_path,
test_index,
deterministic=False,
use_half=False,
verbose=False):
k = step // 1000
test_mels = [np.load(f'{data_path}/mel/{id}.npy') for id in test_index]
maxlen = max([x.shape[1] for x in test_mels])
aligned = [
torch.cat([
torch.FloatTensor(x),
torch.zeros(80, maxlen - x.shape[1] + 1)
],
dim=1) for x in test_mels
]
print(torch.stack(aligned).size())
out = self.forward_generate(torch.stack(aligned).cuda(),
deterministic,
use_half=use_half,
verbose=verbose)
os.makedirs(paths.gen_path(), exist_ok=True)
for i, id in enumerate(test_index):
gt = np.load(f'{data_path}/quant/{id}.npy')
gt = (gt.astype(np.float32) + 0.5) / (2**15 - 0.5)
librosa.output.write_wav(
f'{paths.gen_path()}/{k}k_steps_{i}_target.wav',
gt,
sr=sample_rate)
audio = out[i][:len(gt)].cpu().numpy()
librosa.output.write_wav(
f'{paths.gen_path()}/{k}k_steps_{i}_generated.wav',
audio,
sr=sample_rate)