From b7b9f7a7fc3a2e6a027175e5a84541ca2291fbb5 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sun, 1 Jan 2023 11:36:00 +0100
Subject: Fixed accuracy calc, other improvements

---
 train_dreambooth.py | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

(limited to 'train_dreambooth.py')

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 8fd78f1..1ebcfe3 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -231,6 +231,30 @@ def parse_args():
         default=None,
         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",
@@ -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
 
-- 
cgit v1.2.3-54-g00ecf