Une façon plutôt séparable de faire ceci est d'utiliser
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)
Ici, avec booleans
GPU
et CPU
, nous indiquons si nous souhaitons exécuter notre code avec le GPU ou le CPU en définissant de manière rigide le nombre de GPU et de CPU auxquels la session Tensorflow est autorisée à accéder. Les variables num_GPU
et num_CPU
définissez cette valeur. num_cores
définit ensuite le nombre de cœurs CPU disponibles pour une utilisation via intra_op_parallelism_threads
etinter_op_parallelism_threads
.
La intra_op_parallelism_threads
variable dicte le nombre de threads qu'une opération parallèle dans un seul nœud du graphe de calcul est autorisée à utiliser (intra). Alors que la inter_ops_parallelism_threads
variable définit le nombre de threads accessibles pour les opérations parallèles à travers les nœuds du graphe de calcul (inter).
allow_soft_placement
permet d'exécuter des opérations sur la CPU si l'un des critères suivants est rempli:
il n'y a pas d'implémentation GPU pour l'opération
il n'y a aucun périphérique GPU connu ou enregistré
il est nécessaire de co-localiser avec d'autres entrées du CPU
Tout cela est exécuté dans le constructeur de ma classe avant toute autre opération, et est complètement séparable de tout modèle ou autre code que j'utilise.
Remarque: Cela nécessite tensorflow-gpu
et cuda
/ cudnn
doit être installé car l'option est donnée d'utiliser un GPU.
Réf:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
comme dans une réponse ci