Commutez automatiquement TPU / GPU / CPU avec le modèle Tensorflow.Keras

introduction

Cet article décrit comment écrire un programme capable de lire des informations sur le matériel (principalement des informations d'exécution Colaboratory) et de basculer automatiquement entre TPU et GPU (y compris le processeur) lors de l'utilisation de TensorFlow.keras. C'est résumé. ~~ (Parce qu'il est devenu difficile de commenter manuellement) ~~

Généralement, MNIST avec Keras et TPU publié sur Site officiel (Google Cloud, Cloud TPU Docs) /github/tensorflow/tpu/blob/master/tools/colab/keras_mnist_tpu.ipynb) est résumé. Si vous connaissez tensorflow.keras, il peut être plus facile de lire la source de référence.

Suppléments, mises en garde, etc.

• La vérification du fonctionnement est effectuée sur Colaboratory.
• Peut-être qu'il prend également en charge tensorflow ver2. --Code pour tensorflow.keras. (Pas puer keras ou tensoflow) --Il peut y avoir une erreur dans la description (en particulier l'explication).

Code qui bascule automatiquement entre TPU / GPU / CPU

#Obtenir des informations sur le matériel
import tensorflow as tf

try:
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection
except ValueError:
  tpu = None
  gpus = tf.config.experimental.list_logical_devices("GPU")

if tpu:
  tf.tpu.experimental.initialize_tpu_system(tpu)
  strategy = tf.distribute.experimental.TPUStrategy(tpu, steps_per_run=128) # Going back and forth between TPU and host is expensive. Better to run 128 batches on the TPU before reporting back.
  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
elif len(gpus) > 1:
  strategy = tf.distribute.MirroredStrategy([gpu.name for gpu in gpus])
  print('Running on multiple GPUs ', [gpu.name for gpu in gpus])
elif len(gpus) == 1:
  strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU
  print('Running on single GPU ', gpus[0].name)
else:
  strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU
  print('Running on CPU')
print("Number of accelerators: ", strategy.num_replicas_in_sync)


#Avec stratégie lors de la création, du chargement et de la compilation de modèles.scope()Entourez de
from tensorflow import keras
with strategy.scope():
    model = make_model()

Si vous faites les deux points suivants sans penser à une difficulté particulière, vous pourrez changer TPU / GPU / CPU sans autorisation en fonction du matériel.

1. ʻimport tensorflow as tf~print (" Le nombre d'accélérateurs (abrégé`) est collé par copier-coller
2. Mettez la définition du modèle et la partie chargement du modèle dans la portée de with strategy.scope ():

Quelque chose comme un commentaire

• priorité --Sélectionnez dans l'ordre "TPU> multi GPU> GPU unique> CPU" en fonction du programme.
• tf.distribute.cluster_resolver.TPUClusterResolver()
• Obtenir des informations sur le matériel TPU. Une erreur se produit dans un environnement où le TPU n'est pas disponible.
• tf.config.experimental.list_logical_devices("GPU")
• Obtenez des informations sur le matériel GPU. La valeur de retour est une liste. Si vous donnez "CPU", vous pouvez également obtenir des informations sur le CPU. --ʻSi tpu: ʻaprès --Comment définir chaque appareil. C'est une bonne idée de se rappeler que TPU et multiGPU sont spéciaux et «écrivez ceci».
• tf.distribute.get_strategy()
• Je ne regarde pas les informations sur le matériel, mais je regarde si le tensorflow installé est compatible avec le GPU.
• En d'autres termes, il est (devrait) être déterminé lorsque tensorflow est installé, pas au moment de l'exécution.
• Si vous avez un GPU mais qu'il ne fonctionne pas, passez en revue l'installation de CUDA (Référence: GPU support)

en conclusion

Dans cet article, j'ai expliqué comment utiliser tensorflow.keras pour basculer automatiquement entre TPU / GPU / CPU en fonction de l'état du matériel. Tirez parti du TPU avec tensorflow.keras pour une meilleure vie d'apprentissage en profondeur.