TensorFlow GPU – Basic Operations and Multi-GPU Setup

If you are an ML practitioner looking to use TensorFlow for NVIDIA GPU, this guide walks you through a clean install and a reliable verification. Together, we will install TensorFlow the recommended way, confirm the runtime sees your GPU and prove real GPU execution before scaling training across multiple GPUs.

NOTE: If you do not want to manage drivers and CUDA, you can run the same checks on prebuilt GPU instances by AceCloud and focus on training behavior.

Which Installation Path Should You Choose for Your OS?

Choosing the right path first prevents most “GPU not found” outcomes, because TensorFlow support differs by platform. Here are the four recommended ways:

1. Linux and Ubuntu (recommended): Create a venv, upgrade pip, install tensorflow[and-cuda] and then verify GPU detection in one Python line.
2. Windows: TensorFlow GPU on native Windows stops after TensorFlow 2.10, therefore you should use WSL2 for modern GPU support.
3. macOS: The official pip guide states there is no official GPU support for macOS, therefore plan for CPU-only in that workflow.
4. Version targets to keep you out of trouble: Use Linux driver ≥ 525.60.13, WSL driver ≥ 528.33, CUDA 12.3 and cuDNN 8.9.7 as your baseline

Prerequisites Before Enabling GPU Acceleration

Prerequisites matter because TensorFlow’s GPU runtime expects a compatible NVIDIA driver stack and supported GPU architecture.

1. Supported NVIDIA GPU and driver

First, you should confirm your GPU is in a supported CUDA architecture family as older architectures can fail at load time. TensorFlow’s pip guide lists supported NVIDIA CUDA architectures including 3.5, 5.0, 6.0, 7.0, 7.5, 8.0 and higher, therefore older GPUs can require a source build.

Use this snippet to list GPUs, print useful details and enable memory growth, because it gives you immediate signals without misleading output.

import tensorflow as tf 
 
gpus = tf.config.list_physical_devices("GPU") 
print("GPUs:", gpus) 
 
for gpu in gpus: 
    tf.config.experimental.set_memory_growth(gpu, True) 
    details = tf.config.experimental.get_device_details(gpu) 
    print(f"Device: {gpu.name}") 
    print("Details:", details) 
 
if not gpus: 
    raise SystemExit("No GPU found by TensorFlow. Check driver, WSL2 and CUDA requirements.") 

2. Versions of CUDA and cuDNN

Secondly, you should align to the versions on the TensorFlow install page, because mismatches are a top cause of empty GPU lists and runtime load errors.

Compatibility baseline from the TensorFlow pip guide:

If version matching is slowing you down, a managed GPU stack can reduce time spent reconciling driver and library combinations across machines.

3. Python and packaging tools to use

Next, you must use a clean venv and pip as TensorFlow is officially released to PyPI and the install guide warns against conda installs. The pip guide lists Python 3.9–3.12, therefore you should confirm your interpreter version before installing anything.

How to Install TensorFlow with GPU support on Linux or WSL2?

This section follows the pip guide’s GPU route as it is the fastest supported path for Linux and WSL2 today.

Step 1: Install the NVIDIA driver and confirm it works

We suggest you run nvidia-smi first, because TensorFlow cannot use a GPU that the driver cannot see.

nvidia-smi

Step 2: Create and activate a virtual environment

Next, you should isolate TensorFlow in a venv. This prevents CUDA library conflicts with other Python environments on the same host.

python3 -m venv tf
source tf/bin/activate

Step 3: Upgrade pip

Then, you will have to upgrade pip before installing TensorFlow as the guide calls out pip version requirements and modern wheel support.

pip install –upgrade pip

Step 4: Install TensorFlow with CUDA extras

You should install the GPU build using the CUDA extras, because the pip guide documents this as the supported GPU install command.

pip install “tensorflow[and-cuda]”

Avoid installing TensorFlow with conda as the guide warns it may not have the latest stable version and pip is the official release channel.

Also Read: CUDA cores vs Tensor cores: Choosing the Right GPU for Machine Learning

Step 5: Verify GPU detection

Finally, you should verify with list_physical_devices(“GPU”) as it is the simplest positive signal that TensorFlow can enumerate your CUDA devices.

python3 -c “import tensorflow as tf; print(tf.config.list_physical_devices(‘GPU’))”

NOTE: If you want to validate the same workflow without local driver maintenance, you can run these identical steps on a cloud GPU instance with published GPU plans.

When Should You Install CUDA and cuDNN Manually?

Manual CUDA and cuDNN installs are an advanced path, because they increase the number of version pairs you must keep consistent across upgrades.

To be specific, you should consider manual installs when you need a nonstandard CUDA version with a strict enterprise baseline or a custom TensorFlow build for unsupported architectures.

In those cases, the TensorFlow install page remains your source of truth for supported combinations. This is because it is updated alongside TensorFlow releases.

How to Verify TensorFlow is Using the GPU?

Such verification is highly critical as you can still fall back to CPU execution through placement or build issues. Here are the best ways to verify if TensorFlow is using the GPU:

1. GPU presence check

Since an empty list means TensorFlow cannot even initialize the GPU runtime, we suggest you start with enumeration.

import tensorflow as tf
print(tf.config.list_physical_devices(“GPU”))

2. A correct, non-misleading micro-benchmark

You can even force result materialization before stopping the timer as GPU kernels can execute asynchronously relative to Python timing.

import time 
import tensorflow as tf 
 
device = "/GPU:0" if tf.config.list_physical_devices("GPU") else "/CPU:0" 
print("Using:", device) 
 
with tf.device(device): 
    a = tf.random.uniform([1000, 1000]) 
    b = tf.random.uniform([1000, 1000]) 
 
    # Warm-up to pay one-time setup costs up front. 
    for _ in range(10): 
        _ = tf.linalg.matmul(a, b).numpy() 
 
    start = time.perf_counter() 
    out = tf.linalg.matmul(a, b) 
    _ = out.numpy()  # Forces completion for timing. 
    end = time.perf_counter() 
 
print(f"MatMul wall time: {(end - start) * 1000:.2f} ms") 

3. Device placement logging(optional)

Enable placement logging when results look suspicious, since it shows whether key ops land on GPU or CPU.

import tensorflow as tf
tf.debugging.set_log_device_placement(True)

How to Run a Minimal Keras Model on a GPU Without Breaking Code?

A minimal end-to-end Keras example helps because it exercises data input, model creation, compilation and training, which is where hidden CPU fallbacks appear.

Minimal runnable TF2 example

You can run this as a single file as it includes imports, dataset, model and a short fit call with a visible GPU check.

import tensorflow as tf 
 
print("GPUs:", tf.config.list_physical_devices("GPU")) 
 
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data() 
x_train = (x_train / 255.0).astype("float32") 
x_train = x_train[..., None]  # Add channel dimension. 
 
ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)) 
ds = ds.shuffle(10_000).batch(256).prefetch(tf.data.AUTOTUNE) 
 
model = tf.keras.Sequential([ 
    tf.keras.layers.Input(shape=(28, 28, 1)), 
    tf.keras.layers.Conv2D(32, 3, activation="relu"), 
    tf.keras.layers.MaxPool2D(), 
    tf.keras.layers.Flatten(), 
    tf.keras.layers.Dense(128, activation="relu"), 
    tf.keras.layers.Dense(10) 
]) 
 
model.compile( 
    optimizer="adam", 
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), 
    metrics=["accuracy"] 
) 
 
model.fit(ds, epochs=2) 

A correct custom training loop with GradientTape

In our opinion, you should use a custom loop when you need manual control. This makes sense as it makes step boundaries explicit and simplifies debugging of loss scaling or gradient issues.

import tensorflow as tf 
 
(x_train, y_train), _ = tf.keras.datasets.fashion_mnist.load_data() 
x_train = (x_train / 255.0).astype("float32") 
 
ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)) 
ds = ds.shuffle(10_000).batch(256).prefetch(tf.data.AUTOTUNE) 
 
model = tf.keras.Sequential([ 
    tf.keras.layers.Input(shape=(28, 28)), 
    tf.keras.layers.Flatten(), 
    tf.keras.layers.Dense(256, activation="relu"), 
    tf.keras.layers.Dense(10) 
]) 
 
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) 
optimizer = tf.keras.optimizers.Adam() 
 
@tf.function 
def train_step(x, y): 
    with tf.GradientTape() as tape: 
        logits = model(x, training=True) 
        loss = loss_fn(y, logits) 
    grads = tape.gradient(loss, model.trainable_variables) 
    optimizer.apply_gradients(zip(grads, model.trainable_variables)) 
    return loss 
 
for epoch in range(2): 
    for step, (x, y) in enumerate(ds): 
        loss = train_step(x, y) 
        if step % 50 == 0: 
            print(f"epoch={epoch} step={step} loss={loss.numpy():.4f}") 

If training works locally, multi-GPU is a natural next step. We say that with conviction since data parallelism often provides a straightforward throughput increase for larger batches.

How to Control GPU Memory Growth and Common Performance Factors?

Memory and throughput tuning matters because default behaviors can reserve large GPU memory blocks and hide input bottlenecks.

1. Enable memory growth

You should enable memory growth in shared GPU hosts as it reduces aggressive pre-allocation and improves coexistence with other processes.

import tensorflow as tf 
 
gpus = tf.config.list_physical_devices("GPU") 
for gpu in gpus: 
    tf.config.experimental.set_memory_growth(gpu, True) 
print("GPUs:", gpus) 

2. Use mixed precision when the GPU supports it

You can enable mixed precision on modern NVIDIA GPUs. This can reduce memory pressure and improve training throughput when numerically stable.

import tensorflow as tf
tf.keras.mixed_precision.set_global_policy("mixed_float16")

3. Reduce input stalls with tf.data

You should prefetch and batch efficiently as a slow input pipeline can leave GPUs idle even when your model code is correct. For deeper diagnosis, TensorFlow Profiler can show GPU underutilization causes, because it connects step time to input, compute and kernel scheduling.

What are the Most Common TensorFlow GPU Setup Issues and Fixes?

Most GPU setup failures fall into a few patterns. Therefore, a simple decision flow saves time compared to random reinstall cycles. Here’s the decision tree to follow:

1) tf.config.list_physical_devices(“GPU”) returns [].

First, you should confirm you are on Linux or WSL2, because native Windows GPU support ends after TensorFlow 2.10. Then, verify driver minimums since Linux needs ≥ 525.60.13 and WSL2 needs ≥ 528.33 in the current pip guide.

2) You see “device kernel image is invalid.”

You should check your GPU’s CUDA architecture. This is because TensorFlow wheels only include PTX for the latest supported architecture and older GPUs may require a source build.

3) Training fails with OOM or fragmentation symptoms.

For this, you will have to enable memory growth and reduce batch size. This is suggested since that combination lowers peak allocation pressure and prevents immediate allocation failure.

4) GPU exists, yet the performance looks like CPU.

We recommend you use a micro-benchmark that forces result materialization as asynchronous execution can make naive timers report unrealistic values.

5) You are stuck in version mismatch loops.

Here, you should align to CUDA 12.3 and cuDNN 8.9.7 from the pip guide. We recommend it since that baseline matches the supported wheel expectations.

Best Practices to Follow for TensorFlow GPU

GPU with TensorFlow offers a significant performance boost over its CPU counterpart. However, to get the most out of your GPU, optimize your code and monitor your usage.

Optimizing Performance with TensorFlow GPU

Use TensorFlow Operations Optimized for GPUs – TensorFlow has operations optimized for GPUs, such as matrix multiplication and convolution. By using these operations, you can take full advantage of a GPU’s capabilities and accelerate your computations.

• Batch Your Data: Batching your data means processing multiple inputs at once, which can help to minimize the time spent transferring data between the CPU and GPU. This technique can significantly speed up your training time and reduce GPU memory usage.
• Use Data Augmentation: Data augmentation is the process of generating new training data by applying random transformations to existing data. This technique can help improve your model’s generalization and increase the efficiency of your training process by reducing the need for new data.
• Use Mixed Precision Training: Mixed precision training is the process of using both single-precision and half-precision floating-point data types in your training process. With half-precision data types, you can reduce the memory usage of your model and speed up your training process.

Monitoring GPU usage

• Use GPU Profiling Tools: TensorFlow provides profiling tools to monitor the performance of GPU during training. These tools can identify performance bottlenecks and optimize the code for maximum efficiency.
• Monitor GPU Memory Usage: Monitoring GPU memory usage is essential to avoid running out of memory during training. TensorFlow’s built-in tools monitor memory usage and adjust batch sizes or model architecture (if necessary).

Troubleshooting Common Issues

• Out-of-Memory Errors: Out-of-memory errors occur when the GPU does not have enough memory to process the data. To fix this issue, you can try reducing your batch size or using a smaller model architecture.
• GPU Driver Issues: GPU driver issues can cause TensorFlow to crash or perform poorly. To avoid these issues, it is essential to keep your GPU drivers up to date.
• Incompatible Hardware: Not all GPUs are compatible with TensorFlow and some may not provide the expected performance benefits. Before investing in a new GPU, make sure to check TensorFlow’s hardware compatibility list.

Achieve TensorFlow GPU Virtualization with AceCloud

There you go. You now know how TensorFlow GPU can significantly enhance the performance of machine learning models. But setting up TensorFlow GPUs shouldn’t be this complicated.

With AceCloud, you can accelerate your workloads while using the power of GPU servers. We offer a range of GPU server configurations to provide the best possible performance for TensorFlow workloads.

AceCloud’s easy-to-use interface and reliable infrastructure allow users to quickly set up and deploy their TensorFlow models on the cloud GPU server of their choice. Don’t wait! Use your free consultation session to connect with our cloud GPU experts and accelerate your machine learning workload today.

Author Popular Post

Jason Karlin

author

Industry veteran with over 10 years of experience architecting and managing GPU-powered cloud solutions. Specializes in enabling scalable AI/ML and HPC workloads for enterprise and research applications. Former lead solutions architect for top-tier cloud providers and startups in the AI infrastructure space.

• CUDA cores vs Tensor cores: Choosing the Right GPU for Machine Learning  
• GPUs in Healthcare: Transforming Diagnoses and Treatment with Cutting-Edge Technology
• Sarvam AI vs ChatGPT, Gemini, Krutrim and DeepSeek: Best AI for India?
• Best GPUs For Deep Learning [2025 Updated List]
• Why GPUs for Deep Learning? A Complete Explanation