From 5ecd6a84f288fd1574d5047ce59e911ceef70d3b Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Fri, 30 Aug 2024 23:04:28 -0400 Subject: [PATCH] refactored training params to node --- custom_nodes/rvc_nodes.py | 62 +++++++++++++++++++++++++-------------- lib/train/losses.py | 2 +- lib/train/utils.py | 60 +++++++++++++++---------------------- training_cli.py | 20 ++++++------- 4 files changed, 74 insertions(+), 70 deletions(-) diff --git a/custom_nodes/rvc_nodes.py b/custom_nodes/rvc_nodes.py index 78bbe9a..39c7b5f 100644 --- a/custom_nodes/rvc_nodes.py +++ b/custom_nodes/rvc_nodes.py @@ -8,7 +8,7 @@ import torch from . import BASE_DIR -from ..lib.train.utils import HParams +from ..lib.train.utils import HParams, DEFAULT_TRAINING_PARAMS from ..training_cli import train_model @@ -153,13 +153,9 @@ class LoadRVCModelNode: class RVCNode: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls): - return { "required": { "audio": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),), @@ -213,9 +209,6 @@ class RVCNode: class RVCProcessDatasetNode: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls): @@ -358,10 +351,36 @@ class RVCProcessDatasetNode: ),) -class RVCTrainModelNode: +class RVCTrainParamsNode: - def __init__(self): - pass + @classmethod + def INPUT_TYPES(cls): + + return { + "optional": dict( + batch_size=("INT",dict(default=4,min=1,max=64,step=1)), + c_mel=("FLOAT",dict(default=45.,min=0.,max=100.,step=.1)), + c_kl=("FLOAT",dict(default=1.,min=0.,max=100.,step=.1)), + c_fm=("FLOAT",dict(default=2.,min=0.,max=100.,step=.1)), + c_mfcc=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), + c_lfcc=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), + c_hd=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), + c_sts=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), + c_gp=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), + ) + } + + RETURN_TYPES = ('RVC_TRAINING_PARAMS', ) + RETURN_NAMES = ("rvc_training_params", ) + + FUNCTION = "init" + + CATEGORY = CATEGORY + + def init(self,batch_size=4,c_mel=45,c_kl=1.,c_fm=1.,c_mfcc=0.,c_lfcc=0.,c_hd=0.,c_sts=0.,c_gp=0.,): + return (dict(batch_size=batch_size,c_mel=c_mel,c_kl=c_kl,c_fm=c_fm,c_mfcc=c_mfcc,c_gp=c_gp,c_lfcc=c_lfcc,c_hd=c_hd,c_sts=c_sts),) + +class RVCTrainModelNode: @classmethod def INPUT_TYPES(cls): @@ -374,10 +393,10 @@ class RVCTrainModelNode: return { "required": { "rvc_dataset_pipe": ("RVC_DATASET_PIPE",), + "rvc_training_params": ("RVC_TRAINING_PARAMS",dict(default=DEFAULT_TRAINING_PARAMS)), }, "optional": dict( gpu=(DEVICES, {"default": DEVICES[0]}), - batch_size=("INT",dict(default=4,min=1,max=64,step=1)), total_epoch=("INT",dict(default=100,min=10,max=1000,step=10)), save_every_epoch=("INT",dict(default=0,min=0,max=100)), pretrained_G=(PRETRAINED_G,{"default": PRETRAINED_G[0]}), @@ -388,9 +407,8 @@ class RVCTrainModelNode: train_index=("BOOLEAN",{"default": True}), retrain=("BOOLEAN",{"default": False}), save_best_model=("BOOLEAN",{"default": True}), + best_model_threshold=("INT",dict(default=50,min=20,max=80)), log_every_epoch=("FLOAT",dict(default=1.,min=0.,max=2.,step=.5)), - gradient_lambda=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), - timbre_lambda=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), num_workers=("INT",dict(default=1,min=1,max=16)) ) } @@ -406,8 +424,8 @@ class RVCTrainModelNode: def train_model(self, rvc_dataset_pipe, + rvc_training_params, gpu="0", - batch_size=4, total_epoch=100, save_every_epoch=0, pretrained_G="", @@ -418,15 +436,14 @@ class RVCTrainModelNode: train_index=True, retrain=False, save_best_model=True, + best_model_threshold=50, log_every_epoch=1., - gradient_lambda=0., - timbre_lambda=0., num_workers=1): sample_rate = rvc_dataset_pipe["sample_rate"] name = rvc_dataset_pipe["name"] dataset_dir = rvc_dataset_pipe["dataset_dir"] - cache_name = get_hash(batch_size,pretrained_G,pretrained_D,gradient_lambda,timbre_lambda) + cache_name = get_hash(pretrained_G,pretrained_D,*rvc_training_params.items()) model_dir = os.path.join(output_path,"logs",cache_name) if_f0 = rvc_dataset_pipe["if_f0"] @@ -444,7 +461,6 @@ class RVCTrainModelNode: hparams.pretrainD = model_downloader(pretrained_D) hparams.version = "v2" hparams.gpus = gpu - hparams.train.batch_size = max(batch_size,1) hparams.sample_rate = sample_rate hparams.if_f0 = if_f0 hparams.if_latest = if_save_latest @@ -452,10 +468,10 @@ class RVCTrainModelNode: hparams.if_cache_data_in_gpu = if_cache_gpu hparams.data.training_files = rvc_dataset_pipe["training_files"] hparams.save_best_model = save_best_model + hparams.best_model_threshold = best_model_threshold hparams.log_every_epoch = log_every_epoch - hparams.train.gradient_lambda = gradient_lambda hparams.train.num_workers = num_workers - hparams.train.timbre_lambda = timbre_lambda + hparams.train.update(rvc_training_params) file_index = self.train_index(dataset_dir, sample_rate, name) if train_index else None model_path = os.path.join(BASE_MODELS_DIR,"RVC",f"{name}_{sample_rate}.pth") @@ -531,7 +547,8 @@ NODE_CLASS_MAPPINGS = { "LoadHubertModel": LoadHubertModel, "LoadPitchExtractionParams": LoadPitchExtractionParams, "RVCProcessDatasetNode": RVCProcessDatasetNode, - "RVCTrainModelNode": RVCTrainModelNode + "RVCTrainModelNode": RVCTrainModelNode, + "RVCTrainParamsNode": RVCTrainParamsNode } # A dictionary that contains the friendly/humanly readable titles for the nodes @@ -542,4 +559,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LoadPitchExtractionParams": "🌺Load Pitch Extraction Params", "RVCProcessDatasetNode": "🌺Process RVC Dataset", "RVCTrainModelNode": "🌺Train RVC Model", + "RVCTrainParamsNode": "🌺RVC Training Params" } \ No newline at end of file diff --git a/lib/train/losses.py b/lib/train/losses.py index 41a7dd0..9792fb1 100644 --- a/lib/train/losses.py +++ b/lib/train/losses.py @@ -35,7 +35,7 @@ def feature_loss(fmap_r, fmap_g): gl = gl.float() loss += torch.mean(torch.abs(rl - gl)) - return loss * 2 + return loss def discriminator_loss(disc_real_outputs, disc_generated_outputs): diff --git a/lib/train/utils.py b/lib/train/utils.py index a03e8df..2595f66 100644 --- a/lib/train/utils.py +++ b/lib/train/utils.py @@ -1,4 +1,4 @@ -import os, traceback +import os import glob import sys import argparse @@ -10,6 +10,17 @@ from scipy.io.wavfile import read import torch MATPLOTLIB_FLAG = False +DEFAULT_TRAINING_PARAMS = dict( + batch_size=4, + c_mel=45, + c_kl=1., + c_fm=2., + c_gp=0., + c_mfcc=0., + c_lfcc=0., + c_hd=0., + c_sts=0. +) logging.basicConfig(stream=sys.stdout, level=logging.DEBUG) logger = logging @@ -63,36 +74,6 @@ def load_checkpoint_d(checkpoint_path, combd, sbd, optimizer=None, load_opt=1): logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration)) return model, optimizer, learning_rate, iteration - -# def load_checkpoint(checkpoint_path, model, optimizer=None): -# assert os.path.isfile(checkpoint_path) -# checkpoint_dict = torch.load(checkpoint_path, map_location='cpu') -# iteration = checkpoint_dict['iteration'] -# learning_rate = checkpoint_dict['learning_rate'] -# if optimizer is not None: -# optimizer.load_state_dict(checkpoint_dict['optimizer']) -# # print(1111) -# saved_state_dict = checkpoint_dict['model'] -# # print(1111) -# -# if hasattr(model, 'module'): -# state_dict = model.module.state_dict() -# else: -# state_dict = model.state_dict() -# new_state_dict= {} -# for k, v in state_dict.items(): -# try: -# new_state_dict[k] = saved_state_dict[k] -# except: -# logger.info("%s is not in the checkpoint" % k) -# new_state_dict[k] = v -# if hasattr(model, 'module'): -# model.module.load_state_dict(new_state_dict) -# else: -# model.load_state_dict(new_state_dict) -# logger.info("Loaded checkpoint '{}' (epoch {})" .format( -# checkpoint_path, iteration)) -# return model, optimizer, learning_rate, iteration def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1): assert os.path.isfile(checkpoint_path) checkpoint_dict = torch.load(checkpoint_path, map_location="cpu") @@ -382,10 +363,9 @@ def get_hparams(init=True): hparams.if_cache_data_in_gpu = args.if_cache_data_in_gpu hparams.data.training_files = os.path.join(experiment_dir,"filelist.txt") hparams.save_best_model = True - hparams.train.gradient_lambda = 0. hparams.train.num_workers = 4 - hparams.train.timbre_lambda = 0. hparams.model_path = None + hparams.best_model_threshold = 50 return hparams @@ -449,11 +429,10 @@ def get_logger(model_dir, filename="train.log"): return logger -class HParams: +class HParams(): def __init__(self, **kwargs): for k, v in kwargs.items(): - if isinstance(v, dict): - v = HParams(**v) + if isinstance(v, dict): v = HParams(**v) self[k] = v def keys(self): @@ -483,7 +462,14 @@ class HParams: def get(self, key, default=None): try: return self[key] except: return default + + def update(self, data={}, **kwargs): + data.update(kwargs) + for k, v in data.items(): + if isinstance(v, dict): v = HParams(**v) + if k in self and hasattr(self[k], "update") and isinstance(v, HParams): self[k].update(v) + else: self[k] = v def sync_log_interval(self, dataset_length): - log_interval = getattr(self,"log_every_epoch",1) + log_interval = getattr(self,"log_every_epoch",1.) self.train.log_interval = int(dataset_length*log_interval) \ No newline at end of file diff --git a/training_cli.py b/training_cli.py index 0fa5d9b..98e172f 100644 --- a/training_cli.py +++ b/training_cli.py @@ -127,10 +127,10 @@ def run(rank, n_gpus, hps, device): if os.path.isfile(loss_file): with open(loss_file,"r") as f: data: dict = json.load(f) - least_loss = data.get("least_loss",50) + least_loss = data.get("least_loss",hps.best_model_threshold) best_model_name = data.get("best_model_name","") else: - least_loss = 50 + least_loss = hps.best_model_threshold best_model_name = "" if hps.version == "v1": @@ -481,7 +481,7 @@ def train_and_evaluate( with autocast(enabled=False): gradient_penalty = torch.Tensor([0.]).to(wave.device) - if hps.train.get("gradient_lambda",0.)>0: + if hps.train.get("c_gp",0.)>0: # Compute the gradient penalty # Randomly interpolate between real and generated data size = [1]*wave.ndim @@ -502,10 +502,10 @@ def train_and_evaluate( )[0] gradient = gradient.view(gradient.size(0), -1) gradient_penalty += torch.log(gradient.norm(2,dim=-1)).abs().mean() # force gradnorm=1 - gradient_penalty *= hps.train.gradient_lambda/len(disc_interpolated_output) + gradient_penalty *= hps.train.c_gp/len(disc_interpolated_output) loss_mfcc = torch.Tensor([0.]).to(wave.device) - if hps.train.get("timbre_lambda",0)>0: + if hps.train.get("c_mfcc",0)>0: loss_mfcc = mfcc_loss(wave, gen_wave, SR_MAP[hps.sample_rate], n_fft=hps.data.filter_length, hop_length=hps.data.hop_length, @@ -514,7 +514,7 @@ def train_and_evaluate( norm="slaney", f_min=hps.data.mel_fmin, f_max=hps.data.mel_fmax - ) * hps.train.timbre_lambda + ) * hps.train.c_mfcc loss_mfcc = loss_mfcc.squeeze() gradient_penalty = gradient_penalty.squeeze() @@ -531,9 +531,9 @@ 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.smooth_l1_loss(y_mel, y_hat_mel) * hps.train.c_mel - loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl - loss_fm = feature_loss(fmap_r, fmap_g) + loss_mel = F.smooth_l1_loss(y_mel, y_hat_mel) * hps.train.get("c_mel",45.) + loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.get("c_kl",1.) + loss_fm = feature_loss(fmap_r, fmap_g) * hps.train.get("c_fm",2.) loss_gen, losses_gen = generator_loss(y_d_hat_g) loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl @@ -556,7 +556,7 @@ def train_and_evaluate( else: ckpt = net_g.state_dict() - best_model_name = f"{hps.name}_loss{least_loss:2.0f}" if hps.if_save_latest else f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:2.0f}" + best_model_name = f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:2.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=} ===")