name : Janghoo Lee github : https://github.com/ProtossDragoon contact : [email protected] published date : November, 2021
if 'COLAB_TPU_ADDR' in os.environ: # Check TPU
assert 'COLAB_TPU_ADDR' in os.environ, 'Missing TPU.'
tf_master = 'grpc://{}'.format(os.environ['COLAB_TPU_ADDR'])
TPU_ADDRESS = tf_master
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)
tf.config.experimental_connect_to_cluster(resolver) # initialize the colab tpu
tf.tpu.experimental.initialize_tpu_system(resolver) # <https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/keras_mnist_tpu.ipynb?hl=ko&authuser=2#scrollTo=Hd5zB1G7Y9-7>
TRAINING_PARALLEL_STRATEGY = tf.distribute.TPUStrategy(resolver) # Choose distribution strategy for parallel processing.
tpus = tf.config.list_logical_devices('TPU')
print(f'total {len(tpus)} of TPU devices: {tpus}')
USE_TPU = True
USE_GPU = False
else:
USE_TPU = False
print('TPU Not found')
TRAINING_PARALLEL_STRATEGY = None
device_name = tf.test.gpu_device_name()
if not device_name:
USE_GPU = False
print('GPU device not found.')
else:
USE_GPU = True
!nvidia-smi -L
gpus = tf.config.list_logical_devices('GPU')
print(f'total {len(gpus)} GPU devices: {gpus}')
if USE_TPU:
CURRENT_DEVICE = 'tpu'
elif USE_GPU:
CURRENT_DEVICE = 'gpu'
else:
CURRENT_DEVICE = 'cpu'
INFO:tensorflow:Deallocate tpu buffers before initializing tpu system.
INFO:tensorflow:Deallocate tpu buffers before initializing tpu system.
WARNING:tensorflow:TPU system grpc://10.25.231.146:8470 has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.
WARNING:tensorflow:TPU system grpc://10.25.231.146:8470 has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.
INFO:tensorflow:Initializing the TPU system: grpc://10.25.231.146:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.25.231.146:8470
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8