限制在Tensorflow跑模型時使用的GPU的記憶體上限?

Claire Chang

2023-01-04 pm 5:06

使用tensorflow-gpu結果耗一大堆MEMORY是為什麼

使用 TensorFlow GPU 版本會耗費較多的記憶體，這是正常的。因為 GPU 設備有自己的內存，我們可以使用 GPU 設備加速計算。但是，這意味著 GPU 設備的內存也必須足夠大，以便容納計算所需的資料。
如果GPU的記憶體不夠大，則tensorflow會改將原本要放在GPU的記憶體內的資料放到CPU的記憶體裡面，若是CPU的記憶體也不足夠大，則很有可能會導致程式死掉(因為記憶體不足夠)

改善方案

可考慮的改善方向有以下三點:

模型太大，超出了 GPU 設備的內存限制=>可以考慮使用更大的 GPU 設備或對模型進行優化，以減少模型的大小。(請參見: 如何縮小Tensorflow運算模型時使用的記憶體大小)
程式碼中存在記憶體泄漏。請檢查程式碼，確保正確釋放不再使用的記憶體。
GPU 設備的驅動程序版本過舊或損壞

其他方法

請參考這篇文章: https://starriet.medium.com/tensorflow-2-0-wanna-limit-gpu-memory-10ad474e2528
第一個選項：設置set_memory_growth為真。

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)

第二個選項：將第一個 GPU 的內存使用量限制為 1024MB。gpus只需根據需要更改和的索引memory_limit即可。

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=1024)])
  except RuntimeError as e:
    print(e)

我使用了第二個方法成功解決佔用太大記憶體的問題

使用Task Manager查看GPU使用狀況

使用 Windows 任務管理器檢查 GPU 設備的內存使用情況。要查看 GPU 內存使用情況，請按照以下步驟操作：

在 Windows 任務欄中，右鍵單擊「資源監視器」圖示。
在「資源監視器」窗口中，展開「性能」窗格。
在「性能」窗格中，展開「視覺效果」窗格。
在「視覺效果」窗格中，展開「DirectX 內存使用量」窗格。

「DirectX 內存使用量」只包含 GPU 設備上運行的應用程序所使用的內存，不包含 GPU 設備本身的內存。

如果需要更詳細的信息，可以使用 NVIDIA System Monitor 或其他第三方軟件工具檢查 GPU 設備的內存使用情況。

Claire's Blog