1 files changed, 11 insertions, 4 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index f1dca7f..d2e60ec 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -302,6 +302,12 @@ def parse_args():
         help='Optimizer to use ["adam", "adam8bit", "dadam", "dadan"]'
     )
     parser.add_argument(
+        "--dadaptation_d0",
+        type=float,
+        default=1e-6,
+        help="The d0 parameter for Dadaptation optimizers."
+    )
+    parser.add_argument(
         "--adam_beta1",
         type=float,
         default=0.9,
@@ -535,6 +541,7 @@ def main():
             weight_decay=args.adam_weight_decay,
             eps=args.adam_epsilon,
             decouple=True,
+            d0=args.dadaptation_d0,
         )
 
         args.learning_rate = 1.0
@@ -548,6 +555,7 @@ def main():
             dadaptation.DAdaptAdan,
             weight_decay=args.adam_weight_decay,
             eps=args.adam_epsilon,
+            d0=args.dadaptation_d0,
         )
 
         args.learning_rate = 1.0
@@ -596,10 +604,9 @@ def main():
     datamodule.setup()
 
     num_train_epochs = args.num_train_epochs
-
     if num_train_epochs is None:
-        num_images = math.ceil(len(datamodule.train_dataset) / args.train_batch_size) * args.train_batch_size
-        num_train_epochs = math.ceil(args.num_train_steps / num_images)
+        num_train_epochs = math.ceil(args.num_train_steps / len(datamodule.train_dataset))
+        sample_frequency = math.ceil(num_train_epochs * (args.sample_frequency / args.num_train_steps))
 
     params_to_optimize = (unet.parameters(), )
     if args.train_text_encoder_epochs != 0:
@@ -639,7 +646,7 @@ def main():
         lr_scheduler=lr_scheduler,
         num_train_epochs=num_train_epochs,
         gradient_accumulation_steps=args.gradient_accumulation_steps,
-        sample_frequency=args.sample_frequency,
+        sample_frequency=sample_frequency,
         offset_noise_strength=args.offset_noise_strength,
         # --
         tokenizer=tokenizer,