Added vector dropout

1 files changed, 19 insertions, 5 deletions

diff --git a/train_ti.py b/train_ti.py
index 102c0fa..cacbbc7 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -155,6 +155,12 @@ def parse_args():
         help="Tag dropout probability.",
     )
     parser.add_argument(
+        "--vector_dropout",
+        type=int,
+        default=0.1,
+        help="Vector dropout probability.",
+    )
+    parser.add_argument(
         "--vector_shuffle",
         type=str,
         default="auto",
@@ -526,6 +532,8 @@ def main():
         tokenizer = MultiCLIPTokenizer.from_pretrained(args.tokenizer_name)
     elif args.pretrained_model_name_or_path:
         tokenizer = MultiCLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder='tokenizer')
+    tokenizer.set_use_vector_shuffle(args.vector_shuffle)
+    tokenizer.set_dropout(args.vector_dropout)
 
     # Load models and create wrapper for stable diffusion
     text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder='text_encoder')
@@ -777,6 +785,12 @@ def main():
     num_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
     val_steps = num_val_steps_per_epoch * num_epochs
 
+    def on_train():
+        tokenizer.train()
+
+    def on_eval():
+        tokenizer.eval()
+
     def loop(batch):
         # Convert images to latent space
         latents = vae.encode(batch["pixel_values"]).latent_dist.sample().detach()
@@ -850,8 +864,8 @@ def main():
             train_dataloader,
             val_dataloader,
             loop,
-            on_train=lambda: tokenizer.set_use_vector_shuffle(args.vector_shuffle),
-            on_eval=lambda: tokenizer.set_use_vector_shuffle(False)
+            on_train=on_train,
+            on_eval=on_eval,
         )
         lr_finder.run(end_lr=1e2)
 
@@ -903,7 +917,7 @@ def main():
         disable=not accelerator.is_local_main_process,
         dynamic_ncols=True
     )
-    local_progress_bar.set_description("Epoch X / Y")
+    local_progress_bar.set_description(f"Epoch 1 / {num_epochs}")
 
     global_progress_bar = tqdm(
         range(args.max_train_steps + val_steps),
@@ -922,7 +936,7 @@ def main():
             local_progress_bar.reset()
 
             text_encoder.train()
-            tokenizer.set_use_vector_shuffle(args.vector_shuffle)
+            on_train()
 
             for step, batch in enumerate(train_dataloader):
                 with accelerator.accumulate(text_encoder):
@@ -963,7 +977,7 @@ def main():
             accelerator.wait_for_everyone()
 
             text_encoder.eval()
-            tokenizer.set_use_vector_shuffle(False)
+            on_eval()
 
             cur_loss_val = AverageMeter()
             cur_acc_val = AverageMeter()