LoginSignup
3
4

More than 3 years have passed since last update.

PythonからGPUデバイスが使用可能かチェックする

Posted at

TL;DR

Pythonの処理からデバイス番号を指定して,指定したGPUが使えるかどうかをチェックする方法です.
ここで使えるとは,デバイスが存在することを指します.

GPUデバイスが使用可能かチェックする

NVIDIAからライブラリ等が提供されているわけではないので,nvidia-smiコマンドを叩いた結果から確認します.
(ライブラリがあった場合はご教授お願い致します)

以下のようにします.

"""
GPUデバイスが使用可能か確認する
"""
import subprocess


def check_gpu_device(device: int) -> bool:
    """
    GPUデバイスが使用可能か確認する
    Args:
        device: デバイス番号
    """
    output = subprocess.check_output(["nvidia-smi", "-L"])
    lines = output.decode().split('\n')
    for line in lines:
        if line == "":
            continue
        gpu_device_id = line.split(":")[0]
        if gpu_device_id == f"GPU {device}":
            return True
    return False

-Lオプションを付けることで1行1デバイスで列挙できるので,Pythonから扱いやすくなります.
(--query-gpuでもOKです)

まとめ

PythonからGPUデバイスが使用可能かチェックする処理を紹介しました.
基本的には自分が使っているPCのデバイス番号は把握しているかと思いますが,何かしらのアプリケーションの形で提供している場合には使えるかと思います.
ご参考になれば幸いです.

3
4
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
3
4