GPU周り使い方

pip install -q tensorflow-gpu==2.0.0-rc1

import tensorflow as tf

x = tf.random.uniform([3, 3])

print("利用できるGPUはあるか: "),
print(tf.config.experimental.list_physical_devices("GPU"))

print("テンソルはGPU #0にあるか:  "),
print(x.device.endswith('GPU:0'))

import time

def time_matmul(x):
    start = time.time()
    for loop in range(10):
        tf.matmul(x, x)

    result = time.time()-start
    print("10 loops: {:0.2f}ms".format(1000*result))

# CPUでの実行を強制
print("On CPU:")
with tf.device("CPU:0"):
    x = tf.random.uniform([1000, 1000])
    assert x.device.endswith("CPU:0")
    time_matmul(x)

# GPU #0があればその上での実行を強制
if tf.config.experimental.list_physical_devices("GPU"):
    print("On GPU:")
    with tf.device("GPU:0"): # 2番めのGPUなら GPU:1, 3番目なら GPU:2 など
        x = tf.random.uniform([1000, 1000])
        assert x.device.endswith("GPU:0")
        time_matmul(x)