Cache vs Prefetch: Choosing tf.data Dataset.cache() or Dataset.prefetch() for Limited Host RAM

Direct Answer

With a host RAM limit of 4 GB and many large JPEGs, do not cache decoded tensors in memory. Use .cache() only when it stores the raw encoded bytes (or a file‑backed cache); otherwise rely on .prefetch() to keep the pipeline overlapped without increasing RAM usage. Placing .cache() before .shuffle() will freeze the shuffle order across epochs; placing it after preserves per‑epoch randomness.

Question 1: Does cache() before shuffle() freeze epoch order?

Yes. When .cache() materializes the dataset, the sequence of elements is fixed. Subsequent .shuffle() works on that fixed sequence, producing the same shuffled order every epoch (unless you set reshuffle_each_iteration=False explicitly). To get a new random order each epoch, place .shuffle() before .cache() or omit .cache() entirely.

Question 2: Memory footprints of cache before vs. after decode

The amount of data stored by .cache() depends on its position relative to the decoding map:

Placement	What is cached	Typical size per image
.cache() → .map(decode)	Compressed JPEG bytes	0.5–2 MB
.map(decode) → .cache()	Decoded uint8 tensor (H×W×3)	10–25× larger than the compressed form

In a 4 GB environment, caching decoded images is likely to hit the OOM threshold unless the dataset is very small. Caching the raw bytes or using a file‑backed cache keeps RAM usage bounded.

Question 3: Diagnostics for a true leak vs. allocator caching

1. Measure RSS (Resident Set Size) with top or htop over several epochs. A plateau after the first pass suggests allocator caching; a continuous rise indicates a leak.
2. Query TensorFlow’s allocator: tf.config.experimental.get_memory_info('CPU:0').current. A stable value across epochs supports caching; a steadily increasing value points to a leak.
3. Optionally run tf.profiler.experimental.start() and inspect the memory timeline for a flat vs. stair‑step pattern.

Recommended Pipeline for 4 GB RAM

# Basic pipeline without in‑memory cache
(ds
  .shuffle(buffer_size=10000, reshuffle_each_iteration=True)
  .map(load_and_decode, num_parallel_calls=tf.data.AUTOTUNE)
  .batch(batch_size)
  .prefetch(tf.data.AUTOTUNE))
# If decoding is expensive and disk is fast, add a file cache
# ds = ds.cache('/tmp/train_cache')  # stores raw bytes on disk

Missing Diagnostic Detail

To fine‑tune the decision between a file‑backed cache and pure prefetching, we need the average decoded image size (in bytes) and the total number of examples. This will let us estimate whether the raw‑byte cache will fit comfortably in RAM or if disk‑based caching is necessary.