Jupyter上でKerasを使うときにGPUを無効化する方法

machine-learning
環境設定 Jupyter Keras

研究室のみんなでサーバ上のGPUを共有してたりすると、Jupyter上でKerasを使うときにGPUを無効化したくなる時ありますよね。 その方法を見つけたので備忘録としてメモしておきます。

特に今自分のいる環境では、サーバ群がslurm workload managerというツールで管理されています。 下手にJupyterでKerasを使うとそのサーバ上の全てのGPUを勝手に占有し、他の人のジョブの挙動がおかしくなって（？）しまうのです。 なので、Jupyter上でKerasを使って軽く実験するときには、必ずGPUを使わないように設定する必要がありました。

自分の環境では以下の方法でGPUを無効化できました。

# jupyter notebook上です。
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import keras

Kerasをインポートする前に環境変数を設定する方法が手っ取り早いようです。

念のため自分の環境の情報をメモしておきます。

  • OS: CentOS Linux release 7.2.1511 (Core)
  • Python: 3.5.0
  • Keras: 2.0.6 (バックエンドはtensorflow)
  • tensorflow: 1.2.1
  • jupyterはクライアントからSSHで利用

別の方法もあるようですが、自分の環境ではうまくいきませんでした・・・。

もし同じことで悩んでいたら是非試してみてください。