From eccc07537bf89e6fbd8f46c2855f3e5cd42ff952 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=9D=92=E9=BE=8D=E8=81=96=E8=80=85=40bdsqlsz?=
 <qinglongshengzhe@gmail.com>
Date: Thu, 27 Apr 2023 00:17:32 +0800
Subject: [PATCH] Update train_util.py for DAdapt

---
 library/train_util.py | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/library/train_util.py b/library/train_util.py
index c80c4cca..e7498f07 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -1883,7 +1883,7 @@ def add_optimizer_arguments(parser: argparse.ArgumentParser):
         "--optimizer_type",
         type=str,
         default="",
-        help="Optimizer to use / オプティマイザの種類: AdamW (default), AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, AdaFactor",
+        help="Optimizer to use / オプティマイザの種類: AdamW (default), AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation(DAdaptAdam), DAdaptAdaGrad, DAdaptAdan, DAdaptSGD, AdaFactor",
     )
 
     # backward compatibility
@@ -2448,7 +2448,7 @@ def resume_from_local_or_hf_if_specified(accelerator, args):
 
 
 def get_optimizer(args, trainable_params):
-    # "Optimizer to use: AdamW, AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, Adafactor"
+    # "Optimizer to use: AdamW, AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, DAdaptation(DAdaptAdam), DAdaptAdaGrad, DAdaptAdan, DAdaptSGD, Adafactor"
 
     optimizer_type = args.optimizer_type
     if args.use_8bit_adam:
@@ -2563,6 +2563,35 @@ def get_optimizer(args, trainable_params):
 
         optimizer_class = dadaptation.DAdaptAdam
         optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
+        
+    elif optimizer_type == "DAdaptAdaGrad".lower():
+        try:
+            import dadaptation
+        except ImportError:
+            raise ImportError("No dadaptation / dadaptation がインストールされていないようです")
+        print(f"use D-Adaptation AdaGrad optimizer | {optimizer_kwargs}")
+
+        actual_lr = lr
+        lr_count = 1
+        if type(trainable_params) == list and type(trainable_params[0]) == dict:
+            lrs = set()
+            actual_lr = trainable_params[0].get("lr", actual_lr)
+            for group in trainable_params:
+                lrs.add(group.get("lr", actual_lr))
+            lr_count = len(lrs)
+
+        if actual_lr <= 0.1:
+            print(
+                f"learning rate is too low. If using dadaptation, set learning rate around 1.0 / 学習率が低すぎるようです。1.0前後の値を指定してください: lr={actual_lr}"
+            )
+            print("recommend option: lr=1.0 / 推奨は1.0です")
+        if lr_count > 1:
+            print(
+                f"when multiple learning rates are specified with dadaptation (e.g. for Text Encoder and U-Net), only the first one will take effect / D-Adaptationで複数の学習率を指定した場合（Text EncoderとU-Netなど）、最初の学習率のみが有効になります: lr={actual_lr}"
+            )
+
+        optimizer_class = dadaptation.DAdaptAdaGrad
+        optimizer = optimizer_class(trainable_params, lr=lr, **optimizer_kwargs)
     
     elif optimizer_type == "DAdaptAdan".lower():
         try: