Достаточно отдельный способ сделать это - использовать
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)
Здесь с помощью booleans GPUи CPUмы указываем, хотим ли мы запускать наш код с помощью графического процессора или процессора, жестко определяя количество графических процессоров и процессоров, к которым разрешен доступ сеансу Tensorflow. Переменные num_GPUи num_CPUопределяют это значение. num_coresзатем устанавливает количество ядер ЦП, доступных для использования с помощью intra_op_parallelism_threadsи inter_op_parallelism_threads.
В intra_op_parallelism_threadsдиктует переменное число потоков параллельной работы в одном узле в графе вычисления допускается использование (внутри). В то время как inter_ops_parallelism_threadsпеременная определяет количество потоков, доступных для параллельных операций через узлы графа вычислений (inter).
allow_soft_placement позволяет запускать операции на ЦП, если выполняется любой из следующих критериев:
нет реализации на GPU для операции
нет известных или зарегистрированных устройств GPU
необходимо совмещать с другими входами ЦП
Все это выполняется в конструкторе моего класса перед любыми другими операциями и полностью отделимо от любой модели или другого кода, который я использую.
Примечание. Для этого требуется установка tensorflow-gpuи cuda/ cudnn, поскольку предоставляется возможность использовать графический процессор.
Ссылки:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'как в ответе ниже