TopolearnKnowing where ops run, and stopping TF from grabbing all GPU memory, avoids the most common setup problems.
TensorFlow places operations automatically: if a GPU is visible and the op has a GPU kernel, it runs there. Check what TF sees with tf.config.list_physical_devices("GPU").
Installation notes for current releases:
- On Linux,
pip install "tensorflow[and-cuda]"pulls matching CUDA and cuDNN libraries. - Native Windows GPU support ended after TF 2.10; use WSL2.
- On Apple silicon, the
tensorflow-metalplugin provides GPU acceleration.
By default TensorFlow reserves nearly all GPU memory when it first initialises a GPU. That is efficient for one job but blocks other processes. Call tf.config.experimental.set_memory_growth(gpu, True) at program start (before any op runs) to allocate on demand, or set a hard limit with a logical device configuration.
Force placement with a context manager: with tf.device("/CPU:0"): or "/GPU:1". A tensor's .device attribute shows where it was produced, and tf.debugging.set_log_device_placement(True) logs every op's device.
Data-movement costs apply as in any GPU framework: keep the input pipeline fast (see tf.data performance) and avoid pulling values back to Python with .numpy() inside tight loops.
Example
import tensorflow as tf
gpus = tf.config.list_physical_devices("GPU")
print(gpus) # [] on a CPU-only machine
for gpu in gpus: # must run before the GPU is used
tf.config.experimental.set_memory_growth(gpu, True)
a = tf.random.normal([1000, 1000])
print(a.device) # .../device:GPU:0 if a GPU exists, else .../device:CPU:0
with tf.device("/CPU:0"):
b = tf.matmul(a, a) # pinned to the CPU
print(b.device.endswith("CPU:0")) # TrueRemember
- Placement is automatic; check with
list_physical_devices - Enable memory growth before any GPU op runs
- Pin ops with
tf.device; inspect.device
Watch out for
- Calling
set_memory_growthafter tensors were created, which raises a RuntimeError.