Una forma bastante separable de hacer esto es usar
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)
Aquí, con booleans
GPU
y CPU
, indicamos si nos gustaría ejecutar nuestro código con la GPU o la CPU definiendo rígidamente la cantidad de GPU y CPU a las que se permite acceder a la sesión de Tensorflow. Las variables num_GPU
y num_CPU
definen este valor. num_cores
luego establece el número de núcleos de CPU disponibles para su uso mediante intra_op_parallelism_threads
yinter_op_parallelism_threads
.
La intra_op_parallelism_threads
variable dicta el número de subprocesos que se permite utilizar (intra) una operación paralela en un solo nodo en el gráfico de cálculo. Mientras que la inter_ops_parallelism_threads
variable define el número de subprocesos accesibles para operaciones paralelas a través de los nodos del gráfico de cálculo (inter).
allow_soft_placement
permite que se ejecuten operaciones en la CPU si se cumple alguno de los siguientes criterios:
no hay implementación de GPU para la operación
no hay dispositivos GPU conocidos o registrados
existe la necesidad de coubicar con otras entradas de la CPU
Todo esto se ejecuta en el constructor de mi clase antes que cualquier otra operación, y es completamente separable de cualquier modelo u otro código que utilice.
Nota: Esto requiere tensorflow-gpu
y cuda
/cudnn
para ser instalado porque se da la opción de usar una GPU.
Refs:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
como en una respuesta a continuación