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
# set learning rate
for g in optimizer.param_groups:
g['lr'] = lr
total_iters = len(train_set)
epochs = (total_steps - model.get_step()) // total_iters + 1
total_number_of_batches = len(train_set)
writer = SummaryWriter("runs/{0}-{1}".format(
model_name_prefix,
datetime.now().strftime("%Y%m%d-%H%M%S")))
scheduler = StepLR(optimizer, step_size=1, gamma=0.983)
for e in range(EPOCH, epochs + 1):
start = time.time()
running_loss = 0.
avg_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)
optimizer.step()
running_loss += loss.item()
avg_loss = running_loss / i
speed = i / (time.time() - start)
step = model.get_step()
k = step // 1000
# Write to tensorboard per batch
writer.add_scalar('Epoch loss', loss.item(),
e * total_number_of_batches + i)
msg = f'| Epoch: {e}/{epochs} ({i}/{total_iters}) | Loss: {avg_loss:.4f} | {speed:.1f} steps/s | Step: {k}k | '
stream(msg)
"""
####################### Testing ############################
torch.cuda.empty_cache()
loss_test = 0
for _, (x_test, y_test, m_test) in enumerate(test_set, 1):
x_test, m_test, y_test = x_test.to(device), m_test.to(device), y_test.to(device)
if device.type == 'cuda' and torch.cuda.device_count() > 1:
raise RuntimeError("Unsupported")
else:
y_test_hat = model(x_test, m_test)
if model.mode == 'RAW':
y_test_hat = y_test_hat.transpose(1, 2).unsqueeze(-1)
elif model.mode == 'MOL':
y_test = y_test.float()
y_test = y_test.unsqueeze(-1)
loss_test += loss_func(y_test_hat, y_test).item()
avg_loss_test = loss_test / len(test_set)
msg = f'| Epoch: {e}/{epochs} | Test-Loss: {loss_test:.4f} | Test-AvgLoss: {avg_loss_test:.4f} | '
stream("\n")
stream(msg)
writer.add_scalar('Test loss', loss_test, e)
writer.add_scalar('Average test loss', avg_loss_test, e)
############################################################
"""
# Write to tensorboard per epoch
writer.add_scalar('Running loss', running_loss, e)
writer.add_scalar('Average loss', avg_loss, e)
# Must save latest optimizer state to ensure that resuming training
# doesn't produce artifacts
save_checkpoint('voc',
paths,
model,
optimizer,
name="{0}-epoch-{1}-loss-{2}".format(
model_name_prefix, e, avg_loss),
is_silent=True)
model.log(paths.voc_log, msg)
print(' ')
scheduler.step()
print('Epoch:', e, 'LR:', scheduler.get_lr())
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(' ')
def voc_train_loop(paths: Paths, model: WaveRNN, loss_func, optimizer,
train_set, test_set, init_lr, final_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):
adjust_learning_rate(optimizer, e, epochs, init_lr,
final_lr) # 初始学习率与最终学习率-Begee
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) # x/y: (Batch, sub_bands, T)
######################### MultiBand-WaveRNN #########################
if hp.voc_multiband:
y0 = y[:, 0, :].squeeze(0).unsqueeze(
-1) # y0/y1/y2/y3: (Batch, T, 1)
y1 = y[:, 1, :].squeeze(0).unsqueeze(-1)
y2 = y[:, 2, :].squeeze(0).unsqueeze(-1)
y3 = y[:, 3, :].squeeze(0).unsqueeze(-1)
y_hat = model(x, m) # (Batch, T, num_classes, sub_bands)
if model.mode == 'RAW':
y_hat0 = y_hat[:, :, :, 0].transpose(1, 2).unsqueeze(
-1) # (Batch, num_classes, T, 1)
y_hat1 = y_hat[:, :, :, 1].transpose(1, 2).unsqueeze(-1)
y_hat2 = y_hat[:, :, :, 2].transpose(1, 2).unsqueeze(-1)
y_hat3 = y_hat[:, :, :, 3].transpose(1, 2).unsqueeze(-1)
elif model.mode == 'MOL':
y0 = y0.float()
y1 = y1.float()
y2 = y2.float()
y3 = y3.float()
loss = loss_func(y_hat0, y0) + loss_func(
y_hat1, y1) + loss_func(y_hat2, y2) + loss_func(
y_hat3, y3)
######################### MultiBand-WaveRNN #########################
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).cpu()
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(' ')