プロセスを実行するGPUを指定したい。そして、次のように設定しました。
import tensorflow as tf
with tf.device('/gpu:0'):
a = tf.constant(3.0)
with tf.Session() as sess:
while True:
print sess.run(a)
ただし、2つのGPUの両方にメモリを割り当てます。
| 0 7479 C python 5437MiB
| 1 7479 C python 5437MiB
CUDA_VISIBLE_DEVICES=1
を設定する必要があると思います。または、使用するGPUを選択します。 GPUを1つだけ表示する場合、環境変数の設定に関係なく、tensorflowでそれを/gpu:0
と呼びます。
その環境変数の詳細: https://devblogs.nvidia.com/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/
これを達成するための3つの方法があります。
CUDA_VISIBLE_DEVICES
環境変数を使用します。環境変数CUDA_VISIBLE_DEVICES="1"
を設定すると、デバイス1のみが表示され、CUDA_VISIBLE_DEVICES="0,1"
を設定すると、デバイス0と1が表示されます。 pythonでこれを行うには、os
パッケージをインポートした後にos.environ["CUDA_VISIBLE_DEVICES"]="0,1"
行を追加します。
with tf.device('/gpu:2')
を使用してグラフを作成します。次に、GPUデバイス2を使用して実行します。
config = tf.ConfigProto(device_count = {'GPU': 1})
を使用してからsess = tf.Session(config=config)
を使用します。これはGPUデバイス1を使用します。
TFは、特に指示がない限り、使用可能なすべてのGPUで使用可能なすべてのメモリを割り当てます。以下に、1つ(または少数)のGPUに固執する4つの方法を示します。
バッシュソリューション。pythonまたはjupyterノートブックを起動する前に、ターミナル/コンソールでCUDA_VISIBLE_DEVICES=0,1
を設定します。
Pythonソリューション。セッションを構築する前に次の2行のコードを実行します。
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"
自動化されたソリューション。以下の方法は、他のスクリプトで使用されていないGPUデバイスを自動的に検出し、CUDA_VISIBLE_DEVICESを設定します。セッションを構築する前にmask_unused_gpus
を呼び出す必要があります。現在のメモリ使用量によってGPUを除外します。これにより、コードを変更したり、コンソールパラメータを設定したりせずに、スクリプトの複数のインスタンスを一度に実行できます。
関数:
import subprocess as sp
import os
def mask_unused_gpus(leave_unmasked=1):
ACCEPTABLE_AVAILABLE_MEMORY = 1024
COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"
try:
_output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:]
memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
available_gpus = [i for i, x in enumerate(memory_free_values) if x > ACCEPTABLE_AVAILABLE_MEMORY]
if len(available_gpus) < leave_unmasked: raise ValueError('Found only %d usable GPUs in the system' % len(available_gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, available_gpus[:leave_unmasked]))
except Exception as e:
print('"nvidia-smi" is probably not installed. GPUs are not masked', e)
mask_unused_gpus(2)
制限:一度に複数のスクリプトを開始すると、セッションの構築時にメモリがすぐに割り当てられないため、衝突が発生する可能性があります。問題がある場合は、元の ソースコード:mask_busy_gpus() のようにランダムバージョンを使用できます。
Tensorflow 2.0はさらに別の方法を提案します:
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
# Restrict TensorFlow to only use the first GPU
try:
tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
except RuntimeError as e:
# Visible devices must be set at program startup
print(e)
pythonスクリプトの先頭に追加することにより、GPUオプション設定を変更できます。
gpu_options = tf.GPUOptions(visible_device_list="0")
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
「0」は、使用するGPUの名前です。端末プロンプトでコマンドnvidia-smiを入力すると、GPUのリストを使用できます。
Kerasを使用すると、これら2つの機能によりCPUまたはGPUを選択でき、GPUの場合は使用するメモリの割合を選択できます。
import os
from keras.backend.tensorflow_backend import set_session
import tensorflow as tf
def set_cpu_option():
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["CUDA_VISIBLE_DEVICES"] = ""
def set_gpu_option(which_gpu, fraction_memory):
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = fraction_memory
config.gpu_options.visible_device_list = which_gpu
set_session(tf.Session(config=config))
return
set_gpu_option("0", 0.9)
# or
set_cpu_option()