def gen_from_file(model: WaveRNN, load_path: Path, save_path: Path, batched, target, overlap):
k = model.get_step() // 1000
file_name = load_path.stem
suffix = load_path.suffix
if suffix == ".wav":
wav = load_wav(load_path)
save_wav(wav, save_path/f'__{file_name}__{k}k_steps_target.wav')
mel = melspectrogram(wav)
elif suffix == ".npy":
mel = np.load(load_path)
if mel.ndim != 2 or mel.shape[0] != hp.num_mels:
raise ValueError(f'Expected a numpy array shaped (n_mels, n_hops), but got {wav.shape}!')
_max = np.max(mel)
_min = np.min(mel)
if _max >= 1.01 or _min <= -0.01:
raise ValueError(f'Expected spectrogram range in [0,1] but was instead [{_min}, {_max}]')
else:
raise ValueError(f"Expected an extension of .wav or .npy, but got {suffix}!")
mel = torch.tensor(mel).unsqueeze(0)
batch_str = f'gen_batched_target{target}_overlap{overlap}' if batched else 'gen_NOT_BATCHED'
save_str = save_path/f'__{file_name}__{k}k_steps_{batch_str}.wav'
_ = model.generate(mel, save_str, batched, target, overlap, hp.mu_law)
def gen_testset(model: WaveRNN, test_set, samples, batched, target, overlap, save_path: Path):
k = model.get_step() // 1000
for i, (m, x) in enumerate(test_set, 1):
if i > samples: break
print('\n| Generating: %i/%i' % (i, samples))
x = x[0].numpy()
bits = 16 if hp.voc_mode == 'MOL' else hp.bits
if hp.mu_law and hp.voc_mode != 'MOL':
x = decode_mu_law(x, 2**bits, from_labels=True)
else:
x = label_2_float(x, bits)
save_wav(x, save_path/f'{k}k_steps_{i}_target.wav')
batch_str = f'gen_batched_target{target}_overlap{overlap}' if batched else 'gen_NOT_BATCHED'
save_str = str(save_path/f'{k}k_steps_{i}_{batch_str}.wav')
_ = model.generate(m, save_str, batched, target, overlap, hp.mu_law)
def init_wavernn(args):
# wavernn
print('\n#####################################')
if args.vocoder == 'wavernn' or args.vocoder == 'wr':
wavernn_hp.configure(args.hp_file)
paths = Paths(wavernn_hp.data_path, wavernn_hp.voc_model_id,
wavernn_hp.tts_model_id)
if not args.force_cpu and torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
print('Using device:', device)
print('\nInitialising WaveRNN Model...\n')
# Instantiate WaveRNN Model
voc_model = WaveRNN(rnn_dims=wavernn_hp.voc_rnn_dims,
fc_dims=wavernn_hp.voc_fc_dims,
bits=wavernn_hp.bits,
pad=wavernn_hp.voc_pad,
upsample_factors=wavernn_hp.voc_upsample_factors,
feat_dims=wavernn_hp.num_mels,
compute_dims=wavernn_hp.voc_compute_dims,
res_out_dims=wavernn_hp.voc_res_out_dims,
res_blocks=wavernn_hp.voc_res_blocks,
hop_length=wavernn_hp.hop_length,
sample_rate=wavernn_hp.sample_rate,
mode=wavernn_hp.voc_mode).to(device)
voc_load_path = args.voc_weights if args.voc_weights else paths.voc_latest_weights
voc_model.load(voc_load_path)
voc_k = voc_model.get_step() // 1000
simple_table([
('Vocoder Type', 'WaveRNN'), ('WaveRNN', str(voc_k) + 'k'),
('Generation Mode',
'Batched' if wavernn_hp.voc_gen_batched else 'Unbatched'),
('Target Samples',
wavernn_hp.voc_target if wavernn_hp.voc_gen_batched else 'N/A'),
('Overlap Samples',
wavernn_hp.voc_overlap if wavernn_hp.voc_gen_batched else 'N/A')
])
if args.vocoder == 'griffinlim' or args.vocoder == 'gl':
v_type = args.vocoder
elif (args.vocoder == 'wavernn'
or args.vocoder == 'wr') and wavernn_hp.voc_gen_batched:
v_type = 'wavernn_batched'
else:
v_type = 'wavernn_unbatched'
else:
return None, None, None
return voc_model, voc_k, v_type
def main():
# Parse Arguments
parser = argparse.ArgumentParser(description='Train WaveRNN Vocoder')
parser.add_argument('--lr',
'-l',
type=float,
help='[float] override hparams.py learning rate')
parser.add_argument('--batch_size',
'-b',
type=int,
help='[int] override hparams.py batch size')
parser.add_argument('--force_train',
'-f',
action='store_true',
help='Forces the model to train past total steps')
parser.add_argument('--gta',
'-g',
default=True,
action='store_true',
help='train wavernn on GTA features')
parser.add_argument(
'--force_cpu',
'-c',
action='store_true',
help='Forces CPU-only training, even when in CUDA capable environment')
parser.add_argument('--hp_file',
metavar='FILE',
default='wavernn_vocoder/hparams.py',
help='The file to use for the hyperparameters')
args = parser.parse_args()
hp.configure(args.hp_file) # load hparams from file
if args.lr is None:
args.lr = hp.voc_lr
if args.batch_size is None:
args.batch_size = hp.voc_batch_size
paths = Paths(hp.data_path, hp.voc_model_id, hp.tts_model_id)
batch_size = args.batch_size
force_train = args.force_train
train_gta = args.gta
lr = args.lr
if not args.force_cpu and torch.cuda.is_available():
device = torch.device('cuda')
if batch_size % torch.cuda.device_count() != 0:
raise ValueError(
'`batch_size` must be evenly divisible by n_gpus!')
else:
device = torch.device('cpu')
print('Using device:', device)
print('\nInitialising Model...\n')
# Instantiate WaveRNN Model
voc_model = WaveRNN(rnn_dims=hp.voc_rnn_dims,
fc_dims=hp.voc_fc_dims,
bits=hp.bits,
pad=hp.voc_pad,
upsample_factors=hp.voc_upsample_factors,
feat_dims=hp.num_mels,
compute_dims=hp.voc_compute_dims,
res_out_dims=hp.voc_res_out_dims,
res_blocks=hp.voc_res_blocks,
hop_length=hp.hop_length,
sample_rate=hp.sample_rate,
mode=hp.voc_mode).to(device)
# Check to make sure the hop length is correctly factorised
assert np.cumprod(hp.voc_upsample_factors)[-1] == hp.hop_length
optimizer = optim.Adam(voc_model.parameters())
restore_checkpoint('voc',
paths,
voc_model,
optimizer,
create_if_missing=True)
# train_set, test_set = get_vocoder_datasets(paths.data, batch_size, train_gta)
# wavernn_data = './wavernn_training_data/'
wavernn_data = hp.data_path
train_set, test_set = get_vocoder_datasets(wavernn_data, batch_size,
train_gta)
total_steps = 10_000_000 if force_train else hp.voc_total_steps
simple_table([
('Remaining', str(
(total_steps - voc_model.get_step()) // 1000) + 'k Steps'),
('Batch Size', batch_size), ('LR', lr),
('Sequence Len', hp.voc_seq_len), ('GTA Train', train_gta)
])
loss_func = F.cross_entropy if voc_model.mode == 'RAW' else discretized_mix_logistic_loss
voc_train_loop(paths, voc_model, loss_func, optimizer, train_set, test_set,
lr, total_steps)
print('Training Complete.')
print(
'To continue training increase voc_total_steps in hparams.py or use --force_train'
)
def voc_train_loop(paths: Paths, model: WaveRNN, loss_func, optimizer,
train_set, test_set, lr, total_steps):
# Use same device as model parameters
device = next(model.parameters()).device
for g in optimizer.param_groups:
g['lr'] = lr
total_iters = len(train_set)
epochs = (total_steps - model.get_step()) // total_iters + 1
for e in range(1, epochs + 1):
start = time.time()
running_loss = 0.
for i, (x, y, m) in enumerate(train_set, 1):
x, m, y = x.to(device), m.to(device), y.to(device)
# Parallelize model onto GPUS using workaround due to python bug
if device.type == 'cuda' and torch.cuda.device_count() > 1:
y_hat = data_parallel_workaround(model, x, m)
else:
y_hat = model(x, m)
if model.mode == 'RAW':
y_hat = y_hat.transpose(1, 2).unsqueeze(-1)
elif model.mode == 'MOL':
y = y.float()
y = y.unsqueeze(-1)
loss = loss_func(y_hat, y)
optimizer.zero_grad()
loss.backward()
if hp.voc_clip_grad_norm is not None:
grad_norm = torch.nn.utils.clip_grad_norm_(
model.parameters(), hp.voc_clip_grad_norm)
if np.isnan(grad_norm):
print('grad_norm was NaN!')
optimizer.step()
running_loss += loss.item()
avg_loss = running_loss / i
speed = i / (time.time() - start)
step = model.get_step()
k = step // 1000
if step % hp.voc_checkpoint_every == 0:
gen_testset(model, test_set, hp.voc_gen_at_checkpoint,
hp.voc_gen_batched, hp.voc_target, hp.voc_overlap,
paths.voc_output)
ckpt_name = f'wave_step{k}K'
save_checkpoint('voc',
paths,
model,
optimizer,
name=ckpt_name,
is_silent=True)
msg = f'| Epoch: {e}/{epochs} ({i}/{total_iters}) | Loss: {avg_loss:.4f} | {speed:.1f} steps/s | Step: {k}k | '
stream(msg)
# Must save latest optimizer state to ensure that resuming training
# doesn't produce artifacts
save_checkpoint('voc', paths, model, optimizer, is_silent=True)
model.log(paths.voc_log, msg)
print(' ')
gta = args.gta
if not args.force_cpu and torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
print('Using device:', device)
print('\nInitialising Model...\n')
model = WaveRNN(rnn_dims=hp.voc_rnn_dims,
fc_dims=hp.voc_fc_dims,
bits=hp.bits,
pad=hp.voc_pad,
upsample_factors=hp.voc_upsample_factors,
feat_dims=hp.num_mels,
compute_dims=hp.voc_compute_dims,
res_out_dims=hp.voc_res_out_dims,
res_blocks=hp.voc_res_blocks,
hop_length=hp.hop_length,
sample_rate=hp.sample_rate,
mode=hp.voc_mode).to(device)
paths = Paths(hp.data_path, hp.voc_model_id, hp.tts_model_id)
voc_weights = args.voc_weights if args.voc_weights else paths.voc_latest_weights
model.load(voc_weights)
simple_table([('Generation Mode', 'Batched' if batched else 'Unbatched'),
('Target Samples', target if batched else 'N/A'),
('Overlap Samples', overlap if batched else 'N/A')])