Improved aspect ratio bucketing

1 files changed, 15 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index 323ef10..eb0b8b6 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -143,7 +143,7 @@ def parse_args():
         "--num_buckets",
         type=int,
         default=4,
-        help="Number of aspect ratio buckets in either direction (adds 64 pixels per step).",
+        help="Number of aspect ratio buckets in either direction.",
     )
     parser.add_argument(
         "--progressive_buckets",
@@ -151,6 +151,18 @@ def parse_args():
         help="Include images in smaller buckets as well.",
     )
     parser.add_argument(
+        "--bucket_step_size",
+        type=int,
+        default=64,
+        help="Step size between buckets.",
+    )
+    parser.add_argument(
+        "--bucket_max_pixels",
+        type=int,
+        default=None,
+        help="Maximum pixels per bucket.",
+    )
+    parser.add_argument(
         "--tag_dropout",
         type=float,
         default=0.1,
@@ -718,6 +730,8 @@ def main():
         size=args.resolution,
         num_buckets=args.num_buckets,
         progressive_buckets=args.progressive_buckets,
+        bucket_step_size=args.bucket_step_size,
+        bucket_max_pixels=args.bucket_max_pixels,
         dropout=args.tag_dropout,
         template_key=args.train_data_template,
         valid_set_size=args.valid_set_size,