【TensorFlow】GPUが認識されているか確認する方法

この記事を読んで分かること

利用可能なGPUが存在するか確認する方法

TensorFlowでGPUが利用可能か確認する
1. 1. 認識されているGPUが存在するか確認する
2. 2. 認識されているリソースの一覧を表示
実用的なGPU利用のコード例
まとめ
参考文献

TensorFlowでGPUが利用可能か確認する

利用可能なGPUを確認する方法は2つあります。

1. 認識されているGPUが存在するか確認する

tensorflowで利用できるGPUが存在する場合にはTrueが返ってくる関数です。

import tensorflow as tf
tf.test.is_gpu_available()

True

2. 認識されているリソースの一覧を表示

認識されているCPUとGPUの一覧を表示する関数です。

from tensorflow.python.client import device_lib
device_lib.list_local_devices()

実行結果(クリックして表示)

[name: "/device:CPU:0"
  device_type: "CPU"
  memory_limit: 268435456
  locality {
  }
  incarnation: 255424782033352859
  xla_global_id: -1,
  name: "/device:GPU:0"
  device_type: "GPU"
  memory_limit: 6401949696
  locality {
    bus_id: 1
    links {
    }
  }
  incarnation: 5073644171173579923
  physical_device_desc: "device: 0, name: NVIDIA GeForce RTX 3070, pci bus id: 0000:01:00.0, compute capability: 8.6"
  xla_global_id: 416903419]

この例では、CPUとGPUが1つずつ認識されています。

実用的なGPU利用のコード例

紹介した2つのGPU確認方法を利用して、実用的なコードを作成しました。

tensorflowでGPUを利用する場合はこのようにすると便利になると思います。

import tensorflow as tf
from tensorflow.python.client import device_lib

if tf.test.is_gpu_available():
    gpus = tf.config.experimental.list_physical_devices('GPU')

    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)
    
    print("GPUの読み込みが完了しました")

else:
    print("GPUが存在していません")
    device_lib.list_local_devices()