1. ホーム
  2. python

[解決済み] Tensorflowはjupyter内でCUDA_VISIBLE_DEVICESを設定する。

2023-04-29 11:42:43

質問

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

CUDA_VISIBLE_DEVICES を使用して、私は python ファイルのデバイスを隠すことができますが、ノートブック内でそれを行う方法についてはわかりません。

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

どのように解決するのですか?

ノートブックで環境変数を設定するには 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