减少Tensorflow中矩阵乘法的GPU内存使用

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属性。