Revert

1 files changed, 46 insertions, 6 deletions

diff --git a/train_ti.py b/train_ti.py
index 26ac384..5482326 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -1,7 +1,6 @@
 import argparse
 import datetime
 import logging
-import itertools
 from functools import partial
 from pathlib import Path
 import math
@@ -308,6 +307,26 @@ def parse_args():
         help="Minimum learning rate in the lr scheduler."
     )
     parser.add_argument(
+        "--use_ema",
+        action="store_true",
+        help="Whether to use EMA model."
+    )
+    parser.add_argument(
+        "--ema_inv_gamma",
+        type=float,
+        default=1.0
+    )
+    parser.add_argument(
+        "--ema_power",
+        type=float,
+        default=4/5
+    )
+    parser.add_argument(
+        "--ema_max_decay",
+        type=float,
+        default=0.9999
+    )
+    parser.add_argument(
         "--optimizer",
         type=str,
         default="dadan",
@@ -334,7 +353,7 @@ def parse_args():
     parser.add_argument(
         "--adam_weight_decay",
         type=float,
-        default=1e-2,
+        default=0,
         help="Weight decay to use."
     )
     parser.add_argument(
@@ -432,6 +451,23 @@ 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,
@@ -696,6 +732,13 @@ 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,
+        ema_max_decay=args.ema_max_decay,
         sample_batch_size=args.sample_batch_size,
         sample_num_batches=args.sample_batches,
         sample_num_steps=args.sample_steps,
@@ -757,10 +800,7 @@ def main():
             sample_frequency = math.ceil(num_train_epochs * (sample_frequency / args.num_train_steps))
 
         optimizer = create_optimizer(
-            itertools.chain(
-                text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
-                text_encoder.text_model.embeddings.overlay.parameters(),
-            ),
+            text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
             lr=args.learning_rate,
         )