前言
在使用 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 tfgpus = 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 tfgpus = 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 MBexcept RuntimeError as e:print(e)
3. 手動釋放 GPU 顯存
在訓練或預測結束后，使用 gc 模塊和 TensorFlow 的內存管理函數手動釋放 GPU 顯存。
import tensorflow as tf
import gctf.keras.backend.clear_session()
gc.collect()
4. 使用 with 語句管理上下文
在訓練或預測代碼中使用 with 語句，可以自動管理資源釋放。
import tensorflow as tfdef 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 timedef 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 mpdef 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 資源的高效利用，提升計算任務的穩定性和性能。
