from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
tf.config.experimental.list_physical_devices()
[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'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
len(tf.config.experimental.list_physical_devices('GPU'))
1
tf.debugging.set_log_device_placement(True)
# Place tensors on the CPU
with tf.device('/CPU:0'):
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)
Tensor("MatMul_1:0", shape=(2, 2), dtype=float32)
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
# Create 2 virtual GPUs with 512MB memory each
try:
tf.config.experimental.set_virtual_device_configuration(
gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=512),
tf.config.experimental.VirtualDeviceConfiguration(memory_limit=512)])
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
print(len(gpus), "Physical GPU,", len(logical_gpus), "Logical GPUs")
except RuntimeError as e:
# Virtual devices must be set before GPUs have been initialized
print(e)
1 Physical GPU, 2 Logical GPUs
import time
cpu_slot = 0
gpu_slot = 0
# Using CPU at slot 0
with tf.device('/CPU:' + str(cpu_slot)):
# Starting a timer
start = time.time()
# Doing operations on CPU
A = tf.constant([[3, 2], [5, 2]])
print(tf.eye(2,2))
# Printing how long it took with CPU
end = time.time() - start
print(end)
# Using the GPU at slot 0
with tf.device('/GPU:' + str(gpu_slot)):
# Starting a timer
start = time.time()
# Doing operations on GPU
A = tf.constant([[3, 2], [5, 2]])
print(tf.eye(2,2))
# Printing how long it took with GPU
end = time.time() - start
print(end)
Tensor("eye_2/diag:0", shape=(2, 2), dtype=float32, device=/device:CPU:0) 0.002489328384399414 Tensor("eye_3/diag:0", shape=(2, 2), dtype=float32, device=/device:GPU:0) 0.0018601417541503906