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

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のデバイス番号は把握しているかと思いますが,何かしらのアプリケーションの形で提供している場合には使えるかと思います.
ご参考になれば幸いです.

シェアする

  • このエントリーをはてなブックマークに追加

フォローする