tensorflow - Python インターフェイスで GPU の数を指定する方法は?

Question

tensorflow を開始するSessionと、dekstop (デバイス 0 と 2) の両方の GPU に 100% のメモリがすぐに割り当てられます。

Python APIを使用して GPU を 1 つだけに制限するにはどうすればよいですか?

Answer 1

たとえば、GPU #0 のみを使用するように制限する大雑把な方法は、次のように定義することです。

export CUDA_VISIBLE_DEVICES=0

Answer 2

一部の操作を特定のデバイスでのみ実行する場合は、次の操作を実行できます。

with tf.device('/gpu:0'):
   ...
   # all your operations

デバイスは次の方法で指定されます: "/cpu:0"、"/gpu:0"、"/gpu:1"。詳細については、GPU デバイスの使用に関するトピックをお読みになることをお勧めします。

いくつかのチュートリアルを見たい場合は、cifar10 チュートリアルをご覧ください。