Keras with TensorFlow Backend: Technical Analysis of Flexible CPU and GPU Usage Control

In deep learning development, Keras, as a high-level neural network API, is often combined with the TensorFlow backend to leverage GPU acceleration. However, practical applications may require flexible switching between CPU and GPU based on task needs, such as during code debugging, resource constraints, or specific algorithm optimization. This article systematically analyzes multiple methods to force Keras to use CPU or GPU, based on high-quality Q&A data from technical communities, and explores their underlying mechanisms.

Environment Variable Control Method

The most straightforward approach is to disable GPU and force CPU usage by setting environment variables. This can be implemented in Python code or executed via the command line. For example, before importing Keras or TensorFlow, add the following code:

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ""

Here, CUDA_VISIBLE_DEVICES is set to an empty string, causing the system to not recognize any GPU devices, thus falling back to CPU. A variant is to set the value to "-1", which works in some Windows systems. This method is simple and efficient, requiring no modification to model code, and is suitable for quick switching scenarios.

TensorFlow Session Configuration Method

For finer control, TensorFlow's ConfigProto can be used to configure sessions. This method allows dynamic setting of CPU and GPU counts, as well as thread parallelism parameters. Example code:

import tensorflow as tf
from keras import backend as K

num_cores = 4
if GPU:
    num_GPU = 1
    num_CPU = 1
if CPU:
    num_CPU = 1
    num_GPU = 0

config = tf.ConfigProto(intra_op_parallelism_threads=num_cores,
                        inter_op_parallelism_threads=num_cores,
                        allow_soft_placement=True,
                        device_count = {'CPU' : num_CPU,
                                        'GPU' : num_GPU})
session = tf.Session(config=config)
K.set_session(session)

Here, intra_op_parallelism_threads and inter_op_parallelism_threads control thread usage for CPU cores, while allow_soft_placement allows operations to automatically fall back to CPU if GPU is unavailable. This method offers greater flexibility but requires installation of tensorflow-gpu and CUDA/cuDNN.

Device Scope Control Method

Another approach is to use TensorFlow device scopes to directly specify the device for operations. For example:

import tensorflow as tf
with tf.device('/gpu:0'):
    # Operations on GPU
    model.fit(X, y, epochs=20)
with tf.device('/cpu:0'):
    # Operations on CPU
    x = tf.placeholder(tf.float32, shape=(None, 20, 64))

This method is suitable for mixed-device scenarios but requires manual management of device allocation, which may increase code complexity.

Summary and Recommendations

The choice of method depends on specific needs: the environment variable method is ideal for quick global switching; the session configuration method provides fine-grained control for high-performance computing; and the device scope method is suitable for optimizing specific operations. In practice, it is recommended to select based on task characteristics and system environment to maximize computational efficiency. Additionally, note that these methods may change with updates to TensorFlow and Keras versions, so referring to official documentation and community discussions is advised.