Wie kann verhindert werden, dass Tensorflow den gesamten GPU-Speicher zuweist?

Ich arbeite in einer Umgebung, in der Rechenressourcen gemeinsam genutzt werden, dh wir haben ein paar Servermaschinen, die jeweils mit ein paar Nvidia Titan X-GPUs ausgestattet sind.

Bei kleinen bis mittelgroßen Modellen reichen die 12 GB des Titan X normalerweise aus, damit 2–3 Personen gleichzeitig auf derselben GPU trainieren können. Wenn die Modelle so klein sind, dass ein einzelnes Modell nicht alle Recheneinheiten der GPU voll ausnutzt, kann dies sogar zu einer Beschleunigung im Vergleich zum Ausführen eines Trainingsvorgangs nach dem anderen führen. Selbst in Fällen, in denen der gleichzeitige Zugriff auf die GPU die individuelle Trainingszeit verlangsamt, ist es immer noch schön, die Flexibilität zu haben, mehrere Benutzer gleichzeitig auf der GPU trainieren zu lassen.

Das Problem mit TensorFlow ist, dass es standardmäßig die volle Menge an verfügbarem GPU-Speicher zuweist, wenn es gestartet wird. Selbst für ein kleines zweischichtiges neuronales Netzwerk sehe ich, dass alle 12 GB des GPU-Speichers aufgebraucht sind.

Gibt es eine Möglichkeit, TensorFlow dazu zu bringen, nur beispielsweise 4 GB GPU-Speicher zuzuweisen, wenn man weiß, dass dies für ein bestimmtes Modell ausreicht?

Sie können den Anteil des zuzuweisenden GPU-Speichers festlegen, wenn Sie a erstellen tf.Session durch Passieren von a tf.GPUOptions als Teil der optionalen config Streit:

# Assume that you have 12GB of GPU memory and want to allocate ~4GB:
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)

sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

Das per_process_gpu_memory_fraction fungiert als feste Obergrenze für die Menge an GPU-Speicher, die vom Prozess auf jeder GPU auf demselben Computer verwendet wird. Derzeit wird dieser Bruchteil einheitlich auf alle GPUs auf demselben Computer angewendet; Es gibt keine Möglichkeit, dies pro GPU festzulegen.

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

https://github.com/tensorflow/tensorflow/issues/1578

Für TensorFlow 2.0 und 2.1 (Dokumente):

import tensorflow as tf
tf.config.gpu.set_per_process_memory_growth(True)

Für TensorFlow 2.2+ (Dokumente):

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)

Die Dokumentation listet auch einige weitere Methoden auf:

• Umgebungsvariable setzen TF_FORCE_GPU_ALLOW_GROWTH zu true.
• Verwenden tf.config.experimental.set_virtual_device_configuration um ein hartes Limit für ein virtuelles GPU-Gerät festzulegen.

Hier ein Auszug aus dem Buch Deep Learning with TensorFlow

In einigen Fällen ist es wünschenswert, dass der Prozess nur eine Teilmenge des verfügbaren Speichers zuordnet oder die Speichernutzung nur so weit erhöht, wie sie vom Prozess benötigt wird. TensorFlow bietet zwei Konfiguration Optionen in der Sitzung, um dies zu steuern. Die erste ist die allow_growth Option, die versucht, nur so viel GPU-Speicher basierend auf Laufzeitzuweisungen zuzuweisen, beginnt mit der Zuweisung von sehr wenig Speicher, und wenn Sitzungen ausgeführt werden und mehr GPU-Speicher benötigt wird, erweitern wir den vom TensorFlow-Prozess benötigten GPU-Speicherbereich.

1) Wachstum zulassen: (flexibler)

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)

Die zweite Methode ist per_process_gpu_memory_fraction Option, die den Bruchteil der Gesamtspeichermenge bestimmt each sichtbare GPU sollte zugewiesen werden. Notiz: Es ist keine Freigabe von Speicher erforderlich, es kann sogar die Speicherfragmentierung verschlimmern, wenn dies erledigt ist.

2) Ordnen Sie festen Speicher zu:

Nur zuzuordnen 40% des Gesamtspeichers jeder GPU durch:

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config, ...)

Notiz:
Das ist jedoch nur nützlich, wenn Sie wirklich die Menge an verfügbarem GPU-Speicher für den TensorFlow-Prozess binden möchten.

Verwenden Sie für Tensorflow Version 2.0 und 2.1 das folgende Snippet:

 import tensorflow as tf
 gpu_devices = tf.config.experimental.list_physical_devices('GPU')
 tf.config.experimental.set_memory_growth(gpu_devices[0], True)

Für frühere Versionen folgendes Snippet hat bei mir funktioniert:

import tensorflow as tf
tf_config=tf.ConfigProto()
tf_config.gpu_options.allow_growth=True
sess = tf.Session(config=tf_config)

Alle obigen Antworten gehen von der Ausführung mit a aus sess.run() aufrufen, was in neueren Versionen von TensorFlow eher zur Ausnahme als zur Regel wird.

Bei Verwendung der tf.Estimator Framework (TensorFlow 1.4 und höher) die Möglichkeit, den Bruch an den implizit erstellten weiterzugeben MonitoredTrainingSession ist,

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
trainingConfig = tf.estimator.RunConfig(session_config=conf, ...)
tf.estimator.Estimator(model_fn=..., 
                       config=trainingConfig)

Ähnlich im Eager-Modus (TensorFlow 1.5 und höher),

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
tfe.enable_eager_execution(config=conf)

Änderung: 04.11.2018
Als Beispiel, wenn Sie verwenden sollen tf.contrib.gan.traindann können Sie etwas Ähnliches wie unten verwenden:

tf.contrib.gan.gan_train(........, config=conf)

Sie können verwenden

TF_FORCE_GPU_ALLOW_GROWTH=true

in deinen Umgebungsvariablen.

Im Tensorfluss Code:

bool GPUBFCAllocator::GetAllowGrowthValue(const GPUOptions& gpu_options) {
  const char* force_allow_growth_string =
      std::getenv("TF_FORCE_GPU_ALLOW_GROWTH");
  if (force_allow_growth_string == nullptr) {
    return gpu_options.allow_growth();
}