Training script improvements

1 files changed, 38 insertions, 5 deletions

diff --git a/train_ti.py b/train_ti.py
index b1f6a49..6aa4007 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -93,6 +93,18 @@ def parse_args():
         help="The directory where class images will be saved.",
     )
     parser.add_argument(
+        "--exclude_keywords",
+        type=str,
+        nargs='*',
+        help="Skip dataset items containing a listed keyword.",
+    )
+    parser.add_argument(
+        "--exclude_modes",
+        type=str,
+        nargs='*',
+        help="Exclude all items with a listed mode.",
+    )
+    parser.add_argument(
         "--repeats",
         type=int,
         default=1,
@@ -120,7 +132,8 @@ def parse_args():
         "--seed",
         type=int,
         default=None,
-        help="A seed for reproducible training.")
+        help="A seed for reproducible training."
+    )
     parser.add_argument(
         "--resolution",
         type=int,
@@ -356,6 +369,12 @@ def parse_args():
     if len(args.placeholder_token) != len(args.initializer_token):
         raise ValueError("You must specify --placeholder_token")
 
+    if isinstance(args.exclude_keywords, str):
+        args.exclude_keywords = [args.exclude_keywords]
+
+    if isinstance(args.exclude_modes, str):
+        args.exclude_modes = [args.exclude_modes]
+
     if args.output_dir is None:
         raise ValueError("You must specify --output_dir")
 
@@ -576,11 +595,22 @@ def main():
         weight_dtype = torch.bfloat16
 
     def keyword_filter(item: CSVDataItem):
-        return any(
+        cond1 = any(
             keyword in part
             for keyword in args.placeholder_token
             for part in item.prompt
         )
+        cond2 = args.exclude_keywords is None or not any(
+            keyword in part
+            for keyword in args.exclude_keywords
+            for part in item.prompt
+        )
+        cond3 = args.mode is None or args.mode in item.mode
+        cond4 = args.exclude_modes is None or not any(
+            mode in item.mode
+            for mode in args.exclude_modes
+        )
+        return cond1 and cond2 and cond3 and cond4
 
     def collate_fn(examples):
         prompts = [example["prompts"] for example in examples]
@@ -617,7 +647,6 @@ def main():
         num_class_images=args.num_class_images,
         size=args.resolution,
         repeats=args.repeats,
-        mode=args.mode,
         dropout=args.tag_dropout,
         center_crop=args.center_crop,
         template_key=args.train_data_template,
@@ -769,7 +798,7 @@ def main():
             target, target_prior = torch.chunk(target, 2, dim=0)
 
             # Compute instance loss
-            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none").mean([1, 2, 3]).mean()
+            loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
 
             # Compute prior loss
             prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="mean")
@@ -785,7 +814,7 @@ def main():
 
     if args.find_lr:
         lr_finder = LRFinder(accelerator, text_encoder, optimizer, train_dataloader, val_dataloader, loop)
-        lr_finder.run(min_lr=1e-6, num_train_batches=4)
+        lr_finder.run(min_lr=1e-6, num_train_batches=1)
 
         plt.savefig(basepath.joinpath("lr.png"))
         plt.close()
@@ -798,6 +827,10 @@ def main():
         config = vars(args).copy()
         config["initializer_token"] = " ".join(config["initializer_token"])
         config["placeholder_token"] = " ".join(config["placeholder_token"])
+        if config["exclude_modes"] is not None:
+            config["exclude_modes"] = " ".join(config["exclude_modes"])
+        if config["exclude_keywords"] is not None:
+            config["exclude_keywords"] = " ".join(config["exclude_keywords"])
         accelerator.init_trackers("textual_inversion", config=config)
 
     # Train!