diff --git a/webui_train.py b/webui_train.py index 1463bde..94b1bc3 100644 --- a/webui_train.py +++ b/webui_train.py @@ -30,7 +30,17 @@ def get_path(model_name): return dataset_path, lbl_path, train_path, val_path, config_path -def initialize(model_name, batch_size, epochs, save_every_steps, bf16_run): +def initialize( + model_name, + batch_size, + epochs, + save_every_steps, + bf16_run, + freeze_EN_bert, + freeze_JP_bert, + freeze_ZH_bert, + freeze_style, +): global logger_handler dataset_path, _, train_path, val_path, config_path = get_path(model_name) @@ -42,7 +52,7 @@ def initialize(model_name, batch_size, epochs, save_every_steps, bf16_run): logger_handler = logger.add(os.path.join(dataset_path, file_name)) logger.info( - f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, bf16_run: {bf16_run}" + f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, bf16_run: {bf16_run}, freeze_ZH_bert: {freeze_ZH_bert}, freeze_JP_bert: {freeze_JP_bert}, freeze_EN_bert: {freeze_EN_bert}, freeze_style: {freeze_style}" ) if os.path.isfile(config_path): with open(config_path, "r", encoding="utf-8") as f: @@ -59,6 +69,11 @@ def initialize(model_name, batch_size, epochs, save_every_steps, bf16_run): config["train"]["bf16_run"] = bf16_run config["train"]["eval_interval"] = save_every_steps + config["train"]["freeze_EN_bert"] = freeze_EN_bert + config["train"]["freeze_JP_bert"] = freeze_JP_bert + config["train"]["freeze_ZH_bert"] = freeze_ZH_bert + config["train"]["freeze_style"] = freeze_style + model_path = os.path.join(dataset_path, "models") if os.path.exists(model_path): logger.warning(f"Step 1: {model_path} already exists, so copy it to backup.") @@ -214,11 +229,23 @@ def preprocess_all( num_processes, normalize, trim, + freeze_EN_bert, + freeze_JP_bert, + freeze_ZH_bert, + freeze_style, ): if model_name == "": return False, "Error: モデル名を入力してください" success, message = initialize( - model_name, batch_size, epochs, save_every_steps, bf16_run + model_name, + batch_size, + epochs, + save_every_steps, + bf16_run, + freeze_EN_bert, + freeze_JP_bert, + freeze_ZH_bert, + freeze_style, ) if not success: return False, message @@ -349,14 +376,6 @@ if __name__ == "__main__": info="bfloat16を使うかどうか。新しめのグラボだと学習が早くなるかも、古いグラボだと動かないかも。", value=True, ) - num_processes = gr.Slider( - label="プロセス数", - info="前処理時の並列処理プロセス数、大きすぎるとフリーズするかも", - value=cpu_count() // 2, - minimum=1, - maximum=cpu_count(), - step=1, - ) normalize = gr.Checkbox( label="音声の音量を正規化する(音量の大小が揃っていない場合など)", value=False, @@ -365,6 +384,33 @@ if __name__ == "__main__": label="音声の最初と最後の無音を取り除く", value=False, ) + with gr.Accordion("詳細設定", open=False): + num_processes = gr.Slider( + label="プロセス数", + info="前処理時の並列処理プロセス数、大きすぎるとフリーズするかも", + value=cpu_count() // 2, + minimum=1, + maximum=cpu_count(), + step=1, + ) + gr.Markdown("学習時に特定の部分を凍結させるかどうか") + freeze_EN_bert = gr.Checkbox( + label="英語bert部分を凍結", + value=False, + ) + freeze_JP_bert = gr.Checkbox( + label="日本語bert部分を凍結", + value=False, + ) + freeze_ZH_bert = gr.Checkbox( + label="中国語bert部分を凍結", + value=False, + ) + freeze_style = gr.Checkbox( + label="スタイル部分を凍結", + value=False, + ) + with gr.Column(): preprocess_button = gr.Button(value="自動前処理を実行", variant="primary") info_all = gr.Textbox(label="状況") @@ -397,6 +443,22 @@ if __name__ == "__main__": label="bfloat16を使う", value=True, ) + freeze_EN_bert_manual = gr.Checkbox( + label="英語bert部分を凍結", + value=False, + ) + freeze_JP_bert_manual = gr.Checkbox( + label="日本語bert部分を凍結", + value=False, + ) + freeze_ZH_bert_manual = gr.Checkbox( + label="中国語bert部分を凍結", + value=False, + ) + freeze_style_manual = gr.Checkbox( + label="スタイル部分を凍結", + value=False, + ) with gr.Column(): generate_config_btn = gr.Button(value="実行", variant="primary") info_init = gr.Textbox(label="状況") @@ -467,6 +529,10 @@ if __name__ == "__main__": num_processes, normalize, trim, + freeze_EN_bert, + freeze_JP_bert, + freeze_ZH_bert, + freeze_style, ], outputs=[info_all], ) @@ -478,6 +544,10 @@ if __name__ == "__main__": epochs_manual, save_every_steps_manual, bf16_run_manual, + freeze_EN_bert_manual, + freeze_JP_bert_manual, + freeze_ZH_bert_manual, + freeze_style_manual, ], outputs=[info_init], )