[PYTHON] Automatically switch between TPU / GPU / CPU in Tensorflow.Keras model

Introduction

This article describes how to write a program that can read hardware information (mainly Colaboratory runtime information) and automatically switch between TPU and GPU (including CPU) when using TensorFlow.keras. It is summarized. ~~ (Because it has become troublesome to comment out manually) ~~

Generally, MNIST with Keras and TPU posted on Official Site (Google Cloud, Cloud TPU Docs) /github/tensorflow/tpu/blob/master/tools/colab/keras_mnist_tpu.ipynb) is a summary. If you are familiar with tensorflow.keras, it may be easier to read the reference source.

Supplements, cautions, etc.

--Operation verification is carried out on Colaboratory. --Maybe it also supports tensorflow ver2. --Code for tensorflow.keras. (Not puer keras or tensoflow) --There may be an error in the description (especially the explanation).

Code that automatically switches between TPU / GPU / CPU

#Hardware information acquisition
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)


#With strategy when creating, loading, and compiling models.scope()Surround with
from tensorflow import keras
with strategy.scope():
    model = make_model()

If you do the following two points without thinking about any particular difficulty, you will be able to switch TPU / GPU / CPU without permission depending on the hardware.

1. ʻimport tensorflow as tf ~ print ("Number of accelerators (abbreviation`) is pasted by copy and paste
2. Put the model definition and model loading part in the scope of with strategy.scope ():

Something like a commentary

• priority --Select in the order of "TPU> multi GPU> single GPU> CPU" according to the program.
• tf.distribute.cluster_resolver.TPUClusterResolver() --Obtain TPU hardware information. An error will occur in an environment where TPU is not available.
• tf.config.experimental.list_logical_devices("GPU") --Obtain GPU hardware information. The return value is a list. If you give "CPU", you can also get CPU information. --ʻIf tpu: `after --How to define each device. TPU and multiGPU are special and you should remember that "this is what you write".
• tf.distribute.get_strategy() --We are not looking at the hardware information, but looking at whether the installed tensorflow is GPU compatible. --In other words, it is (should) be determined when tensorflow is installed, not at runtime. --If you have a GPU but it doesn't work, review the CUDA installation (Reference: GPU support)

in conclusion

In this article, I explained how to use tensorflow.keras to automatically switch between TPU / GPU / CPU depending on the hardware situation. Take advantage of TPU with tensorflow.keras for a better deep learning life.