How to switch between CPU and GPU usage in Keras with TensorFlow backend?

Backend
1

I’ve got Keras set up with TensorFlow backend and CUDA installed. Sometimes I want to make Keras use the CPU instead of the GPU. Is there a way to do this without having to set up a separate CPU-only TensorFlow in a virtual environment?

With Theano as the backend, you can use flags to control this. But I’m not sure if there are similar options for TensorFlow when using Keras. Does anyone know if it’s possible to switch between CPU and GPU on the fly? And if so, how do you do it?

Here’s a simple example of what I’m trying to achieve:

import keras

# Hypothetical function to set device
def set_keras_device(device):
    # How to implement this?
    pass

# Train on GPU
set_keras_device('gpu')
model = keras.Sequential([
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='categorical_crossentropy')
model.fit(x_train, y_train, epochs=5)

# Switch to CPU for inference
set_keras_device('cpu')
predictions = model.predict(x_test)

Any help would be appreciated!

2

u can use tf.config.set_visible_devices to switch between cpu and gpu. it’s pretty easy:

import tensorflow as tf

def set_device(device):
    if device == 'gpu':
        tf.config.set_visible_devices(tf.config.list_physical_devices('GPU'), 'GPU')
    else:
        tf.config.set_visible_devices([], 'GPU')

# use gpu
set_device('gpu')
# use cpu
set_device('cpu')

this lets u switch on the fly without separate envs

3

I’ve faced this issue before and found a practical solution using TensorFlow’s device management. You can leverage tf.config.set_visible_devices to control which devices are accessible to TensorFlow. Here’s a concise function that does the job:

import tensorflow as tf

def set_keras_device(device):
    if device == 'gpu':
        tf.config.set_visible_devices(tf.config.list_physical_devices('GPU'), 'GPU')
    else:
        tf.config.set_visible_devices([], 'GPU')

# Usage:
set_keras_device('gpu')  # Use GPU
set_keras_device('cpu')  # Use CPU

This approach allows you to switch between CPU and GPU dynamically without separate environments. It’s efficient and works seamlessly with Keras models. Remember to call this function before creating or using your model to ensure the desired device is set.

4

hey there! have u tried using tf.device? it’s pretty neat for switching between cpu and gpu. like this:

import tensorflow as tf

with tf.device('/CPU:0'):
    # do cpu stuff here
with tf.device('/GPU:0'):
    # do gpu stuff here

what do u think? could this work for ur setup? lemme know if u need more info!