94 lines
2.2 KiB
Python
94 lines
2.2 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from utils.MultiResolutionSTFTLoss import MultiResolutionSTFTLoss
|
|
|
|
# Keep STFT settings as is
|
|
stft_loss_fn = MultiResolutionSTFTLoss(
|
|
fft_sizes=[512, 1024, 2048],
|
|
hop_sizes=[64, 128, 256],
|
|
win_lengths=[256, 512, 1024]
|
|
)
|
|
|
|
|
|
def feature_matching_loss(fmap_r, fmap_g):
|
|
"""
|
|
Computes L1 distance between real and fake feature maps.
|
|
"""
|
|
loss = 0
|
|
for dr, dg in zip(fmap_r, fmap_g):
|
|
for rl, gl in zip(dr, dg):
|
|
rl = rl.detach()
|
|
loss += torch.mean(torch.abs(rl - gl))
|
|
|
|
return loss * 2
|
|
|
|
|
|
def discriminator_loss(disc_real_outputs, disc_generated_outputs):
|
|
"""
|
|
Least Squares GAN Loss (LSGAN) for the Discriminator.
|
|
Objective: Real -> 1, Fake -> 0
|
|
"""
|
|
loss = 0
|
|
r_losses = []
|
|
g_losses = []
|
|
|
|
for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
|
|
r_loss = torch.mean((dr - 1) ** 2)
|
|
g_loss = torch.mean(dg ** 2)
|
|
|
|
loss += (r_loss + g_loss)
|
|
r_losses.append(r_loss.item())
|
|
g_losses.append(g_loss.item())
|
|
|
|
return loss, r_losses, g_losses
|
|
|
|
|
|
def generator_adv_loss(disc_generated_outputs):
|
|
"""
|
|
Least Squares GAN Loss for the Generator.
|
|
Objective: Fake -> 1 (Fool the discriminator)
|
|
"""
|
|
loss = 0
|
|
for dg in zip(disc_generated_outputs):
|
|
dg = dg[0] # Unpack tuple
|
|
loss += torch.mean((dg - 1) ** 2)
|
|
return loss
|
|
|
|
|
|
def discriminator_train(
|
|
high_quality,
|
|
discriminator,
|
|
generator_output
|
|
):
|
|
y_d_rs, y_d_gs, _, _ = discriminator(high_quality, generator_output.detach())
|
|
|
|
d_loss, _, _ = discriminator_loss(y_d_rs, y_d_gs)
|
|
|
|
return d_loss
|
|
|
|
|
|
def generator_train(
|
|
low_quality,
|
|
high_quality,
|
|
generator,
|
|
discriminator,
|
|
generator_output
|
|
):
|
|
y_d_rs, y_d_gs, fmap_rs, fmap_gs = discriminator(high_quality, generator_output)
|
|
|
|
loss_gen_adv = generator_adv_loss(y_d_gs)
|
|
|
|
loss_fm = feature_matching_loss(fmap_rs, fmap_gs)
|
|
|
|
stft_loss = stft_loss_fn(high_quality, generator_output)["total"]
|
|
|
|
lambda_stft = 45.0
|
|
lambda_fm = 2.0
|
|
lambda_adv = 1.0
|
|
|
|
combined_loss = (lambda_stft * stft_loss) + \
|
|
(lambda_fm * loss_fm) + \
|
|
(lambda_adv * loss_gen_adv)
|
|
|
|
return combined_loss, loss_gen_adv
|