Experimental: TI via LoRA

1 files changed, 2 insertions, 22 deletions

diff --git a/train_ti.py b/train_ti.py
index 5482326..0ce0056 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -353,7 +353,7 @@ def parse_args():
     parser.add_argument(
         "--adam_weight_decay",
         type=float,
-        default=0,
+        default=1e-2,
         help="Weight decay to use."
     )
     parser.add_argument(
@@ -451,23 +451,6 @@ def parse_args():
         help="The weight of prior preservation loss."
     )
     parser.add_argument(
-        "--use_emb_decay",
-        action="store_true",
-        help="Whether to use embedding decay."
-    )
-    parser.add_argument(
-        "--emb_decay_target",
-        default=0.4,
-        type=float,
-        help="Embedding decay target."
-    )
-    parser.add_argument(
-        "--emb_decay",
-        default=1e2,
-        type=float,
-        help="Embedding decay factor."
-    )
-    parser.add_argument(
         "--noise_timesteps",
         type=int,
         default=1000,
@@ -732,9 +715,6 @@ def main():
         sample_scheduler=sample_scheduler,
         checkpoint_output_dir=checkpoint_output_dir,
         gradient_checkpointing=args.gradient_checkpointing,
-        use_emb_decay=args.use_emb_decay,
-        emb_decay_target=args.emb_decay_target,
-        emb_decay=args.emb_decay,
         use_ema=args.use_ema,
         ema_inv_gamma=args.ema_inv_gamma,
         ema_power=args.ema_power,
@@ -800,7 +780,7 @@ def main():
             sample_frequency = math.ceil(num_train_epochs * (sample_frequency / args.num_train_steps))
 
         optimizer = create_optimizer(
-            text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
+            text_encoder.text_model.embeddings.overlay.parameters(),
             lr=args.learning_rate,
         )