Add validation split of datasets

library/config_util.py

@@ -85,6 +85,8 @@ class BaseDatasetParams:
   max_token_length: int = None
   resolution: Optional[Tuple[int, int]] = None
   debug_dataset: bool = False
+  validation_seed: Optional[int] = None
+  validation_split: float = 0.0
 
 @dataclass
 class DreamBoothDatasetParams(BaseDatasetParams):
@@ -200,6 +202,8 @@ class ConfigSanitizer:
     "enable_bucket": bool,
     "max_bucket_reso": int,
     "min_bucket_reso": int,
+    "validation_seed": int,
+    "validation_split": float,
     "resolution": functools.partial(__validate_and_convert_scalar_or_twodim.__func__, int),
   }
 
@@ -427,64 +431,89 @@ def generate_dataset_group_by_blueprint(dataset_group_blueprint: DatasetGroupBlu
       dataset_klass = FineTuningDataset
 
     subsets = [subset_klass(**asdict(subset_blueprint.params)) for subset_blueprint in dataset_blueprint.subsets]
-    dataset = dataset_klass(subsets=subsets, **asdict(dataset_blueprint.params))
+    dataset = dataset_klass(subsets=subsets, is_train=True, **asdict(dataset_blueprint.params))
     datasets.append(dataset)
 
-  # print info
-  info = ""
-  for i, dataset in enumerate(datasets):
-    is_dreambooth = isinstance(dataset, DreamBoothDataset)
-    is_controlnet = isinstance(dataset, ControlNetDataset)
-    info += dedent(f"""\
-      [Dataset {i}]
-        batch_size: {dataset.batch_size}
-        resolution: {(dataset.width, dataset.height)}
-        enable_bucket: {dataset.enable_bucket}
-    """)
-
-    if dataset.enable_bucket:
-      info += indent(dedent(f"""\
-        min_bucket_reso: {dataset.min_bucket_reso}
-        max_bucket_reso: {dataset.max_bucket_reso}
-        bucket_reso_steps: {dataset.bucket_reso_steps}
-        bucket_no_upscale: {dataset.bucket_no_upscale}
-      \n"""), "  ")
+  val_datasets:List[Union[DreamBoothDataset, FineTuningDataset, ControlNetDataset]] = []
+  for dataset_blueprint in dataset_group_blueprint.datasets:
+    if dataset_blueprint.params.validation_split <= 0.0:
+      continue
+    if dataset_blueprint.is_controlnet:
+      subset_klass = ControlNetSubset
+      dataset_klass = ControlNetDataset
+    elif dataset_blueprint.is_dreambooth:
+      subset_klass = DreamBoothSubset
+      dataset_klass = DreamBoothDataset
     else:
-      info += "\n"
+      subset_klass = FineTuningSubset
+      dataset_klass = FineTuningDataset
+  
+    subsets = [subset_klass(**asdict(subset_blueprint.params)) for subset_blueprint in dataset_blueprint.subsets]
+    dataset = dataset_klass(subsets=subsets, is_train=False, **asdict(dataset_blueprint.params))
+    val_datasets.append(dataset)
 
-    for j, subset in enumerate(dataset.subsets):
-      info += indent(dedent(f"""\
-        [Subset {j} of Dataset {i}]
-          image_dir: "{subset.image_dir}"
-          image_count: {subset.img_count}
-          num_repeats: {subset.num_repeats}
-          shuffle_caption: {subset.shuffle_caption}
-          keep_tokens: {subset.keep_tokens}
-          caption_dropout_rate: {subset.caption_dropout_rate}
-          caption_dropout_every_n_epoches: {subset.caption_dropout_every_n_epochs}
-          caption_tag_dropout_rate: {subset.caption_tag_dropout_rate}
-          caption_prefix: {subset.caption_prefix}
-          caption_suffix: {subset.caption_suffix}
-          color_aug: {subset.color_aug}
-          flip_aug: {subset.flip_aug}
-          face_crop_aug_range: {subset.face_crop_aug_range}
-          random_crop: {subset.random_crop}
-          token_warmup_min: {subset.token_warmup_min},
-          token_warmup_step: {subset.token_warmup_step},
-      """), "  ")
-
-      if is_dreambooth:
+  # print info
+  def print_info(_datasets):
+    info = ""
+    for i, dataset in enumerate(_datasets):
+      is_dreambooth = isinstance(dataset, DreamBoothDataset)
+      is_controlnet = isinstance(dataset, ControlNetDataset)
+      info += dedent(f"""\
+        [Dataset {i}]
+          batch_size: {dataset.batch_size}
+          resolution: {(dataset.width, dataset.height)}
+          enable_bucket: {dataset.enable_bucket}
+      """)
+  
+      if dataset.enable_bucket:
         info += indent(dedent(f"""\
-          is_reg: {subset.is_reg}
-          class_tokens: {subset.class_tokens}
-          caption_extension: {subset.caption_extension}
-        \n"""), "    ")
-      elif not is_controlnet:
+          min_bucket_reso: {dataset.min_bucket_reso}
+          max_bucket_reso: {dataset.max_bucket_reso}
+          bucket_reso_steps: {dataset.bucket_reso_steps}
+          bucket_no_upscale: {dataset.bucket_no_upscale}
+        \n"""), "  ")
+      else:
+        info += "\n"
+  
+      for j, subset in enumerate(dataset.subsets):
         info += indent(dedent(f"""\
-          metadata_file: {subset.metadata_file}
-        \n"""), "    ")
+          [Subset {j} of Dataset {i}]
+            image_dir: "{subset.image_dir}"
+            image_count: {subset.img_count}
+            num_repeats: {subset.num_repeats}
+            shuffle_caption: {subset.shuffle_caption}
+            keep_tokens: {subset.keep_tokens}
+            caption_dropout_rate: {subset.caption_dropout_rate}
+            caption_dropout_every_n_epoches: {subset.caption_dropout_every_n_epochs}
+            caption_tag_dropout_rate: {subset.caption_tag_dropout_rate}
+            caption_prefix: {subset.caption_prefix}
+            caption_suffix: {subset.caption_suffix}
+            color_aug: {subset.color_aug}
+            flip_aug: {subset.flip_aug}
+            face_crop_aug_range: {subset.face_crop_aug_range}
+            random_crop: {subset.random_crop}
+            token_warmup_min: {subset.token_warmup_min},
+            token_warmup_step: {subset.token_warmup_step},
+        """), "  ")
+  
+        if is_dreambooth:
+          info += indent(dedent(f"""\
+            is_reg: {subset.is_reg}
+            class_tokens: {subset.class_tokens}
+            caption_extension: {subset.caption_extension}
+          \n"""), "    ")
+        elif not is_controlnet:
+          info += indent(dedent(f"""\
+            metadata_file: {subset.metadata_file}
+          \n"""), "    ")
 
-  print(info)
+    print(info)
+
+  print_info(datasets)
+
+  if len(val_datasets) > 0:
+    print("Validation dataset")
+    print_info(val_datasets)
 
   # make buckets first because it determines the length of dataset
   # and set the same seed for all datasets
@@ -494,7 +523,15 @@ def generate_dataset_group_by_blueprint(dataset_group_blueprint: DatasetGroupBlu
     dataset.make_buckets()
     dataset.set_seed(seed)
 
-  return DatasetGroup(datasets)
+  for i, dataset in enumerate(val_datasets):
+    print(f"[Validation Dataset {i}]")
+    dataset.make_buckets()
+    dataset.set_seed(seed)
+
+  return (
+      DatasetGroup(datasets),
+      DatasetGroup(val_datasets) if val_datasets else None
+  )
 
 
 def generate_dreambooth_subsets_config_by_subdirs(train_data_dir: Optional[str] = None, reg_data_dir: Optional[str] = None):