高效管理 TensorFlow 2 GPU 显存的实用指南

前言

在使用 TensorFlow 2 进行训练或预测时，合理管理 GPU 显存至关重要。未能有效管理和释放 GPU 显存可能导致显存泄漏，进而影响后续的计算任务。在这篇文章中，我们将探讨几种方法来有效释放 GPU 显存，包括常规方法和强制终止任务时的处理方法。

一、常规显存管理方法

1. 重置默认图

在每次运行新的 TensorFlow 图时，通过调用 tf.keras.backend.clear_session() 来清除当前的 TensorFlow 图和释放内存。

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

2. 限制 GPU 显存使用

通过设置显存使用策略，可以避免 GPU 显存被占用过多。

• 按需增长显存使用：

  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)

• 限制显存使用量：

  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. 手动释放 GPU 显存

在训练或预测结束后，使用 gc 模块和 TensorFlow 的内存管理函数手动释放 GPU 显存。

import tensorflow as tf
import gc

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

4. 使用 with 语句管理上下文

在训练或预测代码中使用 with 语句，可以自动管理资源释放。

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()

二、强制终止任务时的显存管理

有时我们需要强制终止 TensorFlow 任务以释放 GPU 显存。这种情况下，使用 Python 的 multiprocessing 模块或 os 模块可以有效地管理资源。

1. 使用 multiprocessing 模块

通过在单独的进程中运行 TensorFlow 任务，可以在需要时终止整个进程以释放显存。

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. 使用 os 模块终止进程

通过获取进程 ID 并使用 os 模块，可以强制终止 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()

总结

在使用 TensorFlow 2 进行训练或预测时，合理管理和释放 GPU 显存至关重要。通过重置默认图、限制显存使用、手动释放显存以及使用 with 语句管理上下文，可以有效地避免显存泄漏问题。在需要强制终止任务时，使用 multiprocessing 模块和 os 模块可以确保显存得到及时释放。通过这些方法，可以确保 GPU 资源的高效利用，提升计算任务的稳定性和性能。