Достаточно отдельный способ сделать это - использовать
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'
как в ответе ниже