PR #18160: Use prefetch() after batching // image_dataset.py

Imported from GitHub PR #18160

I think it is better use `prefetch()` after batching. In a pipeline we would want the next batch ready (t+1) while processing current batch (t).
Copybara import of the project:

--
ac4c8ea by Kaan Bıçakcı <[email protected]>:

Move prefetch() to end

Merging this change closes #18160

FUTURE_COPYBARA_INTEGRATE_REVIEW=#18160 from Frightera:update_image_dataset ac4c8ea
PiperOrigin-RevId: 539677136

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