pythonからGPUの枚数を知る方法〜multiprocessingをpytorchで使う際の注意点〜

はじめに

pytorchでmultiprocessingを利用するとCUDAの初期化で怒られることがあります．

RuntimeError: cuda runtime error (3) : initialization error at /pytorch/aten/src/THC/THCGeneral.cpp:50
THCudaCheck FAIL file=/pytorch/aten/src/THC/THCGeneral.cpp line=50 error=3 : initialization error

なぜ怒られたか

色々調査すると，spawnがどうとか，様々な文献が見当たりましたが，自分の場合はtorch.cuda.device_count()を利用していたことが原因だったようです．

ということで，GPUの枚数をtorch.cuda.device_count()抜きで知りたい．

GPUの枚数をpythonから取得

nvidia-smiに頼ります．
linuxの場合は以下の通り．

import subprocess
msg = subprocess.check_output("nvidia-smi --query-gpu=index --format=csv", shell=True)
n_devices = max(0, len(msg.decode().split("\n")) - 2)

pytorchのCUDA初期化問題には皆様も気をつけてください．