def test_melgan_trainable_with_melgan_discriminator(dict_g, dict_d, dict_loss):
# setup
batch_size = 4
batch_length = 4096
args_g = make_melgan_generator_args(**dict_g)
args_d = make_melgan_discriminator_args(**dict_d)
args_loss = make_mutli_reso_stft_loss_args(**dict_loss)
y = torch.randn(batch_size, 1, batch_length)
c = torch.randn(batch_size, args_g["in_channels"],
batch_length // np.prod(
args_g["upsample_scales"]))
model_g = MelGANGenerator(**args_g)
model_d = MelGANMultiScaleDiscriminator(**args_d)
aux_criterion = MultiResolutionSTFTLoss(**args_loss)
optimizer_g = RAdam(model_g.parameters())
optimizer_d = RAdam(model_d.parameters())
# check generator trainable
y_hat = model_g(c)
p_hat = model_d(y_hat)
y, y_hat = y.squeeze(1), y_hat.squeeze(1)
sc_loss, mag_loss = aux_criterion(y_hat, y)
aux_loss = sc_loss + mag_loss
adv_loss = 0.0
for i in range(len(p_hat)):
adv_loss += F.mse_loss(
p_hat[i][-1], p_hat[i][-1].new_ones(p_hat[i][-1].size()))
adv_loss /= (i + 1)
with torch.no_grad():
p = model_d(y.unsqueeze(1))
fm_loss = 0.0
for i in range(len(p_hat)):
for j in range(len(p_hat[i]) - 1):
fm_loss += F.l1_loss(p_hat[i][j], p[i][j].detach())
fm_loss /= (i + 1) * j
loss_g = adv_loss + aux_loss + fm_loss
optimizer_g.zero_grad()
loss_g.backward()
optimizer_g.step()
# check discriminator trainable
y, y_hat = y.unsqueeze(1), y_hat.unsqueeze(1).detach()
p = model_d(y)
p_hat = model_d(y_hat)
real_loss = 0.0
fake_loss = 0.0
for i in range(len(p)):
real_loss += F.mse_loss(
p[i][-1], p[i][-1].new_ones(p[i][-1].size()))
fake_loss += F.mse_loss(
p_hat[i][-1], p_hat[i][-1].new_zeros(p_hat[i][-1].size()))
real_loss /= (i + 1)
fake_loss /= (i + 1)
loss_d = real_loss + fake_loss
optimizer_d.zero_grad()
loss_d.backward()
optimizer_d.step()
def test_hifigan_trainable(dict_g, dict_d, dict_loss):
# setup
batch_size = 4
batch_length = 2**13
args_g = make_hifigan_generator_args(**dict_g)
args_d = make_hifigan_multi_scale_multi_period_discriminator_args(**dict_d)
args_loss = make_mutli_reso_stft_loss_args(**dict_loss)
y = torch.randn(batch_size, 1, batch_length)
c = torch.randn(
batch_size,
args_g["in_channels"],
batch_length // np.prod(args_g["upsample_scales"]),
)
model_g = HiFiGANGenerator(**args_g)
model_d = HiFiGANMultiScaleMultiPeriodDiscriminator(**args_d)
aux_criterion = MultiResolutionSTFTLoss(**args_loss)
feat_match_criterion = FeatureMatchLoss(
average_by_layers=False,
average_by_discriminators=False,
include_final_outputs=True,
)
gen_adv_criterion = GeneratorAdversarialLoss(
average_by_discriminators=False, )
dis_adv_criterion = DiscriminatorAdversarialLoss(
average_by_discriminators=False, )
optimizer_g = torch.optim.AdamW(model_g.parameters())
optimizer_d = torch.optim.AdamW(model_d.parameters())
# check generator trainable
y_hat = model_g(c)
p_hat = model_d(y_hat)
sc_loss, mag_loss = aux_criterion(y_hat, y)
aux_loss = sc_loss + mag_loss
adv_loss = gen_adv_criterion(p_hat)
with torch.no_grad():
p = model_d(y)
fm_loss = feat_match_criterion(p_hat, p)
loss_g = adv_loss + aux_loss + fm_loss
optimizer_g.zero_grad()
loss_g.backward()
optimizer_g.step()
# check discriminator trainable
p = model_d(y)
p_hat = model_d(y_hat.detach())
real_loss, fake_loss = dis_adv_criterion(p_hat, p)
loss_d = real_loss + fake_loss
optimizer_d.zero_grad()
loss_d.backward()
optimizer_d.step()
print(model_d)
print(model_g)
def test_melgan_trainable_with_melgan_discriminator(dict_g, dict_d, dict_loss):
# setup
batch_size = 4
batch_length = 4096
args_g = make_melgan_generator_args(**dict_g)
args_d = make_melgan_discriminator_args(**dict_d)
args_loss = make_mutli_reso_stft_loss_args(**dict_loss)
y = torch.randn(batch_size, 1, batch_length)
c = torch.randn(
batch_size,
args_g["in_channels"],
batch_length // np.prod(args_g["upsample_scales"]),
)
model_g = MelGANGenerator(**args_g)
model_d = MelGANMultiScaleDiscriminator(**args_d)
aux_criterion = MultiResolutionSTFTLoss(**args_loss)
feat_match_criterion = FeatureMatchLoss()
gen_adv_criterion = GeneratorAdversarialLoss()
dis_adv_criterion = DiscriminatorAdversarialLoss()
optimizer_g = RAdam(model_g.parameters())
optimizer_d = RAdam(model_d.parameters())
# check generator trainable
y_hat = model_g(c)
p_hat = model_d(y_hat)
sc_loss, mag_loss = aux_criterion(y_hat, y)
aux_loss = sc_loss + mag_loss
adv_loss = gen_adv_criterion(p_hat)
with torch.no_grad():
p = model_d(y)
fm_loss = feat_match_criterion(p_hat, p)
loss_g = adv_loss + aux_loss + fm_loss
optimizer_g.zero_grad()
loss_g.backward()
optimizer_g.step()
# check discriminator trainable
p = model_d(y)
p_hat = model_d(y_hat.detach())
real_loss, fake_loss = dis_adv_criterion(p_hat, p)
loss_d = real_loss + fake_loss
optimizer_d.zero_grad()
loss_d.backward()
optimizer_d.step()
def test_style_melgan_trainable(dict_g, dict_d, dict_loss, loss_type):
# setup
args_g = make_style_melgan_generator_args(**dict_g)
args_d = make_style_melgan_discriminator_args(**dict_d)
args_loss = make_mutli_reso_stft_loss_args(**dict_loss)
batch_size = 4
batch_length = np.prod(args_g["noise_upsample_scales"]) * np.prod(
args_g["upsample_scales"]
)
y = torch.randn(batch_size, 1, batch_length)
c = torch.randn(
batch_size,
args_g["aux_channels"],
batch_length // np.prod(args_g["upsample_scales"]),
)
model_g = StyleMelGANGenerator(**args_g)
model_d = StyleMelGANDiscriminator(**args_d)
aux_criterion = MultiResolutionSTFTLoss(**args_loss)
gen_adv_criterion = GeneratorAdversarialLoss(loss_type=loss_type)
dis_adv_criterion = DiscriminatorAdversarialLoss(loss_type=loss_type)
optimizer_g = torch.optim.Adam(model_g.parameters())
optimizer_d = torch.optim.Adam(model_d.parameters())
# check generator trainable
y_hat = model_g(c)
p_hat = model_d(y_hat)
adv_loss = gen_adv_criterion(p_hat)
sc_loss, mag_loss = aux_criterion(y_hat, y)
aux_loss = sc_loss + mag_loss
loss_g = adv_loss + aux_loss
optimizer_g.zero_grad()
loss_g.backward()
optimizer_g.step()
# check discriminator trainable
p = model_d(y)
p_hat = model_d(y_hat.detach())
real_loss, fake_loss = dis_adv_criterion(p_hat, p)
loss_d = real_loss + fake_loss
optimizer_d.zero_grad()
loss_d.backward()
optimizer_d.step()
def test_melgan_trainable_with_residual_discriminator(dict_g, dict_d,
dict_loss):
# setup
batch_size = 4
batch_length = 4096
args_g = make_generator_args(**dict_g)
args_d = make_residual_discriminator_args(**dict_d)
args_loss = make_mutli_reso_stft_loss_args(**dict_loss)
y = torch.randn(batch_size, 1, batch_length)
c = torch.randn(batch_size, args_g["in_channels"],
batch_length // np.prod(args_g["upsample_scales"]))
model_g = MelGANGenerator(**args_g)
model_d = ResidualParallelWaveGANDiscriminator(**args_d)
aux_criterion = MultiResolutionSTFTLoss(**args_loss)
optimizer_g = RAdam(model_g.parameters())
optimizer_d = RAdam(model_d.parameters())
# check generator trainable
y_hat = model_g(c)
p_hat = model_d(y_hat)
y, y_hat, p_hat = y.squeeze(1), y_hat.squeeze(1), p_hat.squeeze(1)
adv_loss = F.mse_loss(p_hat, p_hat.new_ones(p_hat.size()))
sc_loss, mag_loss = aux_criterion(y_hat, y)
aux_loss = sc_loss + mag_loss
loss_g = adv_loss + aux_loss
optimizer_g.zero_grad()
loss_g.backward()
optimizer_g.step()
# check discriminator trainable
y, y_hat = y.unsqueeze(1), y_hat.unsqueeze(1).detach()
p = model_d(y)
p_hat = model_d(y_hat)
p, p_hat = p.squeeze(1), p_hat.squeeze(1)
loss_d = F.mse_loss(p, p.new_ones(p.size())) + F.mse_loss(
p_hat, p_hat.new_zeros(p_hat.size()))
optimizer_d.zero_grad()
loss_d.backward()
optimizer_d.step()
