本記事では、PythonからGPUのメモリ使用量を取得する方法を紹介します。

前提条件

nvidia-smiが実行可能であること

本記事で紹介する方法は、Pythonからsubprocess関数を利用してNVIDIA GPUのメモリ使用量を計測します。

nvidia-smiのインストール方法はNVIDIA Documentationをご参照ください。

GPUメモリ取得方法

import subprocess

default_properies = (
  "timestamp",
  "gpu_name",
  #"gpu_uuid",
  "index",
  "memory.total",
  "memory.used",
  "memory.free",
  "utilization.gpu",
  "utilization.memory",
)

def get_gpu_properties(
  cmd_path="nvidia-smi",
  target_properties=default_properies,
  noheader=True,
  nounits=True
  ):
  """
  CUDA GPUのプロパティ情報取得

  Parameters
  ----------
  cmd_path : str
    コマンドラインから"nvidia-smi"を実行する際のパス
  target_properties : obj
    取得するプロパティ情報
    プロパティ情報の詳細は"nvidia-smi --help-query-gpu"で取得可能
  noheader : bool
    skip the first line with column headers
  nounits : bool
    don't print units for numerical values

  Returns
  -------
  gpu_properties : list
    gpuごとのproperty情報
  """
    
  # formatオプション定義
  format_option = "--format=csv"
  if noheader:
      format_option += ",noheader"
  if nounits:
      format_option += ",nounits"

  # コマンド生成
  cmd = '%s --query-gpu=%s %s' % (cmd_path, ','.join(target_properties), format_option)

  # サブプロセスでコマンド実行
  cmd_res = subprocess.check_output(cmd, shell=True)
    
  # コマンド実行結果をオブジェクトに変換
  gpu_lines = cmd_res.decode().split('\n')
  # リストの最後の要素に空行が入るため除去
  gpu_lines = [ line.strip() for line in gpu_lines if line.strip() != '' ]

  # ", "ごとにプロパティ情報が入っているのでdictにして格納
  gpu_properties = [ { k: v for k, v in zip(target_properties, line.split(', ')) } for line in gpu_lines ]

  return gpu_properties

使用例は以下の通りです。

properties = get_gpu_properties()
# GPU Property情報出力
for property in properties:
  print(property)
  print(property["gpu_name"], "[", property["index"], "] used", property["memory.used"], "MiB" )
  # {'timestamp': '2021/10/04 12:40:37.100', 'gpu_name': 'GeForce GTX 1080 Ti', 'index': '1', 'memory.total': '11178', 'memory.used': '2', 'memory.free': '11176', 'utilization.gpu': '0', 'utilization.memory': '0'}
  # GeForce GTX 1080 Ti [ 0 ] used 19 MiB

get_gpu_propertiesでGPUのProperty情報を取得し、property情報はdict型でreturnされます。
GPUメモリの使用量は"memory.used"に格納されています。

また、target_propertiesに設定できるpropertyはnvidia-smi --help-query-gpuで参照可能です。

nvidia-smi --query-gpuの詳細については公式ドキュメントをご参照ください。

GPUメモリ使用量を定期的に取得する

上記get_gpu_propertiesを定期実行することにより、定期的にGPUメモリ使用量を測定できます。
下記のコード例は厳密な定期実行ではない点ご注意ください。概ねintervalで設定した周期で実行します。

import threading
import time

g_IsMonitor = False
g_GpuPropertiesList = []

def monitoring_gpu_task(interval):
  global g_IsMonitor
  global g_GpuPropertiesList

  while g_IsMonitor:
    time.sleep(interval)
    res = get_gpu_properties()
    g_GpuPropertiesList.append( res )

def start_gpu_monitor(interval):
  global g_IsMonitor
  global g_GpuPropertiesList

  g_IsMonitor = True
  g_GpuPropertiesList.clear()

  monitor_thread = threading.Thread(target=monitoring_gpu_task, args=(interval, ))
  monitor_thread.start()

def end_gpu_monitor():
  global g_IsMonitor
  global g_GpuPropertiesList
  g_IsMonitor = False
  return g_GpuPropertiesList

# 1sec間隔でGPUのProperty情報を取得
start_gpu_monitor(1)
time.sleep(3)
properties_list = end_gpu_monitor()
print(properties_list)
# [
#   [
#     {'timestamp': '2021/10/04 15:33:24.015', 'gpu_name': 'GeForce GTX 1080 Ti', 'index': '0', 'memory.total': '11178', 'memory.used': '19', 'memory.free': '11159', 'utilization.gpu': '0', 'utilization.memory': '0'}, 
#     {'timestamp': '2021/10/04 15:33:24.016', 'gpu_name': 'GeForce GTX 1080 Ti', 'index': '1', 'memory.total': '11178', 'memory.used': '2', 'memory.free': '11176', 'utilization.gpu': '0', 'utilization.memory': '0'}
#   ], 
#   [
#     {'timestamp': '2021/10/04 15:33:25.039', 'gpu_name': 'GeForce GTX 1080 Ti', 'index': '0', 'memory.total': '11178', 'memory.used': '19', 'memory.free': '11159', 'utilization.gpu': '0', 'utilization.memory': '0'}, 
#     {'timestamp': '2021/10/04 15:33:25.039', 'gpu_name': 'GeForce GTX 1080 Ti', 'index': '1', 'memory.total': '11178', 'memory.used': '2', 'memory.free': '11176', 'utilization.gpu': '0', 'utilization.memory': '0'}
#   ]
# ]

非機能要件を把握しておくことの重要性

機能面の要件以外を意味する非機能要件はユーザービリティなどに影響します。

本記事で紹介した、メモリ使用量は開発環境と本番環境における実行速度の違いなどに影響を及ぼすため開発段階であらかじめ把握しておくことをお勧めします。

例えば、開発環境では潤沢なGPUメモリを持つマシン上で動作させており、本番環境は運用コストを理由にGPUメモリを削減している場合を想定します。

この時、本番環境では、GPUメモリの解放待ちが発生し、想定より実行速度が遅くなる場合があります。最悪の場合、OutOfMemoryが発生し、処理継続が困難な可能性もあります。

このため、事前に非機能要件を定義しており、それぞれの数値を把握しておくことは開発後期での手戻りを防ぐ一助となります。