高效管理 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