上矩陣乘法減少GPU存儲器使用在Tensorflow

Question

我有形狀的2點矩陣

a: [1, 3, 3, 512, 512] 
    b: [1, 14, 14, 3, 3, 512,1] 

由

output = tf.multiply(a,b) 

獲得的輸出，並單獨塑造

output: [1, 14, 14, 3, 3, 512, 512] 

該操作大約需要5GB的GPU內存。我想知道是否有任何的方式來減少這種內存消耗（因爲我將有在未來同樣的操作更大的矩陣）。

我通過轉換兩個a和b到半精度嘗試瞭解決方法。這也無助於減少內存使用量。

這將是偉大知道是否有任何Tensorflow技巧，我可以利用這樣做要做到這一點。

謝謝！

Answer 1

您可以設置tf.GPUOption控制GPU內存使用情況。默認情況下，即使不構建任何計算圖，tf也會佔用大量GPU內存。

更多見https://www.tensorflow.org/api_docs/python/tf/GPUOptions。通常情況下，你可能要使用allow_growth或per_process_gpu_memory_fraction財產。