Add WandB logging support

[p1atdev]

関連 #418

WandB でログを取得するオプションを追加しました。利用するには事前に pip install wandb で wandb のインストールと wandb login で必要になります。

以下の引数を追加しました。

• --log_with: 使用するログ取得ツールで、現在は tensorboard、wandb、all が選択できます。all ではすべてを同時に利用できます。指定されなかった場合は従来どおり tensorboard が使用されます。
• --log_tracker_name: WandB で表示されるプロジェクト名(トラッカー名)を指定できます。指定されなかった場合は、textual_inversion や dreambooth などの学習するものの名前が使われます。(log_prefix は wandb では適用されません)

log_with で wandb が有効のときに logging_dir が指定されなかった場合は、実行ディレクトリに wandb ディレクトリが作成されてそこに保存されます。そのため .gitignore に wandb ディレクトリを追加しました。

log_tracker_name は全ての学習スクリプト上の accelerate.init_trackers() で設定しているのですが、ログに関するコードが増えたら、ログ用に何かの関数を用意してそこで設定するようにしたほうがよいかもしれないです...

[kohya-ss]

ありがとうございます。大変便利そうですね。

細かい修正を行い（不要なif判定の削除と、wandb用のmakedirsにexist_ok=Trueを追加）、マージしました。

手元ではwandbの動作確認をしていないのですが、--log_withにwandbを指定し、--logging_dirオプションを指定しない場合、wandbのログが出力されないように思えますがいかがでしょうか。なぜなら、学習スクリプト内のログ出力時に以下のようにargs.logging_dirの有無を確認しているからです。

            if args.logging_dir is not None:
                logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler)
                accelerator.log(logs, step=global_step)

もし動くようなら失礼しました。もし動かないようならtrain_util.pyの以下の部分でargs.logging_dirに値を入れればいいかなと思います。

            if logging_dir is not None:
                os.makedirs(logging_dir, exist_ok=True)
                os.environ["WANDB_DIR"] = logging_dir

[p1atdev]

すみません、その部分の処理を完全に見落としていました...
おっしゃる通り、log_with でwandbが指定されて logging_dir が何も指定されない場合、 generate_step_logs などが実行されず、ログが記録されていませんでした。

train_util.py の if args.log_with is not None: のifブロックの最後で args.logging_dir = "" を追加したところ正しくログが記録されるようになりました。

dashing-shadow-1 が修正前のコードで、画像は記録されているがロスなどが記録されていない

    if args.log_with is not None:
        log_with = args.log_with
        if log_with in ["tensorboard", "all"]:
            if logging_dir is None:
                raise ValueError("logging_dir is required when log_with is tensorboard / Tensorboardを使う場合、logging_dirを指定してください")
        if log_with in ["wandb", "all"]:
            try:
                import wandb
            except ImportError:
                raise ImportError("No wandb / wandb がインストールされていないようです")
            if logging_dir is not None:
                os.makedirs(logging_dir, exist_ok=True)
                os.environ["WANDB_DIR"] = logging_dir
            
+       args.logging_dir = ""

log_with が設定されている場合は args.logging_dir が "" に設定されるようになりましたが、これで特に問題がないようであれば再度 PR を送ります。

[kohya-ss]

確認ありがとうございます。

log_withにwandbを指定して、ディレクトリ名を省略したときのディレクトリ名がwandbで固定なら、それをargs.logging_dirに入れるのがいいかもしれませんね。現在、このメソッド以降はargs.logging_dirは使用されていませんが、将来的に、ログ関係で何かの処理を追加したときにここに実際とは異なる値が入っていると、バグを呼びそうです。

（args.logging_dirをログ出力有無の判断に使っているのがそもそも微妙ですが……。）

以下のようになるかと思いますがいかがでしょうか。

        if log_with in ["wandb", "all"]:
            try:
                import wandb
            except ImportError:
                raise ImportError("No wandb / wandb がインストールされていないようです")
            if logging_dir is not None:
                os.makedirs(logging_dir, exist_ok=True)
                os.environ["WANDB_DIR"] = logging_dir
            else:
                args.logging_dir = "wandb"

これで問題ないようでしたら、私の方で修正させていただきます。

[p1atdev]

ありがとうございます。
正しくログが保存されることを確認しました。問題なさそうです。

[kohya-ss]

確認ありがとうございます。関連するPRをいただいているので、マージ時にあわせて修正いたします。