From b841dd78fe1508207e09b6c581ee8d26654d899e Mon Sep 17 00:00:00 2001 From: ddPn08 Date: Sun, 9 Jul 2023 16:31:38 +0900 Subject: [PATCH] make tracker init_kwargs configurable --- fine_tune.py | 6 +++++- library/train_util.py | 6 ++++++ sdxl_train.py | 6 +++++- train_controlnet.py | 6 +++++- train_db.py | 6 +++++- train_network.py | 6 +++++- train_textual_inversion.py | 6 +++++- train_textual_inversion_XTI.py | 5 ++++- 8 files changed, 40 insertions(+), 7 deletions(-) diff --git a/fine_tune.py b/fine_tune.py index 58a6cda0..a3df178b 100644 --- a/fine_tune.py +++ b/fine_tune.py @@ -6,6 +6,7 @@ import gc import math import os from multiprocessing import Value +import toml from tqdm import tqdm import torch @@ -275,7 +276,10 @@ def train(args): prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device) if accelerator.is_main_process: - accelerator.init_trackers("finetuning" if args.log_tracker_name is None else args.log_tracker_name) + init_kwargs = {} + if args.log_tracker_config is not None: + init_kwargs = toml.load(args.log_tracker_config) + accelerator.init_trackers("finetuning" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs) for epoch in range(num_train_epochs): accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}") diff --git a/library/train_util.py b/library/train_util.py index 809f0af0..84807a2a 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -2445,6 +2445,12 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: default=None, help="name of tracker to use for logging, default is script-specific default name / ログ出力に使用するtrackerの名前、省略時はスクリプトごとのデフォルト名", ) + parser.add_argument( + "--log_tracker_config", + type=str, + default=None, + help="path to tracker config file to use for logging / ログ出力に使用するtrackerの設定ファイルのパス", + ) parser.add_argument( "--wandb_api_key", type=str, diff --git a/sdxl_train.py b/sdxl_train.py index dd5b74dd..a7a92c2a 100644 --- a/sdxl_train.py +++ b/sdxl_train.py @@ -5,6 +5,7 @@ import gc import math import os from multiprocessing import Value +import toml from tqdm import tqdm import torch @@ -350,7 +351,10 @@ def train(args): prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device) if accelerator.is_main_process: - accelerator.init_trackers("finetuning" if args.log_tracker_name is None else args.log_tracker_name) + init_kwargs = {} + if args.log_tracker_config is not None: + init_kwargs = toml.load(args.log_tracker_config) + accelerator.init_trackers("finetuning" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs) for epoch in range(num_train_epochs): accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}") diff --git a/train_controlnet.py b/train_controlnet.py index 39ac43e9..988304f6 100644 --- a/train_controlnet.py +++ b/train_controlnet.py @@ -7,6 +7,7 @@ import random import time from multiprocessing import Value from types import SimpleNamespace +import toml from tqdm import tqdm import torch @@ -324,7 +325,10 @@ def train(args): clip_sample=False, ) if accelerator.is_main_process: - accelerator.init_trackers("controlnet_train" if args.log_tracker_name is None else args.log_tracker_name) + init_kwargs = {} + if args.log_tracker_config is not None: + init_kwargs = toml.load(args.log_tracker_config) + accelerator.init_trackers("controlnet_train" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs) loss_list = [] loss_total = 0.0 diff --git a/train_db.py b/train_db.py index 439f4b9d..13fd651b 100644 --- a/train_db.py +++ b/train_db.py @@ -7,6 +7,7 @@ import itertools import math import os from multiprocessing import Value +import toml from tqdm import tqdm import torch @@ -248,7 +249,10 @@ def train(args): prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device) if accelerator.is_main_process: - accelerator.init_trackers("dreambooth" if args.log_tracker_name is None else args.log_tracker_name) + init_kwargs = {} + if args.log_tracker_config is not None: + init_kwargs = toml.load(args.log_tracker_config) + accelerator.init_trackers("dreambooth" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs) loss_list = [] loss_total = 0.0 diff --git a/train_network.py b/train_network.py index f7ee451b..b500ae7e 100644 --- a/train_network.py +++ b/train_network.py @@ -8,6 +8,7 @@ import random import time import json from multiprocessing import Value +import toml from tqdm import tqdm import torch @@ -672,7 +673,10 @@ class NetworkTrainer: prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device) if accelerator.is_main_process: - accelerator.init_trackers("network_train" if args.log_tracker_name is None else args.log_tracker_name) + init_kwargs = {} + if args.log_tracker_config is not None: + init_kwargs = toml.load(args.log_tracker_config) + accelerator.init_trackers("network_train" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs) loss_list = [] loss_total = 0.0 diff --git a/train_textual_inversion.py b/train_textual_inversion.py index 09294048..0ceab18a 100644 --- a/train_textual_inversion.py +++ b/train_textual_inversion.py @@ -3,6 +3,7 @@ import gc import math import os from multiprocessing import Value +import toml from tqdm import tqdm import torch @@ -493,7 +494,10 @@ class TextualInversionTrainer: prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device) if accelerator.is_main_process: - accelerator.init_trackers("textual_inversion" if args.log_tracker_name is None else args.log_tracker_name) + init_kwargs = {} + if args.log_tracker_config is not None: + init_kwargs = toml.load(args.log_tracker_config) + accelerator.init_trackers("textual_inversion" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs) # function for saving/removing def save_model(ckpt_name, embs_list, steps, epoch_no, force_sync_upload=False): diff --git a/train_textual_inversion_XTI.py b/train_textual_inversion_XTI.py index a08c3a82..33fef943 100644 --- a/train_textual_inversion_XTI.py +++ b/train_textual_inversion_XTI.py @@ -386,7 +386,10 @@ def train(args): prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device) if accelerator.is_main_process: - accelerator.init_trackers("textual_inversion" if args.log_tracker_name is None else args.log_tracker_name) + init_kwargs = {} + if args.log_tracker_config is not None: + init_kwargs = toml.load(args.log_tracker_config) + accelerator.init_trackers("textual_inversion" if args.log_tracker_name is None else args.log_tracker_name, init_kwargs=init_kwargs) # function for saving/removing def save_model(ckpt_name, embs, steps, epoch_no, force_sync_upload=False):