Iterate dispositifs cpu et gpu dans Tensorflow

Question

Je suis conscient que Tensorflow peut explicitement placer le calcul sur tous les périphériques par "/cpu0" ou "/gpu0". Cependant, ceci est codé en dur. Existe-t-il un moyen d'itérer tous les périphériques visibles avec l'API intégrée?

Answer 1

Voici ce que vous aimeriez avoir:

import tensorflow as tf 
from tensorflow.python.client import device_lib 

def get_all_devices(): 
    local_device_protos = device_lib.list_local_devices() 
    return [x.name for x in local_device_protos] 

all_devices = get_all_devices() 
for device_name in all_devices: 
    with tf.device(device_name): 
     if "cpu" in device_name: 
      # Do something 
      pass 
     if "gpu" in device_name: 
      # Do something else 
      pass 

code est inspiré de la meilleure réponse ici: How to get current available GPUs in tensorflow?