Keras and TensorFlow expose image augmentation along two independent axes: the API used (preprocessing/augmentation layers vs. tf.image or ImageDataGenerator-style preprocessing) and the timing of execution (offline before training vs. online during it). Choosing the wrong combination wastes disk space, stalls the GPU, or starves the model of variety. This article maps the four resulting configurations onto concrete dataset and compute scenarios, so practitioners preparing image training data can pick the configuration that matches their storage budget, GPU availability, and reproducibility needs.
Keras and TensorFlow expose image augmentation through two distinct APIs. The higher-level option is the Keras preprocessing and augmentation layer family, used inside a keras.Sequential pipeline or applied to a tf.data.Dataset:
layers.Resizing and layers.Rescaling — deterministic geometry and value-range normalization.layers.RandomFlip, layers.RandomRotation, layers.RandomCrop, layers.RandomContrast, layers.RandomZoom — random stochastic transforms that fire only when called in training mode (DataCamp).The lower-level alternative is the tf.image module, which offers fine-grained tensor ops such as tf.image.flip_left_right, tf.image.random_brightness, tf.image.stateless_random_brightness (with an explicit seed tuple), tf.image.central_crop, tf.image.adjust_saturation, and tf.image.rgb_to_grayscale. These ops are usually bundled into a Python function and applied to the training set with Dataset.map(..., num_parallel_calls=AUTOTUNE), giving full control over the graph, seeding, and parallelism (DataCamp).
The second axis is timing:
Dataset.map or as part of the model, generating a fresh random transform per batch or per epoch. Storage stays at the original footprint and variation is effectively unbounded, at the cost of CPU/IO overhead each step (Machine Learning Mastery, Datature).SequentialWhen random augmentation layers are placed inside a keras.Sequential model, they are inactive during Model.evaluate and Model.predict and only fire under Model.fit. Any production inference path must therefore re-apply the deterministic preprocessing (Resizing, Rescaling, and any required deterministic crop) outside the model, otherwise input shape or value range will not match what the trained weights expect (DataCamp).
ImageDataGeneratortf.keras.preprocessing.image.ImageDataGenerator was deprecated in TensorFlow 2.9 (2022). Existing tutorials and codebases still rely on it, but new projects should use Keras preprocessing layers or tf.image ops inside a tf.data pipeline instead (Datature).
Offline augmentation is the simpler of the two workflows to reason about. A standalone script reads each source image once, applies a deterministic or randomly sampled pipeline of tf.image operations (or Keras preprocessing layers, see the DataCamp augmentation guide for the API surface), and writes N augmented variants to disk before Model.fit is ever called. From the trainer's perspective, the augmentation step has already happened — what remains is a static, larger dataset.
The defining traits are listed directly in the comparison table from Machine Learning Mastery's augmentation guide: offline augmentation runs once before training, produces an expanded dataset saved to disk, increases storage cost by the multiplication factor, and exposes the model to a fixed set of variations. Online augmentation, by contrast, produces effectively unbounded variation with no extra storage. That single trade-off — paying disk space in exchange for a frozen, inspectable corpus — is what makes offline augmentation useful or harmful depending on the scenario.
Storage and compute trade-offs. Augmenting 10,000 images at a factor of 10× produces 100,000 files. At typical JPEG sizes that is tens to hundreds of GB, which is still manageable for small projects but becomes painful quickly. What you buy with that storage is amortization: deterministic preprocessing steps such as resize, rescale, and JPEG decode/encode run once instead of every epoch, which matters on intermittent compute such as Kaggle notebooks or spot instances where you do not want to pay CPU time repeatedly for the same transforms.
When offline is the right choice.
Practical concerns for the writer of the offline script.
The trade-off to remember: offline augmentation exchanges disk space and flexibility for reproducibility, inspectability, and amortized compute — a bargain that pays off precisely when the dataset is small, the augmentations are fixed, and the training loop should not be the place where variability is introduced.
Online augmentation generates variation inside the training loop rather than writing extra images to disk. The canonical Keras/TensorFlow pipeline starts with tf.keras.utils.image_dataset_from_directory, which produces a tf.data.Dataset of (image, label) pairs directly from the folder layout. Augmentation is then applied through one of two sub-modes, and the rest of the chain — shuffle, batch, and prefetch — overlaps CPU-side data generation with the GPU's forward and backward pass:
AUTOTUNE = tf.data.AUTOTUNE
train_ds = (
train_ds
.map(augment_fn, num_parallel_calls=AUTOTUNE)
.shuffle(buffer_size)
.batch(batch_size)
.prefetch(AUTOTUNE)
)
The num_parallel_calls=AUTOTUNE argument lets TensorFlow pick a worker count that matches available cores, and the final prefetch(AUTOTUNE) keeps a buffer of ready batches so the GPU never idles waiting for decoded images. As the Shervine blog on Keras data generators notes, the goal of this parallel structure is to keep the bottleneck on the GPU's forward/backward operations, "and not data generation."
The first sub-mode embeds Keras preprocessing/augmentation layers at the head of the model — for example as the first entries of a keras.Sequential. These layers are stochastic only when the model is in training mode: they check an internal training flag, and Keras automatically sets it to False inside Model.evaluate and Model.predict. Per the DataCamp guide on image augmentation, the layers are therefore "inactive during the testing phase" and only fire during Model.fit. If you need augmentation during inference (for example, test-time augmentation), call the model explicitly with model(inputs, training=True).
tf.data.map with tf.image stateless opsThe second sub-mode keeps augmentation outside the model graph and inside a Dataset.map function that calls tf.image operations. For per-sample reproducibility, use the stateless variants such as tf.image.stateless_random_brightness(image, max_delta, seed=seed) and pass an explicit seed tuple of size 2. Different seeds produce different augmentations of the same image, which is useful when you want a specific sample to round-trip identically across runs.
A common misconception is that online augmentation always produces fresh samples each epoch. It only does so when the map function is stochastic. If the map calls deterministic tf.image ops (for example, tf.image.flip_left_right or tf.image.central_crop), the same image always yields the same output, and the dataset sequence is effectively fixed across epochs. Use stochastic ops — Keras augmentation layers or tf.image random/stateless ops — when you need epoch-to-epoch variety.
Online augmentation is the right default when the dataset is large (hundreds of thousands of images or more), when disk is constrained, when variety matters more than bit-for-bit reproducibility, or when each training run benefits from fresh samples.
This configuration fits small datasets (typically under 10k images) where reproducibility and inspectability matter more than raw throughput. A keras.Sequential model containing RandomFlip, RandomRotation, and RandomZoom is applied to the raw images once, the augmented tensors are written to disk (often as NumPy arrays or TFRecords), and a fresh training script loads the expanded dataset with no augmentation layers attached. The pipeline is declarative — it lives inside a saved model graph that can be diffed, re-run, and shared across experiments — and the resulting files can be opened manually to verify what the model actually sees. The cost is storage: the on-disk dataset grows by the augmentation factor, and any change to the augmentation policy requires regenerating the files.
Model.fitThis is the default for most modern Keras classifiers. Layers such as RandomFlip and RandomRotation sit at the top of a Sequential model and execute under Model.fit. It is convenient and integrates with model.save, but it carries a sharp caveat noted in the Keras docs: augmentation layers are inactive outside Model.fit, so Model.evaluate and Model.predict bypass them entirely. The layer graph also cannot be reused outside that model — exporting it for a separate inference graph or a non-Keras consumer requires re-implementing the logic.
tf.image ops, online via tf.data.mapThis is the high-throughput default. A Python function calls stateless ops such as tf.image.stateless_random_brightness and tf.image.random_crop, each receiving an explicit seed tuple (e.g. (i, 0)) that yields per-sample determinism. The function is then attached to the dataset with .map(augment, num_parallel_calls=tf.data.AUTOTUNE), and the pipeline finishes with .batch() and .prefetch(AUTOTUNE) so the GPU is never starved. Variety is effectively unbounded, and storage stays at the original footprint.
tf.image ops, offlineThe same tf.image function used online can be reused offline, often against a fixed seed iterator, to emit TFRecord shards for distribution. This is the right choice when many training jobs will consume a frozen, versioned dataset, or when consumers outside TensorFlow need the augmented files.
In practice, a hybrid is the production sweet spot: resize and rescale offline into a cached, resized dataset (these ops are deterministic and expensive to repeat), then run stochastic flips, rotations, and color jitter online through tf.data.map. This combination keeps storage bounded, leaves the GPU fed, and makes the stochastic part trivially reproducible via stateless seeds.
Offline augmentation removes per-epoch CPU work, but it multiplies both disk usage and training-time I/O by the augmentation factor. A 50,000-image dataset expanded tenfold occupies roughly 500,000 images on disk, and every epoch then has to stream the enlarged corpus. Online augmentation keeps the on-disk footprint at the original size, but moves the cost to per-epoch CPU work. The bottleneck shifts from storage to data loading, which is mitigated by setting num_parallel_calls=tf.data.AUTOTUNE on Dataset.map and following it with prefetch(tf.data.AUTOTUNE) so the GPU rarely waits for the next batch. When CPU augmentation cannot keep up, the GPU starves and step time is dominated by data rather than compute.
training flagIn current Keras, augmentation layers are only active when training=True. They are silently bypassed during Model.evaluate and Model.predict, so a saved model that relies on layer-level augmentation will not apply those transforms at inference. Two remedies preserve correctness: either re-apply the deterministic preprocessing steps (resize, rescale) inside the inference graph and keep only the stochastic augmentation in the training-only layers, or wrap the saved model with a preprocessing layer so it always runs first.
ImageDataGenerator deprecationtf.keras.preprocessing.image.ImageDataGenerator was deprecated in TensorFlow 2.9 (2022). New projects should migrate to tf.keras.layers preprocessing layers for model-integrated augmentation or to tf.data pipelines that call tf.image ops for batch-level augmentation. The legacy API is still widely referenced in older tutorials, but it lacks bounding-box and mask transforms and is slower than specialized libraries.
TTA is the one exception to the rule that augmentation is training-only: applying 2–5 inference passes (typically horizontal flip plus 2–3 scale factors) and averaging the softmax outputs reliably improves accuracy by about 1–3 percent at a 3–5x latency cost. It is appropriate when accuracy matters more than throughput, such as medical diagnosis or competition submissions, but should not be used for validation or test-set metrics that need to reflect real-world performance.
When plotting augmented float tensors with plt.imshow, matplotlib may warn about clipping or display a washed-out image because it expects 8-bit input. Divide the array by 255 before calling imshow so values fall in the expected [0, 1] range.