nodes/SayanoAI-Comfy-RVC
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

improved balancer + loss logic

Browse Source
This commit is contained in:
SayanoAI
2024-10-03 19:20:55 -04:00
parent 427fb7521d
commit 784767c8e0
4 changed files with 120 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_nodes/rvc_nodes.py
+2 -2
View File
@@ -419,7 +419,7 @@ class RVCTrainModelNode:
train_index=("BOOLEAN",{"default": True}),
retrain=("BOOLEAN",{"default": False}),
save_best_model=("BOOLEAN",{"default": True}),
best_model_threshold=("INT",dict(default=30,min=10,max=50)),
best_model_threshold=("INT",dict(default=10,min=1,max=100)),
log_every_epoch=("FLOAT",dict(default=1.,min=0.,max=2.,step=.1)),
num_workers=("INT",dict(default=1,min=1,max=16))
)
@@ -448,7 +448,7 @@ class RVCTrainModelNode:
train_index=True,
retrain=False,
save_best_model=True,
best_model_threshold=50,
best_model_threshold=10,
log_every_epoch=1.,
num_workers=1):
lib/infer_pack/commons.py
+5
View File
@@ -71,6 +71,11 @@ def minmax_scale(tensor: torch.Tensor, eps=1e-8):
tensor = (tensor - tensor.min()) / (tensor.max() - tensor.min() + eps)
return tensor
def compute_correlation(t1: torch.Tensor, t2: torch.Tensor, eps=1e-8):
numerator = (t1 * t2).nan_to_num(0).sum(dim=-1)
denominator = torch.sqrt((t1 ** 2).sum(dim=-1) * (t2 ** 2).sum(dim=-1)+eps).nan_to_num(0)
return (numerator / denominator).nan_to_num(0)
def init_weights(m, mean=0.0, std=0.01):
classname = m.__class__.__name__
if classname.find("Conv") != -1:
lib/train/losses.py
+57 -47
View File
@@ -1,14 +1,15 @@
from typing import Optional
import torch
import torchaudio.transforms as T
import torch.nn.functional as F
from ..utils import gc_collect
from ..infer_pack.commons import median_pool1d, minmax_scale
from ..infer_pack.commons import compute_correlation, median_pool1d, minmax_scale
class LossBalancer:
model: torch.nn.Module
def __init__(self, model, initial_weights={}, historical_losses={}, ema_weights={}, epsilon=1e-8, weights_decay=0., loss_decay=.0, active=True, use_pareto=True, use_norm=False):
def __init__(self, model: torch.nn.Module, initial_weights: dict={}, historical_losses: dict={}, ema_weights: dict={}, epsilon=1e-8, weights_decay=0., loss_decay=.0, active=True, use_pareto=True, use_norm=False):
"""
Initializes the LossBalancer with optional initial weights, partial gradient computation, and EMA.
@@ -29,7 +30,6 @@ class LossBalancer:
self.use_pareto = use_pareto # apply the 20/80 rule to loss balancing
self.use_norm = use_norm
def to_dict(self):
return dict(
epsilon = self.epsilon,
@@ -45,38 +45,41 @@ class LossBalancer:
def update_ema_weights(self, new_weights):
"""Updates the EMA of the weights."""
if not self.ema_weights: self.ema_weights = new_weights
if not self.ema_weights: self.ema_weights = dict(**new_weights)
else: self.ema_weights = {
k: self.weights_decay * self.ema_weights.get(k, 0.) + (1 - self.weights_decay) * new_weights[k]
k: self.weights_decay * self.ema_weights.get(k, 1.) + (1 - self.weights_decay) * new_weights[k]
for k in new_weights
}
return dict(**self.ema_weights)
def update_historical_losses(self, new_losses: dict):
"""Updates the EMA of the weights."""
if not self.historical_losses: self.historical_losses = new_losses
if not self.historical_losses: self.historical_losses = dict(**new_losses)
else:
for k in new_losses:
self.historical_losses[k]=self.loss_decay*self.historical_losses.get(k, 0.) + (1 - self.loss_decay)*new_losses[k]
return dict(**self.historical_losses)
def calculate_loss_slope(self, key, current_loss):
"""Calculates the slope of the loss using the current loss and its EMA."""
ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon
slope = abs(current_loss-ema_loss)/ema_loss # relative loss change
slope = (current_loss-ema_loss)/ema_loss # relative loss change
return slope
def calculate_gradients(self, key, current_loss):
def calculate_gradients(self, key, current_loss, input):
"""Calculates the gradient norm of the current loss wrt the model params."""
self.model.zero_grad() # Clear previous gradients
# Backward pass to output layer
model_params = list(self.model.parameters())[-1]
output_params = torch.autograd.grad(current_loss, model_params, grad_outputs=torch.ones_like(current_loss, device=model_params.device), retain_graph=True, only_inputs=True)[0]
output_params = torch.autograd.grad(current_loss, [input], retain_graph=True, allow_unused=True, materialize_grads=True)[0]
# Compute L2 gradient norm
grad_norm = output_params.view(output_params.size(0),-1).norm(2, dim=-1).mean()
ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon
return grad_norm/ema_loss
# grad_norm = output_params.view(output_params.size(0),-1).norm(2, dim=-1).mean()
if output_params.ndim>1: grad_norm = output_params.view(output_params.size(0), -1).norm(2, dim=-1).mean()
else: grad_norm = output_params.norm(2)
slope = self.calculate_loss_slope(key, current_loss)
return grad_norm * slope
def pareto_normalizer(self, data: dict):
# Sort the data in descending order based on the values
@@ -104,11 +107,11 @@ class LossBalancer:
normalized_weights = [w / total_weight for w in weights]
# Combine the keys with their normalized weights
normalized_data = {sorted_data[i][0]: normalized_weights[i] for i in range(len(sorted_data))}
normalized_data = {sorted_data[i][0]: normalized_weights[i]*len(data) for i in range(len(sorted_data))}
return normalized_data
def on_train_batch_start(self, losses: dict):
def on_train_batch_start(self, losses: dict, input: Optional[torch.Tensor]=None):
"""
Balances the loss terms based on their gradients before each training batch.
@@ -120,26 +123,32 @@ class LossBalancer:
new_weights (dict): Updated weights for each loss term.
"""
if not self.initial_weights: self.initial_weights = {k: 1.0 for k in losses} # initialize weights
if not self.active: return sum(v*self.initial_weights[k] if k in self.initial_weights else v for k,v in losses.items())
if len(losses)==0: return 0. # no losses to balance
if not self.initial_weights: self.initial_weights = {k: 1. for k in losses} # initialize weights
if not self.ema_weights: self.ema_weights = {k: 1. for k in losses} # initialize ema weights
if not self.active:
self.update_historical_losses({k: v.item()*self.initial_weights.get(k,1.) for k,v in losses.items() if v>0})
return sum(v*self.initial_weights.get(k,1.) for k,v in losses.items())
gradients = {}
valid_losses = {}
# Process each loss, skip if the corresponding weight is 0
for key, loss in losses.items():
weight = self.initial_weights[key] if key in self.initial_weights else 1.
weight = self.initial_weights.get(key, 1.)
if weight == 0 or loss == 0 or not isinstance(loss,torch.Tensor): continue # Skip loss with weight 0 or constants
weighted_loss = loss * weight
if self.use_norm:
if self.use_norm and input is not None:
# Compute L2 gradient norm
grad_norm = self.calculate_gradients(key, loss)
gradients[key] = grad_norm.item()
input.requires_grad_(True)
grad_norm = self.calculate_gradients(key, weighted_loss, input)
gradients[key] = max(grad_norm.item(), self.epsilon)
else:
# Use loss plateau detection with EMA
loss_slope = self.calculate_loss_slope(key, loss)
gradients[key] = loss_slope.item() # Smaller slope -> plateau -> lower priority
valid_losses[key] = loss
loss_slope = self.calculate_loss_slope(key, weighted_loss)
gradients[key] = max(loss_slope.item(), self.epsilon)
valid_losses[key] = weighted_loss
if not valid_losses: return torch.tensor(0.0) # If all losses are skipped
@@ -148,7 +157,7 @@ class LossBalancer:
normalized_weights = self.pareto_normalizer(gradients)
else:
total_gradient = sum(gradients.values()) + self.epsilon
normalized_weights = {k: w/total_gradient for k, w in gradients.items()}
normalized_weights = {k: w/total_gradient*len(gradients) for k, w in gradients.items()}
# Update EMA weights
self.update_ema_weights(normalized_weights)
@@ -167,12 +176,16 @@ class LossBalancer:
"""
if weights_decay is not None: self.weights_decay=weights_decay
if loss_decay is not None: self.loss_decay=loss_decay
print(f"|| EMA loss weights: {self.ema_weights} ({self.weights_decay=})")
print(f"|| EMA losses: {self.historical_losses} ({self.loss_decay=})")
weighted_loss = sum(v*self.ema_weights.get(k,0.) for k,v in self.historical_losses.items())
print(f"===> EMA loss: {weighted_loss:.3f} <====")
weights = dict(sorted(self.ema_weights.items(),key=lambda x:x[1],reverse=True))
losses = dict(sorted(self.historical_losses.items(),key=lambda x:x[1],reverse=True))
print(f"|| EMA weights: {weights} ({self.weights_decay=})")
print(f"|| EMA losses: {losses} ({self.loss_decay=})")
print(f"===> Weighted EMA loss: {self.weighted_ema_loss:.3f} <====")
gc_collect()
@property
def weighted_ema_loss(self):
return sum(v*self.ema_weights.get(k,1./len(self.historical_losses)) for k,v in self.historical_losses.items())
def compute_tsi_loss(original_log_magnitude: torch.Tensor, generated_log_magnitude: torch.Tensor, dim=-1, eps=1e-8):
"""
@@ -189,14 +202,12 @@ def compute_tsi_loss(original_log_magnitude: torch.Tensor, generated_log_magnitu
generated_envelope = minmax_scale(generated_envelope, eps=eps)
# Compute the correlation
numerator = (original_envelope * generated_envelope).sum(dim=-1)
denominator = torch.sqrt((original_envelope ** 2).sum(dim=-1) * (generated_envelope ** 2).sum(dim=-1)+eps)
correlation = numerator / denominator
correlation = compute_correlation(original_envelope, generated_envelope, eps=eps)
# Compute the loss as negative correlation
loss = 1-correlation.mean()
loss = 1-correlation.abs()
return loss
return loss.mean()
def compute_envelope(log_magnitude: torch.Tensor, dim=-1, kernel_size=3, eps=1e-8):
"""
@@ -261,9 +272,9 @@ def compute_tefs(audio_signal: torch.Tensor, eps=1e-8):
envelope = minmax_scale(envelope, eps=eps)
# Compute the instantaneous phase from the analytic signal
phase = torch.angle(analytic_signal)
phase = torch.angle(analytic_signal).cos()
return envelope, phase
return envelope.nan_to_num(0), phase.nan_to_num(0)
def compute_harmonics(mag: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percussive_kernel_sizes=[3,7,13], eps=1e-8):
# Calculate log-magnitude spectrograms
@@ -285,7 +296,7 @@ def compute_harmonics(mag: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percu
harmonic = minmax_scale(harmonic, eps=eps)
percussive = minmax_scale(percussive, eps=eps)
return harmonic, percussive
return harmonic.nan_to_num(0), percussive.nan_to_num(0)
def combined_aux_loss(
original_audio: torch.Tensor, generated_audio: torch.Tensor,
@@ -333,14 +344,13 @@ def combined_aux_loss(
# Define loss terms
harmonic_loss = F.l1_loss(generated_harmonics, original_harmonics)
harmonic_loss += F.l1_loss(generated_percussives, original_percussives)
harmonic_loss *= c_hd
else: harmonic_loss = 0
# temporal invariant phase
if c_tsi>0:
freq_tsi = compute_tsi_loss(org_mag,gen_mag,dim=-1, eps=eps)
temp_tsi = compute_tsi_loss(org_mag,gen_mag,dim=-2, eps=eps)
tsi_loss = (freq_tsi+temp_tsi) * c_tsi
tsi_loss = (freq_tsi+temp_tsi)
else: tsi_loss = 0
# Temperol Envelope and Fine Structure Loss
@@ -348,10 +358,9 @@ def combined_aux_loss(
# temporal envelope
gen_te, gen_tfs = compute_tefs(generated_audio, eps=eps)
org_te, org_tfs = compute_tefs(original_audio, eps=eps)
phase_diff = torch.remainder(gen_tfs - org_tfs, 2 * torch.pi)
correlation = 1-compute_correlation(gen_tfs,org_tfs, eps=eps).abs()
tefs_loss = F.l1_loss(gen_te, org_te) + phase_diff.sin().abs().mean() #F.l1_loss(gen_tfs, org_tfs)
tefs_loss*=c_tefs
tefs_loss = F.l1_loss(gen_te, org_te) + correlation.mean()
else: tefs_loss = 0
return harmonic_loss, tefs_loss, tsi_loss
@@ -375,10 +384,12 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens
inputs=interpolated,
grad_outputs=torch.ones_like(output, device=original_audio.device),
retain_graph=True,
only_inputs=True
allow_unused=True,
materialize_grads=True
)[0]
grad_norm = gradients.view(gradients.size(0), -1).norm(2, dim=-1)+eps
loss += torch.log1p((grad_norm - 1) ** 2).mean()
if gradients.ndim>1: grad_norm = gradients.view(gradients.size(0), -1).norm(2, dim=-1).mean()
else: grad_norm = gradients.norm(2)
loss += torch.log1p((grad_norm - 1) ** 2)
return loss/len(disc_interpolated_output)
def feature_loss(fmap_r, fmap_g):
@@ -391,7 +402,6 @@ def feature_loss(fmap_r, fmap_g):
return loss
def discriminator_loss(disc_real_outputs, disc_generated_outputs):
loss = 0
r_losses = []
training_cli.py
+56 -48
View File
@@ -295,8 +295,16 @@ def run(rank, n_gpus, hps, device):
epsilon=hps.train.eps,
active=hps.train.get("use_balancer",False),
use_pareto=hps.train.get("use_pareto",False),
use_norm=not hps.train.get("fast_mode",False)
)
use_norm=not hps.train.get("fast_mode",False),
initial_weights=dict(
loss_gen=1.,
loss_fm=hps.train.get("c_fm",2.),
loss_mel=hps.train.get("c_mel",45.),
loss_kl=hps.train.get("c_kl",1.),
harmonic_loss=hps.train.get("c_hd",0.),
tsi_loss=hps.train.get("c_tsi",0.),
tefs_loss=hps.train.get("c_tefs",0.),
))
cache = []
for epoch in range(epoch_str, hps.train.epochs + 1):
train_loader.batch_sampler.set_epoch(epoch)
@@ -524,14 +532,14 @@ def train_and_evaluate(
# Generator
y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(wave, y_hat)
with autocast(enabled=False):
loss_mel = F.l1_loss(y_mel, y_hat_mel)*hps.train.get("c_mel",0.)
loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask)*hps.train.get("c_kl",0.)
loss_fm = feature_loss(fmap_r, fmap_g)*hps.train.get("c_fm",0.)
loss_mel = F.l1_loss(y_mel, y_hat_mel) #*hps.train.get("c_mel",0.)
loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) #*hps.train.get("c_kl",0.)
loss_fm = feature_loss(fmap_r, fmap_g) #*hps.train.get("c_fm",0.)
harmonic_loss, tefs_loss, tsi_loss = combined_aux_loss(
wave, y_hat,n_mels=hps.data.n_mel_channels,sample_rate=hps.data.sampling_rate,
c_tefs=hps.train.c_tefs if hps.train.get("c_tefs",0.)>0 else 0.,
c_hd=hps.train.c_hd if hps.train.get("c_hd",0.)>0 else 0.,
c_tsi=hps.train.c_tsi if hps.train.get("c_tsi",0.)>0 else 0.,
c_tefs=hps.train.get("c_tefs",0.),
c_hd=hps.train.get("c_hd",0.),
c_tsi=hps.train.get("c_tsi",0.),
n_fft=hps.data.filter_length,
hop_length=hps.data.hop_length,
win_length=hps.data.win_length,
@@ -539,7 +547,6 @@ def train_and_evaluate(
)
loss_gen, losses_gen = generator_loss(y_d_hat_g)
aux_loss = harmonic_loss + tefs_loss + tsi_loss
# loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl + aux_loss
loss_gen_all = balancer.on_train_batch_start(dict(
loss_gen=loss_gen,
loss_fm=loss_fm,
@@ -548,7 +555,7 @@ def train_and_evaluate(
harmonic_loss=harmonic_loss,
tsi_loss=tsi_loss,
tefs_loss=tefs_loss,
))
),input=y_hat)
optim_g.zero_grad()
scaler.scale(loss_gen_all).backward()
@@ -558,44 +565,10 @@ def train_and_evaluate(
scaler.update()
if rank == 0:
# if epoch>=max(hps.total_epoch*.2,20): #start saving best model after 20% or 20 epochs
if loss_gen + loss_fm + loss_mel + loss_kl + aux_loss<least_loss:
least_loss = loss_gen + loss_fm + loss_mel + loss_kl + aux_loss
logger.info(f"[lowest loss] {least_loss=:.3f}: {loss_disc=:.3f} {gradient_penalty=:.3f} <==> {loss_gen=:.3f} {loss_fm=:.3f} {loss_mel=:.3f} {loss_kl=:.3f}")
logger.info(f"[aux] {aux_loss=:.3f}: {harmonic_loss=:.3f}, {tefs_loss=:.3f}, {tsi_loss=:.3f}")
if hps.save_best_model:
if hasattr(net_g, "module"): ckpt = net_g.module.state_dict()
else: ckpt = net_g.state_dict()
best_model_name = f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:.0f}" if hps.save_every_weights else f"{hps.name}_loss{least_loss:2.0f}"
status = save_checkpoint(ckpt,best_model_name,epoch,hps)
logger.info(f"=== saving best model {best_model_name}: {status=} ===")
with open(loss_file,"w") as f:
json.dump(dict(least_loss=least_loss.item(),best_model_name=best_model_name,epoch=epoch,steps=global_step,
loss_weights = dict(balancer.ema_weights),
scalar_dict={
"total/loss/all": commons.serialize_tensor(loss_gen_all+loss_disc_all),
"total/loss/gen_all": commons.serialize_tensor(loss_gen_all),
"total/loss/aux": commons.serialize_tensor(aux_loss),
"total/loss/disc_all": commons.serialize_tensor(loss_disc_all),
"total/loss/gen": commons.serialize_tensor(loss_gen),
"total/loss/disc": commons.serialize_tensor(loss_disc),
"total/loss/fm": commons.serialize_tensor(loss_fm),
"total/loss/mel": commons.serialize_tensor(loss_mel),
"total/loss/kl": commons.serialize_tensor(loss_kl),
"aux/loss/harmonic": commons.serialize_tensor(harmonic_loss),
"aux/loss/tefs": commons.serialize_tensor(tefs_loss),
"aux/loss/tsi": commons.serialize_tensor(tsi_loss),
"gradient/grad_norm_disc": commons.serialize_tensor(grad_norm_d),
"gradient/grad_norm_gen": commons.serialize_tensor(grad_norm_g),
"gradient/gradient_penalty": commons.serialize_tensor(gradient_penalty),
}),f,indent=2)
if hps.train.log_interval>0 and global_step % hps.train.log_interval == 0:
if hps.train.log_interval>0 and global_step % hps.train.log_interval == 0: #tensorboard logging
lr = optim_g.param_groups[0]["lr"]
logger.info(f"Train Epoch: {epoch} [{100.0 * (epoch-1) / hps.total_epoch:.2f}% complete]")
# Amor For Tensorboard display
if loss_mel > 75:
loss_mel = 75
@@ -628,7 +601,7 @@ def train_and_evaluate(
"slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
"slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()),
"all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
"all/diff": utils.plot_spectrogram_to_numpy((y_mel[0]-y_mel[0]).abs().data.cpu().numpy(), cmap="hot")
"slice/diff": utils.plot_spectrogram_to_numpy((y_mel[0]-y_hat_mel[0]).abs().data.cpu().numpy(), cmap="hot")
}
with torch.no_grad():
@@ -677,12 +650,47 @@ def train_and_evaluate(
logger.info(f"saving ckpt {save_name}: {status}")
if rank == 0:
logger.info(f"====> Epoch {epoch} (Total Loss={(loss_disc_all+loss_gen_all).item():.3f}): {global_step=} {lr=:.2E} {epoch_recorder.record()}")
total_loss = balancer.weighted_ema_loss + loss_disc_all
logger.info(f"====> Epoch {epoch} ({total_loss=:.3f}): {global_step=} {lr=:.2E} {epoch_recorder.record()}")
logger.info(f"|| {loss_disc_all=:.3f}: {loss_disc=:.3f}, {gradient_penalty=:.3f}")
logger.info(f"|| {loss_gen_all=:.3f}: {loss_gen=:.3f}, {loss_fm=:.3f}, {loss_mel=:.3f}, {loss_kl=:.3f}")
logger.info(f"|| {harmonic_loss=:.3f}, {tefs_loss=:.3f}, {tsi_loss=:.3f}")
balancer.on_epoch_end(.5 / (1 + np.exp(-10 * (epoch / hps.total_epoch - 0.16)))+.5) #sigmoid scaling of ema
if total_loss<least_loss:
least_loss = total_loss
logger.info(f"[lowest loss] {least_loss=:.3f}: {loss_disc=:.3f} {gradient_penalty=:.3f} <==> {loss_gen=:.3f} {loss_fm=:.3f} {loss_mel=:.3f} {loss_kl=:.3f}")
logger.info(f"[aux] {aux_loss=:.3f}: {harmonic_loss=:.3f}, {tefs_loss=:.3f}, {tsi_loss=:.3f}")
if hps.save_best_model:
if hasattr(net_g, "module"): ckpt = net_g.module.state_dict()
else: ckpt = net_g.state_dict()
best_model_name = f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:.0f}" if hps.save_every_weights else f"{hps.name}_loss{least_loss:2.0f}"
status = save_checkpoint(ckpt,best_model_name,epoch,hps)
logger.info(f"=== saving best model {best_model_name}: {status=} ===")
with open(loss_file,"w") as f:
json.dump(dict(least_loss=least_loss.item(),best_model_name=best_model_name,epoch=epoch,steps=global_step,
loss_weights = dict(balancer.ema_weights),
scalar_dict={
"total/loss/all": commons.serialize_tensor(loss_gen_all+loss_disc_all),
"total/loss/gen_all": commons.serialize_tensor(loss_gen_all),
"total/loss/aux": commons.serialize_tensor(aux_loss),
"total/loss/disc_all": commons.serialize_tensor(loss_disc_all),
"total/loss/gen": commons.serialize_tensor(loss_gen),
"total/loss/disc": commons.serialize_tensor(loss_disc),
"total/loss/fm": commons.serialize_tensor(loss_fm),
"total/loss/mel": commons.serialize_tensor(loss_mel),
"total/loss/kl": commons.serialize_tensor(loss_kl),
"aux/loss/harmonic": commons.serialize_tensor(harmonic_loss),
"aux/loss/tefs": commons.serialize_tensor(tefs_loss),
"aux/loss/tsi": commons.serialize_tensor(tsi_loss),
"gradient/grad_norm_disc": commons.serialize_tensor(grad_norm_d),
"gradient/grad_norm_gen": commons.serialize_tensor(grad_norm_g),
"gradient/gradient_penalty": commons.serialize_tensor(gradient_penalty),
}),f,indent=2)
if epoch >= hps.total_epoch and rank == 0:
logger.info("Training is done. The program is closed.")