如何限制 TensorFlow 的 GPU 記憶體分配？

限制TensorFlow GPU 記憶體分配

TensorFlow 的預設行為在啟動時分配全部可用GPU 內存，這在共享計算環境中提出了挑戰。當多個使用者在同一 GPU 上執行並發訓練時，必須防止記憶體消耗過多。

解決方案：GPU 記憶體分數

為了解決這個問題，TensorFlow提供指定要分配的 GPU 記憶體部分的選項。透過在 tf.GPUOptions 物件中設定 per_process_gpu_memory_fraction 字段，您可以限制記憶體消耗。以下是一個範例：

# Restrict memory allocation to 4GB on a 12GB GPU
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)

# Create a session with the GPU options
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

此方法為同一台電腦上所有 GPU 上的目前進程提供了 GPU 記憶體所使用的硬性上限。但請注意，該比例在所有 GPU 上統一應用，並且沒有針對每個 GPU 記憶體分配的選項。

以上是如何限制 TensorFlow 的 GPU 記憶體分配？的詳細內容。更多資訊請關注PHP中文網其他相關文章！