Posted on

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

使用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 設備的內存使用情況。