From 4ae5180c3a8fa83e0f59cc7f62c3944cd5327ba9 Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Fri, 6 Sep 2024 00:21:13 -0400 Subject: [PATCH] fixed loss bugs --- lib/infer_pack/commons.py | 4 +- lib/train/losses.py | 92 ++++++++++++++++++++++++++----------- lib/train/mel_processing.py | 6 +-- training_cli.py | 10 ++-- 4 files changed, 77 insertions(+), 35 deletions(-) diff --git a/lib/infer_pack/commons.py b/lib/infer_pack/commons.py index 842339f..2d100ef 100644 --- a/lib/infer_pack/commons.py +++ b/lib/infer_pack/commons.py @@ -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) diff --git a/lib/train/losses.py b/lib/train/losses.py index ff0967b..3ab7bf7 100644 --- a/lib/train/losses.py +++ b/lib/train/losses.py @@ -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) diff --git a/lib/train/mel_processing.py b/lib/train/mel_processing.py index 1c871ab..837239f 100644 --- a/lib/train/mel_processing.py +++ b/lib/train/mel_processing.py @@ -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) diff --git a/training_cli.py b/training_cli.py index ed611d0..d962ec4 100644 --- a/training_cli.py +++ b/training_cli.py @@ -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()