nvidia-smiでpythonからGPU温度を監視

GPUが古くなったせいか、nvidiaドライバの自動制御があってもGPU温度が高くなるとGPUが勝手に止まるようになってきたので、nvidia-smiをpythonからたたいて温度情報を取り出し、温度が高くなったらGPU計算をいったん停止するスクリプトを書いてみる。

nvidia-smi

LinuxでGPU使っているときにGPUのメモリ量などを常時モニタするために便利なnvidia-smi。基本コマンドで得られる情報は引数指定することで個々の値を取ってくることができる。nvidia-smi -lとすれば、常時モニタもできる。下の図の黄色い四角でくくった部分が温度。

nvidia-smiで温度情報を得る

温度情報のみをとってくるにはターミナルで以下のコマンドを打つ。

nvidia-smi -q -d temperature

出力はこちら

==============NVSMI LOG==============

Timestamp                           : Fri May  1 10:12:17 2020
Driver Version                      : 440.64.00
CUDA Version                        : 10.2

Attached GPUs                       : 1
GPU 00000000:01:00.0
    Temperature
        GPU Current Temp            : 37 C
        GPU Shutdown Temp           : 101 C
        GPU Slowdown Temp           : 96 C
        GPU Max Operating Temp      : N/A
        Memory Current Temp         : N/A
        Memory Max Operating Temp   : N/A

これだとまだ情報量が多いので、GPUの温度だけ取ってくるには以下のコマンド。これで温度情報だけを取り出せる。

nvidia-smi --query-gpu=temperature.gpu --format=csv,noheader

上記のコマンドの出力は37とだけ表示される。

＃複数GPUの場合は試してないので、どういう表示になるかはわかりません。

参考までにメモリの使用量や空きのみを調べるには下記のコマンドを打てばよい。

nvidia-smi --query-gpu=memory.free --format=csv,noheader
nvidia-smi --query-gpu=memory.used --format=csv,noheader

温度は数値だけだけど、こちらはなぜか55 MiBというように単位が返ってくるので、pythonで読み出すときはsplitとか使って数値部分だけを読み出す必要がある。

pythonからのGPU温度読み出し

pythonでnvidia-smiコマンドから温度情報のみを取り出す。コマンドラインの情報を取り出すのにsubprocessを使う。subprocessのコマンドはスペースではなく、コマンドライン引数のtupleかリストで入れることに注意。

subprocessでは文字列が返ってくるので、int()で数値に変換する。温度は数値のみが返ってくるので、そのままintに入れても大丈夫。

import subprocess

def getTemp():
    cmd = 'nvidia-smi', '--query-gpu=temperature.gpu', '--format=csv,noheader'
    temperature = int(subprocess.check_output(cmd))
    return temperature

温度が取ってこれるようになったので、温度に下限値と上限値を設けて、上限値を超えたら計算をやめてGPUを冷やし、下限値に達したら計算を再開するための温度モニタを行う。

import time

def coolGPU(lower=50, upper=75):
    temp = getTemp()
    print('GPU temperature: %d C' % temp)
    if temp > upper:
        print('cooling GPU')
        while (temp > lower):
            temp = getTemp()
            print('GPU temperature: %d C' % temp)
            time.sleep(1)

この関数を学習処理のループに入れることで温度を監視しながら学習できるようになる。上記のスクリプトはGPUの温度が75度より大きくなったら50度まで何もしないで冷やすとうようなことをしている。

例えば、下記のように使う。

for iepoch in range(nepoch):
    for i, data in enumerate(train_loader):
        x, y = data
        x = x.to(device)
        y = y.to(device)
        y_ = model.forward(x)
        loss = torch.mean((y_-y)**2)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if i%10==0: # 10イテレーションおきに温度監視
            coolGPU()

今回はGPUの温度が下がるまで何もしないというようにしたけど、PytorchならGPU温度が下がるまでCPUで学習を続けるというように改良してもいいかも。他にも、複数GPUで交互に学習するみたいなこともできそう。

普通に学習するより少し手間がかかるけど、これで GPUの寿命も伸びるかもしれない。

参考サイト

nvidia.custhelp.com