Partage de technologie

한어Русский языкEnglishFrançaisIndonesianSanskrit日本語DeutschPortuguêsΕλληνικάespañolItalianoSuomalainenLatina

Préface

Lorsque vous utilisez TensorFlow 2 pour l'entraînement ou la prédiction, une bonne gestion de la mémoire GPU est cruciale. Le fait de ne pas gérer et libérer efficacement la mémoire GPU peut entraîner des fuites de mémoire, ce qui peut affecter les tâches informatiques ultérieures. Dans cet article, nous explorerons plusieurs façons de libérer efficacement de la mémoire GPU, à la fois de manière conventionnelle et lorsqu'une tâche est forcée de se terminer.

1. Méthodes conventionnelles de gestion de la mémoire vidéo

1. Réinitialiser l'image par défaut

Chaque fois que vous exécutez un nouveau graphique TensorFlow, en appelant tf.keras.backend.clear_session() pour effacer le graphique TensorFlow actuel et libérer de la mémoire.

import tensorflow as tf
tf.keras.backend.clear_session()

2. Limiter l'utilisation de la mémoire GPU

En définissant la politique d'utilisation de la mémoire vidéo, vous pouvez éviter que la mémoire vidéo du GPU ne soit trop occupée.

• Augmentez l'utilisation de la mémoire vidéo à la demande：

  import tensorflow as tf
  
  gpus = tf.config.experimental.list_physical_devices('GPU')
  if gpus:
      try:
          for gpu in gpus:
              tf.config.experimental.set_memory_growth(gpu, True)
      except RuntimeError as e:
          print(e)
  

• Limiter l'utilisation de la mémoire vidéo：

  import tensorflow as tf
  
  gpus = tf.config.experimental.list_physical_devices('GPU')
  if gpus:
      try:
          tf.config.experimental.set_virtual_device_configuration(
              gpus[0],
              [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)])  # 限制为 4096 MB
      except RuntimeError as e:
          print(e)
  

3. Libérez manuellement la mémoire GPU

Après l'entraînement ou la prédiction, utilisez gc Les modules et les fonctions de gestion de la mémoire de TensorFlow libèrent manuellement la mémoire GPU.

import tensorflow as tf
import gc

tf.keras.backend.clear_session()
gc.collect()

4. Utiliser with Contexte de gestion des instructions

Utilisé dans le code de formation ou de prédiction with instruction pour gérer automatiquement la libération des ressources.

import tensorflow as tf

def train_model():
    with tf.device('/GPU:0'):
        model = tf.keras.models.Sequential([
            tf.keras.layers.Dense(64, activation='relu', input_shape=(32,)),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        model.compile(optimizer='adam', loss='categorical_crossentropy')
        # 假设 X_train 和 y_train 是训练数据
        model.fit(X_train, y_train, epochs=10)

train_model()

2. Gestion de la mémoire vidéo lors de l'arrêt forcé des tâches

Parfois, nous devons terminer de force la tâche TensorFlow pour libérer de la mémoire GPU.Dans ce cas, utilisez Pythonmultiprocessing module ouos Les modules peuvent gérer efficacement les ressources.

1. Utiliser multiprocessing module

En exécutant des tâches TensorFlow dans des processus distincts, l'ensemble du processus peut être supprimé pour libérer de la mémoire vidéo en cas de besoin.

import multiprocessing as mp
import tensorflow as tf
import time

def train_model():
    model = tf.keras.models.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=(32,)),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='categorical_crossentropy')
    # 假设 X_train 和 y_train 是训练数据
    model.fit(X_train, y_train, epochs=10)

if __name__ == '__main__':
    p = mp.Process(target=train_model)
    p.start()
    time.sleep(60)  # 例如，等待60秒
    p.terminate()
    p.join()  # 等待进程完全终止

2. Utiliser os le module termine le processus

En obtenant l'ID du processus et en utilisant os Module qui peut mettre fin de force au processus TensorFlow.

import os
import signal
import tensorflow as tf
import multiprocessing as mp

def train_model():
    pid = os.getpid()
    with open('pid.txt', 'w') as f:
        f.write(str(pid))

    model = tf.keras.models.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=(32,)),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='categorical_crossentropy')
    # 假设 X_train 和 y_train 是训练数据
    model.fit(X_train, y_train, epochs=10)

if __name__ == '__main__':
    p = mp.Process(target=train_model)
    p.start()
    time.sleep(60)  # 例如，等待60秒
    with open('pid.txt', 'r') as f:
        pid = int(f.read())
    os.kill(pid, signal.SIGKILL)
    p.join()

Résumer

Lorsque vous utilisez TensorFlow 2 à des fins d'entraînement ou de prédiction, il est crucial de gérer et de libérer correctement la mémoire GPU.En réinitialisant la carte par défaut, en limitant l'utilisation de la mémoire vidéo, en libérant manuellement la mémoire vidéo et en utilisantwith Le contexte de gestion des instructions peut efficacement éviter les problèmes de fuite de mémoire.Lorsque vous devez mettre fin de force à une tâche, utilisezmultiprocessing modules etos Le module peut garantir que la mémoire vidéo est libérée à temps. Grâce à ces méthodes, une utilisation efficace des ressources GPU peut être garantie et la stabilité et les performances des tâches informatiques peuvent être améliorées.