【Python】使用するGPUを指定するコマンド【GPU】

複数枚のGPUがあるスパコンを使ってディープラーニングの学習をおこなうときに、すべてのGPUは使わずに一部だけを使いたいときがあると思う。

たいていは、複数枚使うことで、それぞれにメモリが割り振られて処理が早くなるけど、多すぎると逆に効率が悪くなる場合もあるので枚数を指定できるようになっておくのが良い。

使用するGPUを指定するには以下のオプションをコマンド実行時に追加すればよい。

CUDA_VISIBLE_DEVICES=0,1

これは、GPUの０番と１番の２枚を使うようにする場合。

CUDA_VISIBLE_DEVICES=0,1 python object_detection/model_main.py --logtostderr --model_dir=model --pipeline_config_path=pipeline.config

これが使用するGPUを指定して学習をするときの実行コマンドの例。

ここでは、Object Detection APIを使って学習させている。