Skip to content

Commit

Permalink
Added cli argument for wandb session name
Browse files Browse the repository at this point in the history
  • Loading branch information
nirweingarten committed Jan 3, 2024
1 parent 1a36f9d commit ab71630
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 0 deletions.
6 changes: 6 additions & 0 deletions library/train_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -2935,6 +2935,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(
"--wandb_run_name",
type=str,
default=None,
help="The name of the specific wandb session / wandb ログに表示される特定の実行の名前",
)
parser.add_argument(
"--log_tracker_config",
type=str,
Expand Down
2 changes: 2 additions & 0 deletions train_network.py
Original file line number Diff line number Diff line change
Expand Up @@ -684,6 +684,8 @@ def train(self, args):

if accelerator.is_main_process:
init_kwargs = {}
if args.wandb_run_name:
init_kwargs['wandb'] = {'name': args.wandb_run_name}
if args.log_tracker_config is not None:
init_kwargs = toml.load(args.log_tracker_config)
accelerator.init_trackers(
Expand Down

0 comments on commit ab71630

Please sign in to comment.