TensorFlow-GPU 2.3.0环境下Keras model.fit未利用XLA_GPU设备的排查及指定GPU训练方法

Alright, let's dig into why your GPU isn't being used during training, even though tf.config.get_visible_devices() shows an XLA_GPU. I'll break this into troubleshooting steps first, then cover how to explicitly target a GPU for training.

Troubleshooting the Root Cause #

First, a quick clarification: The XLA_GPU device you see is an XLA-accelerated virtual placeholder, not your actual physical GPU. If TensorFlow isn't listing a PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU') entry, that's a critical red flag—it means TensorFlow isn't properly detecting your hardware GPU. Here's how to diagnose:

• Verify CUDA/cuDNN compatibility: TensorFlow 2.3.0 requires CUDA 10.1 and cuDNN 7.6.5 exactly. Double-check your versions: run nvcc --version to confirm CUDA, and verify cuDNN files are present in your CUDA library directory. Mismatched versions are the #1 cause of GPU detection failures.

• Check for physical GPU detection: Run this code to confirm if TensorFlow sees your actual GPU:

  print(tf.config.list_physical_devices('GPU'))
  

  If this returns an empty list, your GPU isn't being recognized at all—only the XLA placeholder. Fix the CUDA/cuDNN setup first before moving forward.

• Inspect accidental CPU placement: Did you wrap any part of your model building or training code in with tf.device('/CPU:0'):? Even a single misplaced context manager can force the entire model to run on CPU.

• Validate TensorFlow GPU functionality: Run these standard checks to confirm GPU support works:

  print(tf.test.is_gpu_available(cuda_only=True))
  print(tf.test.gpu_device_name())
  

  If these return False or an empty string, TensorFlow's GPU initialization is broken.

How to Explicitly Specify a GPU for Training #

Once you've confirmed TensorFlow can see your physical GPU, here are 3 reliable ways to force training to use it:

1. Use Environment Variables (Before Running Your Script) #

Restrict TensorFlow to your target GPU using the CUDA_VISIBLE_DEVICES variable. For example, to use the first GPU (index 0):

export CUDA_VISIBLE_DEVICES=0  # Linux/macOS
# Windows Command Prompt:
set CUDA_VISIBLE_DEVICES=0
# Windows PowerShell:
$env:CUDA_VISIBLE_DEVICES=0

This makes TensorFlow only see the specified GPU, so all operations will default to it.

2. Use TensorFlow's Device Configuration API #

Add this code at the very start of your script (before initializing models or data):

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # Set only the first GPU as visible
        tf.config.set_visible_devices(gpus[0], 'GPU')
        # Optional: Enable memory growth to avoid allocating all GPU memory at once
        tf.config.experimental.set_memory_growth(gpus[0], True)
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(f"{len(gpus)} Physical GPUs, {len(logical_gpus)} Logical GPU")
    except RuntimeError as e:
        # This error occurs if devices are initialized before setting visibility
        print(f"Error setting GPU visibility: {e}")

3. Wrap Training in a Device Context #

Explicitly place the model and training loop on your GPU using tf.device:

# Build your model inside the GPU context to ensure variables are placed there
with tf.device('/GPU:0'):
    model = tf.keras.Sequential([
        # Your layers here
    ])
    model.compile(optimizer='adam', loss='mse')

# Run training inside the same context
with tf.device('/GPU:0'):
    model.fit(x_train, y_train, epochs=10, batch_size=32)

A quick note: TensorFlow 2.x's eager execution should auto-place operations on GPU if available, but explicit placement avoids edge cases where the framework falls back to CPU.

内容的提问来源于stack exchange，提问作者KurtRao