TensorFlow【GPUの利用】

TensorFlowは、自動的に(優先的に)GPUを利用しているらしいですが、自分の環境でGPUを利用できているかを確認します。

www.tensorflow.org

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)

参考にしました

テンソルと演算  |  TensorFlow Core

Tensorflowで使用するGPUのメモリを制限したいとき - Qiita