как заставить Tensorflow2.2 обращаться к графическому процессору на кластере графических процессоров databricks

Question

Мой вопрос имеет отношение к моему предыдущему вопросу. как включить отображение графического процессора для среды выполнения ML на модулях данных?

 6.6 ML, spark 2.4.5, GPU, Scala 2.11  
 Keras version : 2.2.5

 nvidia-smi
 NVIDIA-SMI 440.64.00    Driver Version: 440.64.00    CUDA Version: 10.2    

Я создал новый пост, потому что мне нужно опубликовать больше кода.

Я установил Tensorflow2 .2 на блоках данных, запустив сценарий оболочки: https://docs.databricks.com/applications/deep-learning/single-node-training/tensorflow.html#install -tensorflow-22-on-databricks-runtime-66-ml & language-GPU

Я пытаюсь запустить некоторый Tensorflow (2.2 ) код на кластере графического процессора Databricks (p2.xlarge).

Код python3 находится по адресу:

https://www.tensorflow.org/guide/gpu 

Мой код запускается из записной книжки databricks:

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
# 0 

import tensorflow as tf
print("Tensorflow version : {}".format(tf.__version__))
print("Num XLA_GPUs Available: ", len(tf.config.experimental.list_physical_devices('XLA_GPU')))
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
print("phy devices Available: ", tf.config.experimental.list_physical_devices())

# Tensorflow version : 2.2.0
# Num XLA_GPUs Available:  1
# Num GPUs Available:  0
# phy devices Available:  [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), 
# PhysicalDevice(name='/physical_device:XLA_CPU:0', device_type='XLA_CPU'), PhysicalDevice(name='/physical_device:XLA_GPU:0', device_type='XLA_GPU')]

# use gpu https://www.tensorflow.org/guide/gpu
tf.debugging.set_log_device_placement(True)

 # Create some tensors
 a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
 b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
 c = tf.matmul(a, b)

 print(c)

вывод из приведенного выше кода:

tf.Tensor([[22. 28.][49. 64.]], shape=(2, 2), dtype=float32)

Но он работает только на CPU, а не на GPU.

gpus = tf.config.experimental.list_physical_devices('XLA_GPU')
 
phy_devices = tf.config.list_physical_devices()
print(phy_devices)

 visible_devices = tf.config.get_visible_devices()
 print('visible_devices : ' + str(visible_devices))

 if not gpus:
    print('no gpus')
 else:
    print(gpus)
    print(gpus[0])
    # Restrict TensorFlow to only use the first GPU
try:
    tf.config.experimental.set_visible_devices(gpus[0], 'XLA_GPU') # exception poped here !!!
    logical_gpus = tf.config.experimental.list_logical_devices('XLA_GPU')
    print('logical_gpus : ' + str(logical_gpus))
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
 except RuntimeError as e:
   # Visible devices must be set before GPUs have been initialized
   print(e)

вывод:

visible_devices : [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')][PhysicalDevice(name='/physical_device:XLA_GPU:0', device_type='XLA_GPU')]
PhysicalDevice(name='/physical_device:XLA_GPU:0', device_type='XLA_GPU')
Visible devices cannot be modified after being initialized

Как сделать сделать доступ к тензорному потоку GPU?