Notebook環境でGPUメモリ使用量をリアルタイム監視する

TL;DR

GPUメモリの使用量をすぐ取得できるようなPython関数をつくってみた

はじめに

GPUメモリの利用状況を確認するためには nvidia-smi や nvidia-smi -q -d MEMORY などの各種コマンドを利用できます。

$ nvidia-smi -q -d MEMORY

>>> 
==============NVSMI LOG==============

Timestamp                           : Sat Nov 30 16:02:56 2019
Driver Version                      : 418.67
CUDA Version                        : 10.1

Attached GPUs                       : 1
GPU 00000000:00:04.0
    FB Memory Usage
        Total                       : 11441 MiB
        Used                        : 0 MiB
        Free                        : 11441 MiB
    BAR1 Memory Usage
        Total                       : 16384 MiB
        Used                        : 2 MiB
        Free                        : 16382 MiB

さらに watch コマンドと組み合わせると指定秒毎に情報を取得することができます。

$ watch -n 10 'nvidia-smi -q -d MEMORY'  # 10秒ごとに表示を繰り返す

これらの機能は, GPUメモリを溢れさせないためのモニタリングに適しています。

しかし問題なのは, どんな時でも使える訳ではなさそうという点です。

機械学習をシェルから走らせる場合にはいいかもしれませんが, Jupyter Notebook, Google Colab, Kaggle KernelなどのNotebook環境でこれをやろうとすると監視の無限ループだけが回り続けて他のセルが実行できなくなるので使い物になりません。

実装してみる

さっそく実装してみました。コードはこちらです。

import subprocess
import shlex

def gpuinfo():
    """
    Returns size of total GPU RAM and used GPU RAM.

    Parameters
    ----------
    None

    Returns
    -------
    info : dict
        Total GPU RAM in integer for key 'total_MiB'.
        Used GPU RAM in integer for key 'used_MiB'.
    """

    command = 'nvidia-smi -q -d MEMORY | sed -n "/FB Memory Usage/,/Free/p" | sed -e "1d" -e "4d" -e "s/ MiB//g" | cut -d ":" -f 2 | cut -c2-'
    commands = [shlex.split(part) for part in command.split(' | ')]
    for i, cmd in enumerate(commands):
        if i==0:
            res = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        else:
            res = subprocess.Popen(cmd, stdin=res.stdout, stdout=subprocess.PIPE)
    total, used = map(int, res.communicate()[0].decode('utf-8').strip().split('\n'))
    info = {'total_MiB':total, 'used_MiB':used}
    return info

実行してみます。

In:  gpuinfo()
Out: {'total_MiB': 11441, 'used_MiB': 0}

GPUのメモリのサイズと使用量が取得できました。
自動的に繰り返すとまではいきませんが, これでNotebookの好きな箇所でGPU使用状況を確認することができます。

表示を洗練させてみる

ログに残しやすいよう, 少し手を加えてメッセージ文字列を出力させてみます。

def gpulife():
    """
    Returns GPU usage information in string.

    Parameters
    ----------
    None

    Returns
    -------
    msg: str
    """

    def _gpuinfo():
        command = 'nvidia-smi -q -d MEMORY | sed -n "/FB Memory Usage/,/Free/p" | sed -e "1d" -e "4d" -e "s/ MiB//g" | cut -d ":" -f 2 | cut -c2-'
        commands = [shlex.split(part) for part in command.split(' | ')]
        for i, cmd in enumerate(commands):
            if i==0:
                res = subprocess.Popen(cmd, stdout=subprocess.PIPE)
            else:
                res = subprocess.Popen(cmd, stdin=res.stdout, stdout=subprocess.PIPE)
        return tuple(map(int, res.communicate()[0].decode('utf-8').strip().split('\n')))

    total, used = _gpuinfo()
    percent = int(used / total * 100)
    msg = 'GPU RAM Usage: {} {}/{} MiB ({:.1f}%)'.format('|' * (percent // 5) + '.' * (20 - percent // 5), used, total, used/total*100)
    return msg

このようになりました。

In:  gpulife()
Out: GPU RAM Usage: .................... 0/11441 MiB (0.0%)

なお, 上記のコードはいまのところ1GPUの場合にしか対応していません。
複数GPUの場合はそのままでは使えませんので, ご注意ください。
また, 動作確認はGoogle Colabでしか行っていません。その点もご容赦ください。

詳細

nvidia-smi -q -d MEMORY の出力結果をストリームエディタで加工しているだけです。
シェルのコマンドは subprocess モジュールを使えばPythonで実行できますが, パイプラインがある場合は単に subprocess.run() に渡すだけでは実行できないため, Python上でパイプラインを構築しています。

参考資料

公式ドキュメント subprocess --- サブプロセス管理
[Python]可読性を上げるための、docstringの書き方を学ぶ（NumPyスタイル）