Compartir tecnología

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

Prefacio

Cuando se utiliza TensorFlow 2 para entrenamiento o predicción, la gestión adecuada de la memoria de la GPU es crucial. No administrar y liberar eficazmente la memoria de la GPU puede provocar pérdidas de memoria, lo que puede afectar las tareas informáticas posteriores. En este artículo, exploraremos varias formas de liberar eficazmente la memoria de la GPU, tanto de forma convencional como cuando se fuerza la finalización de una tarea.

1. Métodos convencionales de gestión de memoria de vídeo.

1. Restablecer la imagen predeterminada

Cada vez que ejecuta un nuevo gráfico de TensorFlow, llamando tf.keras.backend.clear_session() para borrar el gráfico actual de TensorFlow y liberar memoria.

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

2. Limite el uso de memoria de la GPU

Al configurar la política de uso de la memoria de video, puede evitar que la memoria de video de la GPU se ocupe demasiado.

• Aumente el uso de memoria de video bajo demanda：

  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)
  

• Limitar el uso de la memoria de vídeo：

  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. Liberar manualmente la memoria de la GPU

Después del entrenamiento o la predicción, use gc El módulo y las funciones de administración de memoria de TensorFlow liberan manualmente la memoria de la GPU.

import tensorflow as tf
import gc

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

4. uso with contexto de gestión de declaraciones

Utilizado en código de entrenamiento o predicción. with declaración para gestionar automáticamente la liberación de recursos.

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. Gestión de la memoria de vídeo al finalizar tareas por la fuerza

A veces necesitamos finalizar por la fuerza la tarea de TensorFlow para liberar memoria de la GPU.En este caso, utilice Pythonmultiprocessing módulo oos Los módulos pueden gestionar los recursos de manera eficiente.

1. uso multiprocessing módulo

Al ejecutar tareas de TensorFlow en procesos separados, se puede finalizar todo el proceso para liberar memoria de video cuando sea necesario.

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. uso os El módulo finaliza el proceso.

Al obtener el ID del proceso y utilizar os Módulo que puede finalizar con fuerza el proceso de 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()

Resumir

Cuando se utiliza TensorFlow 2 para entrenamiento o predicción, es fundamental administrar y liberar correctamente la memoria de la GPU.Al restablecer el mapa predeterminado, limitar el uso de la memoria de video, liberar manualmente la memoria de video y usarwith El contexto de gestión de declaraciones puede evitar eficazmente problemas de pérdida de memoria.Cuando necesite finalizar una tarea a la fuerza, utilicemultiprocessing módulos yos El módulo puede garantizar que la memoria de video se libere a tiempo. A través de estos métodos, se puede garantizar la utilización eficiente de los recursos de la GPU y se puede mejorar la estabilidad y el rendimiento de las tareas informáticas.