1 files changed, 9 insertions, 1 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 9c1e41c..a70c80e 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -16,6 +16,7 @@ from slugify import slugify
 from util import load_config, load_embeddings_from_dir
 from data.csv import VlpnDataModule, keyword_filter
 from training.functional import train, get_models
+from training.lr import plot_metrics
 from training.strategy.dreambooth import dreambooth_strategy
 from training.optimization import get_scheduler
 from training.util import save_args
@@ -524,6 +525,10 @@ def main():
             args.train_batch_size * accelerator.num_processes
         )
 
+    if args.find_lr:
+        args.learning_rate = 1e-6
+        args.lr_scheduler = "exponential_growth"
+
     if args.use_8bit_adam:
         try:
             import bitsandbytes as bnb
@@ -602,11 +607,12 @@ def main():
         warmup_exp=args.lr_warmup_exp,
         annealing_exp=args.lr_annealing_exp,
         cycles=args.lr_cycles,
+        end_lr=1e2,
         train_epochs=args.num_train_epochs,
         warmup_epochs=args.lr_warmup_epochs,
     )
 
-    trainer(
+    metrics = trainer(
         strategy=dreambooth_strategy,
         project="dreambooth",
         train_dataloader=datamodule.train_dataloader,
@@ -634,6 +640,8 @@ def main():
         sample_image_size=args.sample_image_size,
     )
 
+    plot_metrics(metrics, output_dir.joinpath("lr.png"))
+
 
 if __name__ == "__main__":
     main()