複数GPUの簡単切り替え法

使用するGPUを指定する3つの方法

複数のGPUを利用する時に、どのGPUを使うか指定しないと、全てを占領してしまう。
そこで、使うGPUを指定する方法を調べたら3つあることがわかった。

1. Python コード内

import os

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

2. スクリプト実行時

コード内に書くより、こっちの方が個人的にオススメ

$ env CUDA_VISIBLE_DEVICES=0,1 python train.py

3. configに登録

fishの場合

~/.config/fish/config.fish

function pyconda0
    nohup env CUDA_VISIBLE_DEVICES=0 python $argv &
end

function pyconda1
    nohup env CUDA_VISIBLE_DEVICES=1 python $argv &
end

function pyconda2
    nohup env CUDA_VISIBLE_DEVICES=2 python $argv &
end

# GPU:2を使って学習する時
$ pyconda2 mnist.py

bashの場合

~/.bash_aliases

pyconda () {
    env CUDA_VISIBLE_DEVICES=$1 python $2
}

# GPU:1 と GPU:2を使って学習する時
$ pyconda 1,2 mnist.py

良いGPUライフを！！