feat: use specific val set for oxford

Author: AshishKumar4

flaxdiff/data/sources/images.py

@@ -79,6 +79,19 @@ def load_labels(sample):
 # TFDS Image Source
 # ----------------------------------------------------------------------------------
 
+def get_oxford_valset(text_encoder):
+    # Construct a validation set by the prompts for consistency
+    val_prompts = ['water tulip', ' a water lily', ' a water lily', ' a photo of a rose', ' a photo of a rose', ' a water lily', ' a water lily', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a marigold', ' a water lily', ' a photo of a sunflower', ' a photo of a lotus', ' columbine', ' columbine', ' an orchid', ' an orchid', ' an orchid', ' a water lily', ' a water lily', ' a water lily', ' columbine', ' columbine', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a lotus', ' a photo of a lotus', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a rose', ' a photo of a rose', ' a photo of a rose', ' orange dahlia', ' orange dahlia', ' a lenten rose', ' a lenten rose', ' a water lily', ' a water lily', ' a water lily', ' a water lily', ' an orchid', ' an orchid', ' an orchid', ' hard-leaved pocket orchid', ' bird of paradise', ' bird of paradise', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a globe-flower', ' a photo of a globe-flower', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a ruby-lipped cattleya', ' a photo of a ruby-lipped cattleya', ' a photo of a lovely rose', ' a water lily', ' a osteospermum', ' a osteospermum', ' a water lily', ' a water lily', ' a water lily', ' a red rose', ' a red rose']
+    val_prompts *= 10
+
+    def get_val_dataset(batch_size=128):
+        for i in range(0, len(val_prompts), batch_size):
+            prompts = val_prompts[i:i + batch_size]
+            tokens = text_encoder.tokenize(prompts)
+            yield {"text": tokens}
+
+    return get_val_dataset, len(val_prompts)
+    
 class ImageTFDSSource(DataSource):
     """Data source for TensorFlow Datasets (TFDS) image datasets."""
 

training.py

@@ -575,17 +575,11 @@ def main(args):
         print("Distributed Training enabled")
     print(f"Training on {CONFIG['dataset']['name']} dataset with {batches} samples")
 
-    # Construct a validation set by the prompts
-    val_prompts = ['water tulip', ' a water lily', ' a water lily', ' a photo of a rose', ' a photo of a rose', ' a water lily', ' a water lily', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a marigold', ' a water lily', ' a photo of a sunflower', ' a photo of a lotus', ' columbine', ' columbine', ' an orchid', ' an orchid', ' an orchid', ' a water lily', ' a water lily', ' a water lily', ' columbine', ' columbine', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a lotus', ' a photo of a lotus', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a rose', ' a photo of a rose', ' a photo of a rose', ' orange dahlia', ' orange dahlia', ' a lenten rose', ' a lenten rose', ' a water lily', ' a water lily', ' a water lily', ' a water lily', ' an orchid', ' an orchid', ' an orchid', ' hard-leaved pocket orchid', ' bird of paradise', ' bird of paradise', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a globe-flower', ' a photo of a globe-flower', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a ruby-lipped cattleya', ' a photo of a ruby-lipped cattleya', ' a photo of a lovely rose', ' a water lily', ' a osteospermum', ' a osteospermum', ' a water lily', ' a water lily', ' a water lily', ' a red rose', ' a red rose']
-
-    def get_val_dataset(batch_size=8):
-        for i in range(0, len(val_prompts), batch_size):
-            prompts = val_prompts[i:i + batch_size]
-            tokens = text_encoder.tokenize(prompts)
-            yield {"text": tokens}
-
-    data['test'] = get_val_dataset
-    data['test_len'] = len(val_prompts)
+    if dataset_name == 'oxford_flowers102':
+        from flaxdiff.data.sources.images import get_oxford_valset
+        val, val_len = get_oxford_valset(text_encoder)
+        data['val_len'] = len(val)
+        data['val'] = val
 
     final_state = trainer.fit(
         data,