Fixed accuracy calc, other improvements

1 files changed, 29 insertions, 1 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 8fd78f1..1ebcfe3 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -232,6 +232,30 @@ def parse_args():
         help="Number of restart cycles in the lr scheduler (if supported)."
     )
     parser.add_argument(
+        "--lr_warmup_func",
+        type=str,
+        default="cos",
+        help='Choose between ["linear", "cos"]'
+    )
+    parser.add_argument(
+        "--lr_warmup_exp",
+        type=int,
+        default=1,
+        help='If lr_warmup_func is "cos", exponent to modify the function'
+    )
+    parser.add_argument(
+        "--lr_annealing_func",
+        type=str,
+        default="cos",
+        help='Choose between ["linear", "half_cos", "cos"]'
+    )
+    parser.add_argument(
+        "--lr_annealing_exp",
+        type=int,
+        default=3,
+        help='If lr_annealing_func is "half_cos" or "cos", exponent to modify the function'
+    )
+    parser.add_argument(
         "--use_ema",
         action="store_true",
         default=True,
@@ -760,6 +784,10 @@ def main():
         lr_scheduler = get_one_cycle_schedule(
             optimizer=optimizer,
             num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+            warmup=args.lr_warmup_func,
+            annealing=args.lr_annealing_func,
+            warmup_exp=args.lr_warmup_exp,
+            annealing_exp=args.lr_annealing_exp,
         )
     elif args.lr_scheduler == "cosine_with_restarts":
         lr_scheduler = get_cosine_with_hard_restarts_schedule_with_warmup(
@@ -913,7 +941,7 @@ def main():
         else:
             loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
 
-        acc = (model_pred == latents).float().mean()
+        acc = (model_pred == target).float().mean()
 
         return loss, acc, bsz