diff --git a/src/utils.py b/src/utils.py index f8d008f..82ea25b 100755 --- a/src/utils.py +++ b/src/utils.py @@ -42,12 +42,12 @@ MODELS['dvae.pth'] = "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/370 WHISPER_MODELS = ["tiny", "base", "small", "medium", "large", "large-v2"] WHISPER_SPECIALIZED_MODELS = ["tiny.en", "base.en", "small.en", "medium.en"] WHISPER_BACKENDS = ["openai/whisper", "lightmare/whispercpp", "m-bain/whisperx"] - VOCODERS = ['univnet', 'bigvgan_base_24khz_100band', 'bigvgan_24khz_100band'] GENERATE_SETTINGS_ARGS = None -EPOCH_SCHEDULE = [ 9, 18, 25, 33 ] +LEARNING_RATE_SCHEMES = {"Multistep": "MultiStepLR", "Cos. Annealing": "CosineAnnealingLR_Restart"} +LEARNING_RATE_SCHEDULE = [ 9, 18, 25, 33 ] args = None tts = None @@ -1215,7 +1215,7 @@ def calc_iterations( epochs, lines, batch_size ): iterations = int(epochs * lines / float(batch_size)) return iterations -def schedule_learning_rate( iterations, schedule=EPOCH_SCHEDULE ): +def schedule_learning_rate( iterations, schedule=LEARNING_RATE_SCHEDULE ): return [int(iterations * d) for d in schedule] def optimize_training_settings( **kwargs ): @@ -1378,14 +1378,15 @@ def save_training_settings( **kwargs ): settings['optimizer'] = 'adamw' if settings['gpus'] == 1 else 'adamw_zero' - LEARNING_RATE_SCHEMES = ["MultiStepLR", "CosineAnnealingLR_Restart"] if 'learning_rate_scheme' not in settings or settings['learning_rate_scheme'] not in LEARNING_RATE_SCHEMES: - settings['learning_rate_scheme'] = LEARNING_RATE_SCHEMES[0] + settings['learning_rate_scheme'] = "Multistep" + + settings['learning_rate_scheme'] = LEARNING_RATE_SCHEMES[settings['learning_rate_scheme']] learning_rate_schema = [f"default_lr_scheme: {settings['learning_rate_scheme']}"] if settings['learning_rate_scheme'] == "MultiStepLR": if not settings['learning_rate_schedule']: - settings['learning_rate_schedule'] = EPOCH_SCHEDULE + settings['learning_rate_schedule'] = LEARNING_RATE_SCHEDULE elif isinstance(settings['learning_rate_schedule'],str): settings['learning_rate_schedule'] = json.loads(settings['learning_rate_schedule']) @@ -1395,38 +1396,30 @@ def save_training_settings( **kwargs ): learning_rate_schema.append(f" lr_gamma: 0.5") elif settings['learning_rate_scheme'] == "CosineAnnealingLR_Restart": epochs = settings['epochs'] - restarts = int(epochs / 2) + restarts = settings['learning_rate_restarts'] + restart_period = int(epochs / restarts) - if 'learning_rate_period' not in settings: - settings['learning_rate_period'] = [ iterations_per_epoch for x in range(epochs) ] if 'learning_rate_warmup' not in settings: settings['learning_rate_warmup'] = 0 if 'learning_rate_min' not in settings: - settings['learning_rate_min'] = 1e-07 - if 'learning_rate_restarts' not in settings: - settings['learning_rate_restarts'] = [ iterations_per_epoch * (x+1) * 2 for x in range(restarts) ] # [52, 104, 156, 208] + settings['learning_rate_min'] = 1e-08 + + if 'learning_rate_period' not in settings: + settings['learning_rate_period'] = [ iterations_per_epoch * restart_period for x in range(epochs) ] + + settings['learning_rate_restarts'] = [ iterations_per_epoch * (x+1) * restart_period for x in range(restarts) ] # [52, 104, 156, 208] + if 'learning_rate_restart_weights' not in settings: settings['learning_rate_restart_weights'] = [ ( restarts - x - 1 ) / restarts for x in range(restarts) ] # [.75, .5, .25, .125] settings['learning_rate_restart_weights'][-1] = settings['learning_rate_restart_weights'][-2] * 0.5 learning_rate_schema.append(f" T_period: {settings['learning_rate_period']}") - learning_rate_schema.append(f" warmup: !!float {settings['learning_rate_warmup']}") + learning_rate_schema.append(f" warmup: {settings['learning_rate_warmup']}") learning_rate_schema.append(f" eta_min: !!float {settings['learning_rate_min']}") learning_rate_schema.append(f" restarts: {settings['learning_rate_restarts']}") learning_rate_schema.append(f" restart_weights: {settings['learning_rate_restart_weights']}") settings['learning_rate_scheme'] = "\n".join(learning_rate_schema) - """ - if resume_state: - settings['pretrain_model_gpt'] = f"# {settings['pretrain_model_gpt']}" - else: - settings['resume_state'] = f"# resume_state: './training/{voice}/training_state/#.state'" - - # also disable validation if it doesn't make sense to do it - if settings['dataset_path'] == settings['validation_path'] or not os.path.exists(settings['validation_path']): - settings['validation_enabled'] = 'false' - """ - if settings['resume_state']: settings['source_model'] = f"# pretrain_model_gpt: {settings['source_model']}" settings['resume_state'] = f"resume_state: {settings['resume_state']}'" @@ -1815,10 +1808,6 @@ def save_args_settings(): f.write(json.dumps(settings, indent='\t') ) # super kludgy )`; -def set_generate_settings_arg_order(args): - global GENERATE_SETTINGS_ARGS - GENERATE_SETTINGS_ARGS = args - def import_generate_settings(file="./config/generate.json"): global GENERATE_SETTINGS_ARGS diff --git a/src/webui.py b/src/webui.py index edea0a2..281f3ac 100755 --- a/src/webui.py +++ b/src/webui.py @@ -277,7 +277,6 @@ def setup_gradio(): for i in range(len(GENERATE_SETTINGS_ARGS)): arg = GENERATE_SETTINGS_ARGS[i] GENERATE_SETTINGS[arg] = None - set_generate_settings_arg_order(GENERATE_SETTINGS_ARGS) with gr.Blocks() as ui: with gr.Tab("Generate"): @@ -402,11 +401,23 @@ def setup_gradio(): with gr.Column(): TRAINING_SETTINGS["epochs"] = gr.Number(label="Epochs", value=500, precision=0) with gr.Row(): - with gr.Column(): - TRAINING_SETTINGS["learning_rate"] = gr.Slider(label="Learning Rate", value=1e-5, minimum=0, maximum=1e-4, step=1e-6) - TRAINING_SETTINGS["text_ce_lr_weight"] = gr.Slider(label="Text_CE LR Ratio", value=0.01, minimum=0, maximum=1) + TRAINING_SETTINGS["learning_rate"] = gr.Slider(label="Learning Rate", value=1e-5, minimum=0, maximum=1e-4, step=1e-6) + TRAINING_SETTINGS["text_ce_lr_weight"] = gr.Slider(label="Text_CE LR Ratio", value=0.01, minimum=0, maximum=1) - TRAINING_SETTINGS["learning_rate_schedule"] = gr.Textbox(label="Learning Rate Schedule", placeholder=str(EPOCH_SCHEDULE)) + with gr.Row(): + lr_schemes = list(LEARNING_RATE_SCHEMES.keys()) + TRAINING_SETTINGS["learning_rate_scheme"] = gr.Radio(lr_schemes, label="Learning Rate Scheme", value=lr_schemes[0], type="value") + TRAINING_SETTINGS["learning_rate_schedule"] = gr.Textbox(label="Learning Rate Schedule", placeholder=str(LEARNING_RATE_SCHEDULE), visible=True) + TRAINING_SETTINGS["learning_rate_restarts"] = gr.Number(label="Learning Rate Restarts", value=4, precision=0, visible=False) + + TRAINING_SETTINGS["learning_rate_scheme"].change( + fn=lambda x: ( gr.update(visible=x == lr_schemes[0]), gr.update(visible=x == lr_schemes[1]) ), + inputs=TRAINING_SETTINGS["learning_rate_scheme"], + outputs=[ + TRAINING_SETTINGS["learning_rate_schedule"], + TRAINING_SETTINGS["learning_rate_restarts"], + ] + ) with gr.Row(): TRAINING_SETTINGS["batch_size"] = gr.Number(label="Batch Size", value=128, precision=0) TRAINING_SETTINGS["gradient_accumulation_size"] = gr.Number(label="Gradient Accumulation Size", value=4, precision=0)