fixed loss bugs

2024-09-06 00:21:13 -04:00
parent 4807a8caf8
commit 4ae5180c3a
4 changed files with 77 additions and 35 deletions
@@ -163,7 +163,7 @@ def generate_path(duration, mask):
    return path


-def clip_grad_value_(parameters, clip_value, norm_type=2):
+def clip_grad_value_(parameters, clip_value, norm_type=2, batch_size=1):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = list(filter(lambda p: p.grad is not None, parameters))
@@ -173,7 +173,7 @@ def clip_grad_value_(parameters, clip_value, norm_type=2):

    total_norm = 0
    for p in parameters:
-        param_norm = p.grad.data.norm(norm_type)
+        param_norm = p.grad.data.norm(norm_type)/batch_size
        total_norm += param_norm.item() ** norm_type
        if clip_value is not None:
            p.grad.data.clamp_(min=-clip_value, max=clip_value)
@@ -3,23 +3,63 @@ import torchaudio
 import torch.nn.functional as F
 from ..infer_pack.commons import median_pool1d

-def compute_temporal_envelope(log_magnitude: torch.Tensor, kernel_size=31):
+def log1p_autocorrelation(log_magnitude, max_lag=None):
+    """
+    Computes the autocorrelation of log-magnitude spectrogram.
+    :param log_magnitude: Log magnitude spectrogram (batch, channels, time)
+    :param max_lag: Maximum lag for which to compute autocorrelation.
+    :return: Autocorrelation of the log-magnitude spectrogram.
+    """
+    if max_lag is None:
+        max_lag = log_magnitude.size(-1) // 2  # Default to half the signal length

-    # Apply median pooling to the generated audio
-    pooled_stft = median_pool1d(log_magnitude, kernel_size, stride=kernel_size//2+1)
+    # Ensure the signal is zero-mean
+    log_magnitude = log_magnitude - log_magnitude.mean(dim=-1, keepdim=True)
    
-    # Compute differences on the median-pooled audio
-    diff = F.normalize(torch.diff(pooled_stft, dim=-1, n=1),dim=-1)
+    # Compute autocorrelation using FFT-based convolution
+    padded_log_magnitude = F.pad(log_magnitude, (0, max_lag))
+    # autocorr = torch.corrcoef(padded_log_magnitude)
+    autocorr = F.conv1d(padded_log_magnitude, log_magnitude.flip(dims=[-1]))
    
-    # Return the mean of the differences as the loss
-    return diff.sum(dim=-2).nan_to_num(0)
+    # Take only the relevant part and normalize
+    autocorr = F.normalize(autocorr[:, :max_lag + 1],p=2,dim=-1)
+    
+    return autocorr.nan_to_num(0)
+
+def compute_phase_loss(original_log_magnitude, generated_log_magnitude, max_lag=None):
+    """
+    Computes a time-shift invariant loss based on autocorrelation for log-magnitude data.
+    :param original_log_magnitude: Original log-magnitude spectrogram (batch, channels, time)
+    :param generated_log_magnitude: Generated log-magnitude spectrogram (batch, channels, time)
+    :param max_lag: Maximum lag for which to compute autocorrelation.
+    :return: Loss value.
+    """
+    # Compute autocorrelations
+    original_autocorr = log1p_autocorrelation(original_log_magnitude, max_lag)
+    generated_autocorr = log1p_autocorrelation(generated_log_magnitude, max_lag)
+
+    # Compute the loss between autocorrelations
+    loss = F.smooth_l1_loss(generated_autocorr, original_autocorr)
+    
+    return loss
+    
+def compute_temporal_envelope(log_magnitude: torch.Tensor, kernel_size=31):
+    """Compute the temporal envelope loss using median pooling and second-order differences."""
+    
+    # Apply median pooling to the log-magnitude spectrogram
+    pooled_stft = median_pool1d(log_magnitude, kernel_size, stride=kernel_size // 2 + 1)
+    
+    # Compute second-order differences on the median-pooled audio
+    diff = torch.diff(pooled_stft, dim=-1, n=2)
+    
+    return diff.nan_to_num(0)

 def compute_cepstrals(log_magnitude: torch.Tensor, kernel_size=31):

    pooled_magnitude = median_pool1d(log_magnitude, kernel_size, stride=kernel_size//2+1)

    # Compute the cepstrum
-    cepstral_coefficients = torch.fft.ifft(pooled_magnitude, dim=-2).real
+    cepstral_coefficients = torch.fft.ifft(pooled_magnitude, dim=-2).abs()
    
    return cepstral_coefficients.sum(dim=-2).nan_to_num(0)

@@ -28,25 +68,23 @@ def compute_sts_loss(gen_stft: torch.Tensor, org_stft: torch.Tensor, kernel_size
    # B?N?T
    org_mag = torch.log1p(org_stft.abs())
    gen_mag = torch.log1p(gen_stft.abs())
-
-    gen_temp = compute_temporal_envelope(gen_mag, kernel_size=kernel_size)
-    org_temp = compute_temporal_envelope(org_mag, kernel_size=kernel_size)
-    temporal_loss = 1-F.cosine_similarity(gen_temp, org_temp, dim=-1).mean()
-
-    gen_spec = compute_cepstrals(org_mag,kernel_size=kernel_size)
-    org_spec = compute_cepstrals(gen_mag,kernel_size=kernel_size)
-
-    spectral_loss = F.smooth_l1_loss(gen_spec, org_spec)
-
-    org_phase = org_stft.angle()
-    gen_phase = gen_stft.angle()
+    org_temp = torch.log1p(org_stft.abs().sum(-2,keepdim=True))
+    gen_temp = torch.log1p(gen_stft.abs().sum(-2,keepdim=True))

    # temporal invariant phase
-    org_phase = median_pool1d(org_stft.angle(),3,stride=3).mean(-1)
-    gen_phase = median_pool1d(gen_stft.angle(),3,stride=3).mean(-1)
-    phase_loss = F.smooth_l1_loss(gen_phase,org_phase)+F.smooth_l1_loss(1-torch.cos(gen_phase-org_phase),torch.zeros_like(gen_phase))
+    phase_loss = compute_phase_loss(org_temp,gen_temp)

-    return temporal_loss, spectral_loss, phase_loss.nan_to_num(0)
+    # temporal envelope
+    gen_temp = compute_temporal_envelope(gen_mag, kernel_size=kernel_size)
+    org_temp = compute_temporal_envelope(org_mag, kernel_size=kernel_size)
+    temporal_loss = F.smooth_l1_loss(gen_temp, org_temp)
+
+    # spectral
+    gen_spec = compute_cepstrals(gen_mag,kernel_size=kernel_size)
+    org_spec = compute_cepstrals(org_mag,kernel_size=kernel_size)
+    spectral_loss = F.smooth_l1_loss(gen_spec, org_spec)
+
+    return temporal_loss, spectral_loss, phase_loss

 def compute_harmonics(stft_matrix: torch.Tensor, kernel_size=31):
    # Calculate log-magnitude spectrograms
@@ -80,6 +118,7 @@ def combined_aux_loss(
            win_length=win_length,
            window=hann_window,
            return_complex=True,
+            onesided=True,
            center=False)
        original_stft = torch.stft(
            original_audio.view(-1,original_audio.size(-1)),
@@ -88,6 +127,7 @@ def combined_aux_loss(
            win_length=win_length,
            window=hann_window,
            return_complex=True,
+            onesided=True,
            center=False)
    
    # Harmonic Loss
@@ -112,7 +152,7 @@ def combined_aux_loss(
    # MFCC Loss
    if c_mfcc>0:
        mfcc_transform = torchaudio.transforms.MFCC(
-            sample_rate=sample_rate, n_mfcc=n_mfcc, 
+            sample_rate=sample_rate, n_mfcc=n_mfcc, log_mels=True,
            melkwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length, n_mels=n_filter, f_min=f_min, f_max=f_max)
        ).to(original_audio.device)
        original_mfcc = mfcc_transform(original_audio)
@@ -123,7 +163,7 @@ def combined_aux_loss(
    # LFCC Loss
    if c_lfcc>0:
        lfcc_transform = torchaudio.transforms.LFCC(
-            sample_rate=sample_rate, n_lfcc=n_lfcc, n_filter=n_filter, f_min=f_min, f_max=f_max,
+            sample_rate=sample_rate, n_lfcc=n_lfcc, n_filter=n_filter, f_min=f_min, f_max=f_max,log_lf=True,
            speckwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length)
        ).to(original_audio.device)
        original_lfcc = lfcc_transform(original_audio)
@@ -74,7 +74,7 @@ def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False)
    y = y.squeeze(1)

    # Complex Spectrogram :: (B, T) -> (B, Freq, Frame, RealComplex=2)
-    spec = torch.stft(
+    spec = torch.view_as_real(torch.stft(
        y,
        n_fft,
        hop_length=hop_size,
@@ -84,8 +84,8 @@ def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False)
        pad_mode="reflect",
        normalized=False,
        onesided=True,
-        return_complex=False,
-    )
+        return_complex=True,
+    ))

    # Linear-frequency Linear-amplitude spectrogram :: (B, Freq, Frame, RealComplex=2) -> (B, Freq, Frame)
    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
@@ -3,6 +3,7 @@ import json
 import os
 import shutil
 import traceback
+import numpy as np

 from tqdm import tqdm
 from .lib.audio import SR_MAP
@@ -120,9 +121,10 @@ def train_model(hps: "utils.HParams"):

 def run(rank, n_gpus, hps, device):
    print(f"{__name__=}")
-    global global_step, least_loss, loss_file, best_model_name
+    global global_step, least_loss, loss_file, best_model_name, gradient_clip_value
    global_step = 0
    loss_file = os.path.join(hps.model_dir,"losses.json")
+    gradient_clip_value = np.log(hps.train.c_gp/hps.train.learning_rate) if hps.train.get("c_gp",0.)>0 else None

    if os.path.isfile(loss_file):
        with open(loss_file,"r") as f:
@@ -314,7 +316,7 @@ def train_and_evaluate(
        writer, _ = writers

    train_loader.batch_sampler.set_epoch(epoch)
-    global global_step, least_loss, loss_file, best_model_name
+    global global_step, least_loss, loss_file, best_model_name, gradient_clip_value

    net_g.train()
    net_d.train()
@@ -488,7 +490,7 @@ def train_and_evaluate(
            optim_d.zero_grad()
            scaler.scale(loss_disc_all).backward()
            scaler.unscale_(optim_d)
-            grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
+            grad_norm_d = commons.clip_grad_value_(net_d.parameters(), gradient_clip_value, batch_size=hps.train.batch_size)
            scaler.step(optim_d)

        with autocast(enabled=hps.train.fp16_run):
@@ -518,7 +520,7 @@ def train_and_evaluate(
            optim_g.zero_grad()
            scaler.scale(loss_gen_all).backward()
            scaler.unscale_(optim_g)
-            grad_norm_g = commons.clip_grad_value_(net_g.parameters(), None)
+            grad_norm_g = commons.clip_grad_value_(net_g.parameters(), gradient_clip_value, batch_size=hps.train.batch_size)
            scaler.step(optim_g)
            scaler.update()