Python のマルチスレッド処理とGPUメモリ使用率の測定

ちょっと Deep Learning の推論処理のパフォーマンスを測定してほしいとのことでどうしたもんかと思ったんですが、マルチスレッドでひたすら処理中に GPU を測定するコマンドを発行すればいいんじゃねという脳筋的解決策で取り組んでみた。

Python初心者ですが、Pythonでも意外と簡単にマルチスレッドが作れるらしい。

import threading

class MTPerform(threading.Thread):
    def __init__(self, gpu, interval=0.1):
        super(MTPerform, self).__init__()
        self.interval = interval
        self.gpu = gpu
        self.unit = 'gpu' if gpu >= 0 else 'cpu'
        self.terminate = False
    
    def run(self):
        while True:
            self.measure(unit=self.unit)
            if self.terminate: break

    def stop(self):
        self.terminate = True

threding.Thread を継承して、やらせたい処理は run() をオーバーライドすればいいいと。いや、この処理どうやって止めりゃいいんだ… 無理やり止めるか…(脳筋)

実際の GPU の測定は、以下によると、GPU の状態は nvidia-smi のコマンドを使えば csv 形式に落とせるらしい。

qiita.com

attribute = ['memory.used','utilization.memory']

def get_gpu_info(self):
    
    cmd = 'nvidia-smi --id={} --query-gpu={} --format=csv,noheader,nounits'.format(
                self.gpu, ','.join(attribute))
    out = subprocess.check_output(cmd, shell=True)
    str_out = out.split(',')
        
    mem = int(str_out[0])
    use = int(str_out[1])
        
    return (mem, use)

–id= で GPU を指定、–query-gpu= は欲しい情報(今は memory.used, utilization.memory)、
–format= で csv, noheader, nounits(単位無し)を出力に指定。

ついでに CPU も測定できるように拡張しておくと、psutil を使えば上の GPU と同じような情報が抜けそうだったのでついでに準備してみる。

www.blog.umentu.work

attribute = ['memory.used','utilization.memory']

def get_gpu_info(self):
    
    cmd = 'nvidia-smi --id={} --query-gpu={} --format=csv,noheader,nounits'.format(
                self.gpu, ','.join(attribute))
    out = subprocess.check_output(cmd, shell=True)
    str_out = out.split(',')
        
    mem = int(str_out[0])
    use = int(str_out[1])
        
    return (mem, use)

あとは前回作った chainer 2.0 の MNIST のやつに測定用の関数をねじこんでみる。こちらで start() の関数を呼べばマルチスレッドで処理してくれるらしい。