Help us understand the problem. What is going on with this article?

Tensorflow系のライブラリで使用するGPUを指定する

More than 1 year has passed since last update.

はじめに

Tensorflowでは,デフォルトの設定で回すと搭載されているGPUすべてのメモリを確保してしまいます.TensorflowをバックエンドにしたKerasも同様です.個人のデスクトップPCで回す場合には問題ありませんが,複数のGPUを搭載した共有のGPUサーバを利用する場合,大変迷惑です.
以下に使用するGPUを指定するサンプルコードをまとめました.

Tensorflow

import tensorflow as tf
config = tf.ConfigProto(
    gpu_options=tf.GPUOptions(
        visible_device_list="0,1", # specify GPU number
        allow_growth=True
    )
)
sess = tf.Session(config=config)

Sessionを確保するときにconfigを渡してやります.
allow_growth=Trueを指定することによって,指定したGPUのメモリを最初にすべて確保するのではなく,必要に応じて確保するように設定できます.
あまりうれしくはありませんが他に空いているGPUがない場合,メモリが空いていれば複数のプログラムを回すことが可能です.

keras

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto(
    gpu_options=tf.GPUOptions(
        visible_device_list="2", # specify GPU number
        allow_growth=True
    )
)
set_session(tf.Session(config=config))

(https://github.com/fchollet/keras/issues/1538)

基本的にはバックエンドとして動作するTensorflowの設定を変更します.keras.backend.tensorflow_backend.set_sessionを用います.
注意点としては,コードの内部でtf.Sessionを取得しているようなケースでは,先頭に上記を記述しても意味がありません(とくに自分で書いていないコードを走らせる場合).

(nvidia-docker)

nvidia-dockerにはコンテナに対してマウントするGPUをを設定するオプションがあります.

$ NV_GPU=3,4 nvidia-docker run HOGEHOGE

NV_GPUによって設定できます.

ymfj
Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away