TensorFlow
GPU Memory
tf.device
CPU
Machine Learning

Tensorflow allocating GPU memory when using tf.device'/cpu0'

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

TensorFlow can still initialize GPU contexts even when operations are placed on CPU device scopes. Device visibility and runtime initialization behavior differ from per-op placement. If you need strict CPU-only execution, disable GPU visibility explicitly.

A reliable implementation should remain understandable during troubleshooting and upgrades. That requires explicit assumptions, clear boundaries, and verifiable behavior under both normal and failure conditions.

Core Sections

1. Understand device scope versus runtime initialization

A CPU device scope affects operation placement, but TensorFlow may still probe or reserve GPU resources at startup. This surprises many users monitoring GPU memory usage.

python
1import tensorflow as tf
2
3with tf.device('/CPU:0'):
4    x = tf.constant([1.0, 2.0, 3.0])
5    y = x * 2.0
6print(y)

The baseline should be intentionally small and deterministic. A compact first version is easier to test, easier to reason about, and faster to review when teams iterate.

2. Disable GPU visibility for strict CPU runs

Set visible devices to CPU-only before creating tensors or models. This prevents GPU initialization and memory reservation.

python
1import tensorflow as tf
2
3tf.config.set_visible_devices([], 'GPU')
4
5# optional confirmation
6print(tf.config.get_visible_devices())
7
8x = tf.random.uniform((1024, 1024))
9print(tf.reduce_mean(x))

After baseline correctness, harden around edge cases and integration boundaries. Explicit validation, timeout handling, and predictable error semantics make downstream behavior safer.

3. Use memory growth when GPU is needed but limited

If GPU remains enabled, configure memory growth to avoid reserving all memory upfront. This supports shared-GPU development environments more safely.

Operationally, define what success looks like in measurable terms and record baseline metrics before rollout. This makes post-change evaluation objective rather than anecdotal.

Include at least one representative production-like test, one malformed-input test, and one dependency-failure test in CI. Repeatable coverage prevents regressions introduced by dependency changes or refactors.

Keep ownership and escalation paths clear. When incidents happen, responders should know who owns the code path, what logs and metrics to inspect first, and how to execute a safe rollback or fallback mode.

Before release, confirm recovery mechanics in practice. A rollback strategy that is never rehearsed is often too slow under pressure, while a validated recovery workflow can reduce outage impact dramatically.

A complete engineering solution also includes explicit contracts for ownership, inputs, and failure semantics. Document what callers may send, which errors are retriable, and what actions operators should take when dependencies degrade. Clear contracts reduce ambiguity between teams and prevent divergent behavior in different services that rely on the same pattern.

Testing should represent real constraints rather than toy inputs only. Add one production-like scenario, one malformed-input scenario, and one dependency-failure scenario with deterministic assertions. Keep these checks in continuous integration so every change verifies behavior against the same baseline. This practice catches regressions early and reduces the chance of late surprises during rollout.

Observability should be focused and intentional. Emit concise logs for key branch decisions, include request identifiers for traceability, and track metrics tied to user impact such as latency percentiles, error rates, and retry outcomes. Focused telemetry helps teams distinguish application defects from infrastructure instability quickly during incidents.

Before deployment, prepare rollback and fallback options that can be executed quickly. Feature toggles, staged rollout, and a validated reversion workflow significantly reduce operational risk when real traffic reveals assumptions that were not visible in development. Recovery planning in advance is a core reliability practice and should be rehearsed periodically.

Finally, keep runbook notes near the implementation and update them as behavior evolves. Short, current documentation dramatically improves handoffs and lowers on-call resolution time.

Common Pitfalls

  • Assuming CPU placement scope alone disables all GPU runtime behavior.
  • Calling visibility configuration after TensorFlow has already initialized devices.
  • Ignoring mixed library imports that initialize TensorFlow before config code.
  • Interpreting startup GPU probing as full workload GPU execution.
  • Running shared GPU systems without memory growth or resource limits.

Summary

  • CPU scope does not guarantee zero GPU initialization side effects.
  • Disable GPU visibility early for strict CPU-only execution.
  • Configure memory growth when GPUs are shared and still required.
  • Initialize TensorFlow device policy before creating model objects.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.