Keras model fit doens't use XLA_GPU device

Viewed 53

tensorflow-gpu==2.3.0 is properly installed. GPU device can be discovered.

tf.config.get_visible_devices() shows the GPU device correctly:

[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:XLA_GPU:0', device_type='XLA_GPU')]

However, after I build my model using Keras and try model.fit, I found the GPU is not utilized at all.

the execution shows:

95/237 [===========>..................] - ETA: 3s - loss: 2883.6201 - mse: 2883.6201 Executing op __inference_train_function_56696 in device /job:localhost/replica:0/task:0/device:CPU:0 Executing op __inference_train_function_56696 in device /job:localhost/replica:0/task:0/device:CPU:0 97/237 [===========>..................] - ETA: 3s - loss: 2878.0310 - mse: 2878.0310 Executing op __inference_train_function_56696 in device /job:localhost/replica:0/task:0/device:CPU:0 Executing op __inference_train_function_56696 in device /job:localhost/replica:0/task:0/device:CPU:0 99/237 [===========>..................] - ETA: 3s - loss: 2876.5935 - mse: 2876.5935 Executing op __inference_train_function_56696 in device /job:localhost/replica:0/task:0/device:CPU:0

What can I do to find out what's going on under the hood? How can I run model fit on a sepcified GPU?

0 Answers
Related