Python tf.config.experimental.get_memory_info用法及代碼示例

獲取所選設備的內存信息，作為字典。

用法

tf.config.experimental.get_memory_info(
    device
)

參數

device 獲取內存信息的設備字符串，例如"GPU:0","TPU:0".看https://www.tensorflow.org/api_docs/python/tf/device用於指定設備字符串。

帶有鍵 'current' 和 'peak' 的字典，分別指定當前和峰值內存使用情況。

拋出

ValueError 未找到具有設備名稱的設備，例如 '"nonexistent"'。
ValueError 設備名稱無效，例如 '"GPU"'、'"CPU:GPU"'、'"CPU:"'。
ValueError 與設備名稱匹配的多個設備。
ValueError 未跟蹤內存統計信息，例如 '"CPU:0"'。

這個函數返回一個包含設備內存使用信息的字典。例如：

if tf.config.list_physical_devices('GPU'):
  # Returns a dict in the form {'current':<current mem usage>,
  #                             'peak':<peak mem usage>}
  tf.config.experimental.get_memory_info('GPU:0')

當前返回以下鍵：

'current' ：設備當前使用的內存，以字節為單位。
'peak' ：設備在程序運行期間使用的峰值內存，以字節為單位。可以使用 tf.config.experimental.reset_memory_stats 重置。

未來可能會添加更多鍵，包括device-specific 鍵。

目前僅支持 GPU 和 TPU。如果在 CPU 設備上調用，則會引發異常。

對於 GPU，TensorFlow 將默認分配所有內存，除非更改為 tf.config.experimental.set_memory_growth 。 dict 僅指定 TensorFlow 實際使用的當前和峰值內存，而不是 TensorFlow 在 GPU 上分配的內存。

相關用法

注：本文由純淨天空篩選整理自tensorflow.org大神的英文原創作品 tf.config.experimental.get_memory_info。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。