控制 GPU 記憶體分配

預設情況下,TensorFlow 預先分配 GPU 卡的整個記憶體(這可能會導致 CUDA_OUT_OF_MEMORY 警告)。

要改變這一點,就有可能

  • 使用 per_process_gpu_memory_fraction 配置選項更改預分配的記憶體百分比,

    介於 0 和 1 之間的值,指示
    為每個程序預分配的可用 GPU 記憶體的哪一部分。1 表示
    預分配所有 GPU 記憶體,0.5 表示該程序
    分配~50%的可用 GPU 記憶體。

  • 使用 allow_growth config 選項禁用預分配。記憶體分配隨著使用量的增長而增長。

    如果為 true,則分配器不預先分配整個指定的
    GPU 記憶體區域,而是從小開始並根據需要增長。

例如:

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
sess = tf.Session(config=config) as sess:

要麼

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess= tf.Session(config=config):

在配置選項的資訊在這裡