1 files changed, 11 insertions, 4 deletions

diff --git a/train_ti.py b/train_ti.py
index b7ea5f3..902f508 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -320,6 +320,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,
@@ -659,6 +665,7 @@ def main():
             weight_decay=args.adam_weight_decay,
             eps=args.adam_epsilon,
             decouple=True,
+            d0=args.dadaptation_d0,
         )
     elif args.optimizer == 'dadan':
         try:
@@ -670,6 +677,7 @@ def main():
             dadaptation.DAdaptAdan,
             weight_decay=args.adam_weight_decay,
             eps=args.adam_epsilon,
+            d0=args.dadaptation_d0,
         )
     else:
         raise ValueError(f"Unknown --optimizer \"{args.optimizer}\"")
@@ -690,7 +698,6 @@ def main():
         no_val=args.valid_set_size == 0,
         strategy=textual_inversion_strategy,
         gradient_accumulation_steps=args.gradient_accumulation_steps,
-        sample_frequency=args.sample_frequency,
         checkpoint_frequency=args.checkpoint_frequency,
         milestone_checkpoints=not args.no_milestone_checkpoints,
         global_step_offset=global_step_offset,
@@ -759,10 +766,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))
 
         optimizer = create_optimizer(
             text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
@@ -792,6 +798,7 @@ def main():
             optimizer=optimizer,
             lr_scheduler=lr_scheduler,
             num_train_epochs=num_train_epochs,
+            sample_frequency=sample_frequency,
             # --
             sample_output_dir=sample_output_dir,
             placeholder_tokens=placeholder_tokens,