web-dev-qa-db-ja.com

Tensorflowはjupyter内でCUDA_VISIBLE_DEVICESを設定しました

2つのGPUがあり、ipynbを介して2つの異なるネットワークを同時に実行したいのですが、最初のノートブックは常に両方のGPUを割り当てます。

CUDA_VISIBLE_DEVICESを使用すると、pythonファイルのデバイスを非表示にできますが、ノートブック内でどのように非表示にするかはわかりません。

とにかく、同じサーバーで実行されているノートブックに異なるGPUを隠す方法はありますか?

51
Tim

os.environを使用して、ノートブックに環境変数を設定できます。 TensorFlowを初期化してからTensorFlowを最初のGPUに制限する前に、以下を実行します。

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="0"

TFに正しいデバイスが表示されていることを再確認できます

from tensorflow.python.client import device_lib
print device_lib.list_local_devices()

notebook_util のようなユーティリティモジュールから使用する傾向があります

import notebook_util
notebook_util.pick_gpu_lowest_memory()
import tensorflow as tf
110

magics を使用するだけで、インポートせずに高速に実行できます。

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

すべてのenv変数は文字列なので、"を使用する必要はありません。 env-variableがセットアップされていることを確認するには、%env <name_of_var>を実行します。または、%envでそれらすべてを確認します。

20
Salvador Dali