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

Tensorflowのメモリ確保をChainerっぽくする

More than 3 years have passed since last update.

Tensorflow (v0.12時点) はデフォルトではCUDAから見えている全GPUのメモリを専有してしまいます。また、GPUを使い学習したモデルをCPUで使うことができません。これを、Chainerと同じように、

  • with device: ...のような形で簡単にGPU/CPUを切り替え
  • メモリ使用にあわせて自動的にGPUメモリを増減

することを考えます。

そのような設定にするためには、tf.Sessionを下記のように作ります。

config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True

with tf.Session(config=config) as sess:
    with tf.device(my_device_option):
        model = MyGreatModel()
    sess.run(tf.global_variables_initializer())
    ...

tf.ConfigProtoには他にもオプションがありますが、それはに詳しいです。

koreyou
自然言語処理、応用機械学習、ロボットに興味がある企業リサーチャーです。 載せているコードは他に表記がなければThe UnlicenseとCC-0のデュアルライセンスにて公開しています。 発言は個人の見解であって所属組織を代表するものではありません。
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