PMLE Scaling Prototypes into ML Models Practice Question

You are training a TensorFlow model on Vertex AI using a custom container with a single Tesla T4 GPU. You notice that training is slower than expected, and GPU utilization is consistently below 20%. Profiling shows that the input pipeline is the bottleneck. Which change should you make to improve GPU utilization?

⚠ Common exam trap

The trap here is assuming that a faster GPU or more memory will fix low GPU utilization, when the issue is actually data starvation from an inefficient input pipeline.

Answer choices

Why each option matters

Answer the question above first, then reveal the full breakdown to understand why each option is right or wrong.

Correct answer & explanation

✓

Use tf.data with prefetching and parallel data extraction to overlap data loading with GPU computation.

When the input pipeline is the bottleneck, the GPU waits for data. Using tf.data with prefetching and parallel extraction overlaps CPU preprocessing with GPU computation, ensuring the GPU is continuously fed. This is the standard TensorFlow optimization for such scenarios and directly improves GPU utilization.

Answer analysis

Option-by-option breakdown

For each option: why learners choose it and why it is or isn't the right answer here.

✗
Switch to a larger GPU instance with more memory to reduce data loading overhead.
Why it's wrong here
Upgrading to a larger GPU does not address the input pipeline bottleneck. The GPU is already underutilized, so more GPU memory or compute will not help; the CPU is struggling to feed data fast enough. This approach increases cost without solving the root cause, and GPU utilization may remain low.
✗
Move the dataset to a local SSD on the training VM to reduce I/O latency.
Why it's wrong here
While local SSD can reduce I/O latency, the bottleneck is likely due to single-threaded data loading and preprocessing, not disk speed. Without parallelizing data extraction and prefetching, the GPU will still wait for the CPU to prepare batches. Using tf.data with parallelization is a more direct and effective solution.
✗
Increase the batch size to fully utilize GPU memory.
Why it's wrong here
Increasing batch size may improve GPU utilization if the GPU is underutilized due to small batches, but the profiling indicates the input pipeline is the bottleneck. Larger batches would increase the load on the input pipeline, potentially worsening the bottleneck. Without fixing data loading, GPU utilization may remain low or even decrease due to increased waiting time.
✓
Use tf.data with prefetching and parallel data extraction to overlap data loading with GPU computation.
Why this is correct
The tf.data API with prefetching and parallel extraction allows data preprocessing to occur on CPU while the GPU computes on the previous batch, effectively overlapping I/O and compute. This directly addresses the input pipeline bottleneck, increasing GPU utilization and reducing training time. It is the recommended approach for optimizing input pipelines in TensorFlow.

Go deeper

PMLE Scaling Prototypes into ML Models Practice Question

Why each option matters

Option-by-option breakdown

Related to this question