def train(args):
os.makedirs(args.checkpoint_dir, exist_ok=True)
logging = GetLogging(args.logfile)
train_dataset = CustomerDataset(args.input,
upsample_factor=hop_length,
local_condition=True,
global_condition=False)
device = torch.device("cuda" if args.use_cuda else "cpu")
generator, discriminator = create_model(args)
print(generator)
print(discriminator)
num_gpu = torch.cuda.device_count() if args.use_cuda else 1
global_step = 0
g_parameters = list(generator.parameters())
g_optimizer = optim.Adam(g_parameters, lr=args.g_learning_rate)
d_parameters = list(discriminator.parameters())
d_optimizer = optim.Adam(d_parameters, lr=args.d_learning_rate)
writer = Writer(args.checkpoint_dir, sample_rate=sample_rate)
generator.to(device)
discriminator.to(device)
if args.resume is not None:
restore_step = attempt_to_restore(generator, discriminator,
g_optimizer, d_optimizer,
args.resume, args.use_cuda, logging)
global_step = restore_step
customer_g_optimizer = Optimizer(g_optimizer, args.g_learning_rate,
global_step, args.warmup_steps,
args.decay_learning_rate)
customer_d_optimizer = Optimizer(d_optimizer, args.d_learning_rate,
global_step, args.warmup_steps,
args.decay_learning_rate)
criterion = nn.MSELoss().to(device)
stft_criterion = MultiResolutionSTFTLoss()
for epoch in range(args.epochs):
collate = CustomerCollate(upsample_factor=hop_length,
condition_window=args.condition_window,
local_condition=True,
global_condition=False)
train_data_loader = DataLoader(train_dataset,
collate_fn=collate,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=True,
pin_memory=True)
#train one epoch
for batch, (samples, conditions) in enumerate(train_data_loader):
start = time.time()
batch_size = int(conditions.shape[0] // num_gpu * num_gpu)
samples = samples[:batch_size, :].to(device)
conditions = conditions[:batch_size, :, :].to(device)
losses = {}
if num_gpu > 1:
g_outputs = parallel(generator, (conditions, ))
else:
g_outputs = generator(conditions)
sc_loss, mag_loss = stft_criterion(g_outputs.squeeze(1),
samples.squeeze(1))
g_loss = sc_loss + mag_loss
losses['sc_loss'] = sc_loss.item()
losses['mag_loss'] = mag_loss.item()
losses['g_loss'] = g_loss.item()
customer_g_optimizer.zero_grad()
g_loss.backward()
nn.utils.clip_grad_norm_(g_parameters, max_norm=0.5)
customer_g_optimizer.step_and_update_lr()
time_used = time.time() - start
logging.info(
"Step: {} --sc_loss: {:.3f} --mag_loss: {:.3f} --Time: {:.2f} seconds"
.format(global_step, sc_loss, mag_loss, time_used))
if global_step % args.checkpoint_step == 0:
save_checkpoint(args, generator, discriminator, g_optimizer,
d_optimizer, global_step, logging)
if global_step % args.summary_step == 0:
writer.logging_loss(losses, global_step)
target = samples.cpu().detach()[0, 0].numpy()
predict = g_outputs.cpu().detach()[0, 0].numpy()
writer.logging_audio(target, predict, global_step)
writer.logging_histogram(generator, global_step)
writer.logging_histogram(discriminator, global_step)
global_step += 1
def train(args):
os.makedirs(args.checkpoint_dir, exist_ok=True)
train_dataset = CustomerDataset(
args.input,
upsample_factor=hop_length,
local_condition=True,
global_condition=False)
device = torch.device("cuda" if args.use_cuda else "cpu")
generator, discriminator = create_model(args)
print(generator)
print(discriminator)
num_gpu = torch.cuda.device_count() if args.use_cuda else 1
global_step = 0
g_parameters = list(generator.parameters())
g_optimizer = optim.Adam(g_parameters, lr=args.g_learning_rate)
d_parameters = list(discriminator.parameters())
d_optimizer = optim.Adam(d_parameters, lr=args.d_learning_rate)
writer = SummaryWriter(args.checkpoint_dir)
generator.to(device)
discriminator.to(device)
if args.resume is not None:
restore_step = attempt_to_restore(generator, discriminator, g_optimizer,
d_optimizer, args.resume, args.use_cuda)
global_step = restore_step
customer_g_optimizer = Optimizer(g_optimizer, args.g_learning_rate,
global_step, args.warmup_steps, args.decay_learning_rate)
customer_d_optimizer = Optimizer(d_optimizer, args.d_learning_rate,
global_step, args.warmup_steps, args.decay_learning_rate)
stft_criterion = MultiResolutionSTFTLoss().to(device)
criterion = nn.MSELoss().to(device)
for epoch in range(args.epochs):
collate = CustomerCollate(
upsample_factor=hop_length,
condition_window=args.condition_window,
local_condition=True,
global_condition=False)
train_data_loader = DataLoader(train_dataset, collate_fn=collate,
batch_size=args.batch_size, num_workers=args.num_workers,
shuffle=True, pin_memory=True)
#train one epoch
for batch, (samples, conditions) in enumerate(train_data_loader):
start = time.time()
batch_size = int(conditions.shape[0] // num_gpu * num_gpu)
samples = samples[:batch_size, :].to(device)
conditions = conditions[:batch_size, :, :].to(device)
z = torch.randn(batch_size, args.z_dim).to(device)
losses = {}
if num_gpu > 1:
g_outputs = parallel(generator, (conditions, z))
else:
g_outputs = generator(conditions, z)
#train discriminator
if global_step > args.discriminator_train_start_steps:
if num_gpu > 1:
real_outputs, fake_outputs = \
parallel(discriminator, (samples, g_outputs.detach(), conditions))
else:
real_outputs, fake_outputs = \
discriminator(samples, g_outputs.detach(), conditions)
fake_loss, real_loss = [], []
for (fake_output, real_output) in zip(fake_outputs, real_outputs):
fake_loss.append(criterion(fake_output, torch.zeros_like(fake_output)))
real_loss.append(criterion(real_output, torch.ones_like(real_output)))
#fake_loss = sum(fake_loss) / 10.0
#real_loss = sum(real_loss) / 10.0
fake_loss = sum(fake_loss)
real_loss = sum(real_loss)
d_loss = fake_loss + real_loss
customer_d_optimizer.zero_grad()
d_loss.backward()
nn.utils.clip_grad_norm_(d_parameters, max_norm=0.5)
customer_d_optimizer.step_and_update_lr()
else:
d_loss = torch.Tensor([0])
fake_loss = torch.Tensor([0])
real_loss = torch.Tensor([0])
losses['fake_loss'] = fake_loss.item()
losses['real_loss'] = real_loss.item()
losses['d_loss'] = d_loss.item()
#train generator
if num_gpu > 1:
_, fake_outputs = parallel(discriminator, (samples, g_outputs, conditions))
else:
_, fake_outputs = discriminator(samples, g_outputs, conditions)
adv_loss = []
for fake_output in fake_outputs:
adv_loss.append(criterion(fake_output, torch.ones_like(fake_output)))
#adv_loss = sum(adv_loss) / 10.0
adv_loss = sum(adv_loss)
sc_loss, mag_loss = stft_criterion(g_outputs.squeeze(1), samples.squeeze(1))
if global_step > args.discriminator_train_start_steps:
g_loss = adv_loss * args.lamda_adv + sc_loss + mag_loss
else:
g_loss = sc_loss + mag_loss
losses['adv_loss'] = adv_loss.item()
losses['sc_loss'] = sc_loss
losses['mag_loss'] = mag_loss
losses['g_loss'] = g_loss.item()
customer_g_optimizer.zero_grad()
g_loss.backward()
nn.utils.clip_grad_norm_(g_parameters, max_norm=0.5)
customer_g_optimizer.step_and_update_lr()
time_used = time.time() - start
if global_step > args.discriminator_train_start_steps:
print("Step: {} --adv_loss: {:.3f} --real_loss: {:.3f} --fake_loss: {:.3f} --sc_loss: {:.3f} --mag_loss: {:.3f} --Time: {:.2f} seconds".format(
global_step, adv_loss, real_loss, fake_loss, sc_loss, mag_loss, time_used))
else:
print("Step: {} --sc_loss: {:.3f} --mag_loss: {:.3f} --Time: {:.2f} seconds".format(global_step, sc_loss, mag_loss, time_used))
global_step += 1
if global_step % args.checkpoint_step == 0:
save_checkpoint(args, generator, discriminator,
g_optimizer, d_optimizer, global_step)
if global_step % args.summary_step == 0:
for key in losses:
writer.add_scalar('{}'.format(key), losses[key], global_step)
def train(args):
os.makedirs(args.checkpoint_dir, exist_ok=True)
os.makedirs(args.ema_checkpoint_dir, exist_ok=True)
train_dataset = CustomerDataset(args.input,
upsample_factor=hop_length,
local_condition=True,
global_condition=False)
device = torch.device("cuda" if args.use_cuda else "cpu")
generator, discriminator = create_model(args)
print(generator)
print(discriminator)
num_gpu = torch.cuda.device_count() if args.use_cuda else 1
global_step = 0
g_parameters = list(generator.parameters())
g_optimizer = optim.Adam(g_parameters, lr=args.g_learning_rate)
d_parameters = list(discriminator.parameters())
d_optimizer = optim.Adam(d_parameters, lr=args.d_learning_rate)
writer = SummaryWriter(args.checkpoint_dir)
generator.to(device)
discriminator.to(device)
if args.resume is not None:
restore_step = attempt_to_restore(generator, discriminator,
g_optimizer, d_optimizer,
args.resume, args.use_cuda)
global_step = restore_step
ema = ExponentialMovingAverage(args.ema_decay)
register_model_to_ema(generator, ema)
customer_g_optimizer = Optimizer(g_optimizer, args.g_learning_rate,
global_step, args.warmup_steps,
args.decay_learning_rate)
customer_d_optimizer = Optimizer(d_optimizer, args.d_learning_rate,
global_step, args.warmup_steps,
args.decay_learning_rate)
criterion = nn.MSELoss().to(device)
for epoch in range(args.epochs):
collate = CustomerCollate(upsample_factor=hop_length,
condition_window=args.condition_window,
local_condition=True,
global_condition=False)
train_data_loader = DataLoader(train_dataset,
collate_fn=collate,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=True,
pin_memory=True)
#train one epoch
for batch, (samples, conditions) in enumerate(train_data_loader):
start = time.time()
batch_size = int(conditions.shape[0] // num_gpu * num_gpu)
samples = samples[:batch_size, :].to(device)
conditions = conditions[:batch_size, :, :].to(device)
z = torch.randn(batch_size, args.z_dim).to(device)
#train generator
if num_gpu > 1:
g_outputs = parallel(generator, (conditions, z))
_, fake_outputs, real_features, fake_features = \
parallel(discriminator, (samples, g_outputs, conditions))
else:
g_outputs = generator(conditions, z)
_, fake_outputs, real_features, fake_features = \
discriminator(samples, g_outputs, conditions)
g_d_loss = []
for fake_output in fake_outputs:
target = torch.ones_like(fake_output).to(device)
g_d_loss.append(criterion(fake_output, target))
feature_loss = feature_loss_calculate(real_features, fake_features)
g_loss = feature_loss * args.lamda + sum(g_d_loss)
customer_g_optimizer.zero_grad()
g_loss.backward()
nn.utils.clip_grad_norm_(g_parameters, max_norm=0.5)
customer_g_optimizer.step_and_update_lr()
#train discriminator
g_outputs = g_outputs.detach()
if num_gpu > 1:
real_outputs, fake_outputs, _, _ = \
parallel(discriminator, (samples, g_outputs, conditions))
else:
real_outputs, fake_outputs, _, _ = \
discriminator(samples, g_outputs, conditions)
fake_loss, real_loss = [], []
for (fake_output, real_output) in zip(fake_outputs, real_outputs):
fake_target = torch.zeros_like(fake_output).to(device)
real_target = torch.ones_like(real_output).to(device)
fake_loss.append(criterion(fake_output, fake_target))
real_loss.append(criterion(real_output, real_target))
d_loss = sum(fake_loss) + sum(real_loss)
customer_d_optimizer.zero_grad()
d_loss.backward()
nn.utils.clip_grad_norm_(d_parameters, max_norm=0.5)
customer_d_optimizer.step_and_update_lr()
global_step += 1
print(
"Step: {} --g_loss: {:.3f} --d_loss: {:.3f} --Time: {:.2f} seconds"
.format(global_step, g_loss, d_loss,
float(time.time() - start)))
print(feature_loss.item(), sum(g_d_loss).item(), d_loss.item())
if ema is not None:
apply_moving_average(generator, ema)
if global_step % args.checkpoint_step == 0:
save_checkpoint(args, generator, discriminator, g_optimizer,
d_optimizer, global_step, ema)
if global_step % args.summary_step == 0:
writer.add_scalar("g_loss", g_loss.item(), global_step)
writer.add_scalar("d_loss", d_loss.item(), global_step)
def train(args):
os.makedirs(args.checkpoint_dir, exist_ok=True)
os.makedirs(args.ema_checkpoint_dir, exist_ok=True)
train_dataset = WaveRNNDataset(args.input,
upsample_factor=hop_length,
local_condition=True,
global_condition=False)
device = torch.device("cuda" if args.use_cuda else "cpu")
model = create_model(args)
print(model)
num_gpu = torch.cuda.device_count() if args.use_cuda else 1
model.train(mode=True)
global_step = 0
parameters = list(model.parameters())
optimizer = optim.Adam(parameters, lr=args.learning_rate)
writer = SummaryWriter(args.checkpoint_dir)
model.to(device)
if args.resume is not None:
restore_step = attempt_to_restore(model, optimizer, args.resume,
args.use_cuda)
global_step = restore_step
ema = ExponentialMovingAverage(args.ema_decay)
register_model_to_ema(model, ema)
customer_optimizer = Optimizer(optimizer, args.learning_rate, global_step,
args.warmup_steps, args.decay_learning_rate)
criterion = nn.NLLLoss().to(device)
for epoch in range(args.epochs):
collate = WaveRNNCollate(upsample_factor=hop_length,
condition_window=args.condition_window,
local_condition=True,
global_condition=False)
train_data_loader = DataLoader(train_dataset,
collate_fn=collate,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=True,
pin_memory=True)
#train one epoch
for batch, (coarse, fine, condition) in enumerate(train_data_loader):
start = time.time()
batch_size = int(condition.shape[0] // num_gpu * num_gpu)
coarse = coarse[:batch_size, :].to(device)
fine = fine[:batch_size, :].to(device)
condition = condition[:batch_size, :, :].to(device)
inputs = torch.cat([
coarse[:, :-1].unsqueeze(-1), fine[:, :-1].unsqueeze(-1),
coarse[:, 1:].unsqueeze(-1)
],
dim=-1)
inputs = 2 * inputs.float() / 255 - 1.0
if num_gpu > 1:
out_c, out_f, _ = parallel(model, (inputs, condition))
else:
out_c, out_f, _ = model(inputs, condition)
loss_c = criterion(out_c.transpose(1, 2).float(), coarse[:, 1:])
loss_f = criterion(out_f.transpose(1, 2).float(), fine[:, 1:])
loss = loss_c + loss_f
global_step += 1
customer_optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(parameters, max_norm=0.5)
customer_optimizer.step_and_update_lr()
model.after_update()
if ema is not None:
apply_moving_average(model, ema)
print(
"Step: {} --loss_c: {:.3f} --loss_f: {:.3f} --Lr: {:g} --Time: {:.2f} seconds"
.format(global_step, loss_c, loss_f, customer_optimizer.lr,
float(time.time() - start)))
if global_step % args.checkpoint_step == 0:
save_checkpoint(args, model, optimizer, global_step, ema)
if global_step % args.summary_step == 0:
writer.add_scalar("loss", loss.item(), global_step)
writer.add_scalar("loss_c", loss_c.item(), global_step)
writer.add_scalar("loss_f", loss_f.item(), global_step)