Merge pull request #18160 from Frightera:update_image_dataset

PiperOrigin-RevId: 539682475
keras-team · Jun 12, 2023 · f1eca72 · f1eca72
2 parents 8cabbae + ac4c8ea
commit f1eca72
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/keras/utils/image_dataset.py b/keras/utils/image_dataset.py
@@ -269,8 +269,6 @@ def image_dataset_from_directory(
             interpolation=interpolation,
             crop_to_aspect_ratio=crop_to_aspect_ratio,
         )
-        train_dataset = train_dataset.prefetch(tf.data.AUTOTUNE)
-        val_dataset = val_dataset.prefetch(tf.data.AUTOTUNE)
 
         if batch_size is not None:
             if shuffle:
@@ -286,6 +284,9 @@ def image_dataset_from_directory(
                     buffer_size=1024, seed=seed
                 )
 
+        train_dataset = train_dataset.prefetch(tf.data.AUTOTUNE)
+        val_dataset = val_dataset.prefetch(tf.data.AUTOTUNE)
+
         # Users may need to reference `class_names`.
         train_dataset.class_names = class_names
         val_dataset.class_names = class_names
@@ -314,7 +315,7 @@ def image_dataset_from_directory(
             interpolation=interpolation,
             crop_to_aspect_ratio=crop_to_aspect_ratio,
         )
-        dataset = dataset.prefetch(tf.data.AUTOTUNE)
+
         if batch_size is not None:
             if shuffle:
                 # Shuffle locally at each iteration
@@ -324,6 +325,8 @@ def image_dataset_from_directory(
             if shuffle:
                 dataset = dataset.shuffle(buffer_size=1024, seed=seed)
 
+        dataset = dataset.prefetch(tf.data.AUTOTUNE)
+
         # Users may need to reference `class_names`.
         dataset.class_names = class_names