[PYTHON] Automatisches Wechseln zwischen TPU / GPU / CPU im Tensorflow.Keras-Modell

Einführung

In diesem Artikel wird beschrieben, wie Sie ein Programm schreiben, das Hardwareinformationen (hauptsächlich Colaboratory-Laufzeitinformationen) lesen und bei Verwendung von TensorFlow.keras automatisch zwischen TPU und GPU (einschließlich CPU) wechseln kann. Es ist zusammengefasst. ~~ (Weil es schwierig geworden ist, manuell zu kommentieren) ~~

Im Allgemeinen wird MNIST mit Keras und TPU auf Offizielle Website (Google Cloud, Cloud-TPU-Dokumente) veröffentlicht. /github/tensorflow/tpu/blob/master/tools/colab/keras_mnist_tpu.ipynb) wird zusammengefasst. Wenn Sie mit tensorflow.keras vertraut sind, ist es möglicherweise einfacher, die Referenzquelle zu lesen.

Ergänzungen, Vorsichtsmaßnahmen usw.

• Die Überprüfung des Betriebs wird im Labor durchgeführt.
• Möglicherweise unterstützt es auch Tensorflow ver2. --Code für tensorflow.keras. (Nicht Puer Keras oder Tensoflow)
• Möglicherweise liegt ein Fehler in der Beschreibung vor (insbesondere in der Erklärung).

Code, der automatisch zwischen TPU / GPU / CPU umschaltet

#Hardware-Informationen abrufen
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)


#Mit Strategie beim Erstellen, Laden und Kompilieren von Modellen.scope()Umgeben mit
from tensorflow import keras
with strategy.scope():
    model = make_model()

Wenn Sie die folgenden zwei Punkte ausführen, ohne über bestimmte Schwierigkeiten nachzudenken, können Sie je nach Hardware ohne Erlaubnis zwischen TPU / GPU / CPU wechseln.

1. import tensorflow as tf ~ print ("Anzahl der Beschleuniger (abgekürzt wird durch Kopieren und Einfügen eingefügt))
2. Fügen Sie die Modelldefinition und den Modellladeteil in den Bereich von with Strategy.scope (): ein

So etwas wie ein Kommentar

• Priorität --Wählen Sie in der Reihenfolge "TPU> Multi-GPU> Einzel-GPU> CPU" je nach Programm.
• tf.distribute.cluster_resolver.TPUClusterResolver()
• Informationen zur TPU-Hardware erhalten. In einer Umgebung, in der TPU nicht verfügbar ist, tritt ein Fehler auf.
• tf.config.experimental.list_logical_devices("GPU")
• Informationen zur GPU-Hardware erhalten. Der Rückgabewert ist eine Liste. Wenn Sie "CPU" angeben, können Sie auch CPU-Informationen abrufen. --Nach if tpu:
• So definieren Sie jedes Gerät. TPU und MultiGPU sind etwas Besonderes, und Sie sollten sich daran erinnern, dass "dies das ist, was Sie schreiben".
• tf.distribute.get_strategy()
• Wir prüfen nicht die Hardwareinformationen, sondern, ob der installierte Tensorflow GPU-kompatibel ist.
• Mit anderen Worten, es wird (sollte) bestimmt, wann der Tensorflow installiert ist, nicht zur Laufzeit.
• Wenn Sie eine GPU haben, diese aber nicht funktioniert, überprüfen Sie die CUDA-Installation (Referenz: GPU-Unterstützung)

abschließend

In diesem Artikel habe ich erklärt, wie mit tensorflow.keras je nach Hardwarestatus automatisch zwischen TPU / GPU / CPU gewechselt wird. Nutzen Sie TPU mit tensorflow.keras für ein besseres Deep-Learning-Leben.