tf.linalg.eigh在GPU上非常慢 - 正常吗？

Posted 2021-03-31

所以我刚刚找到了减慢我在GPU上的代码的罪魁祸首：tf.linalg.eigh()。

这个想法非常简单：我创建 - 比方说 - 87.000 4x4 Hermitian矩阵，并希望得到它们的特征值和特征向量。为此，我有一个形状[87.000,4,4]的占位符matrix，我喂入tf.linalg.eigh(matrix)。我运行会话并将这些矩阵作为输入（矩阵具有数据类型complex64）并希望作为输出特征值。

这需要一个不到0.04秒的8核CPU，而GPU需要19秒 - 其中NumPy需要大约0.4秒。

所以我的问题是：为什么GPU上的tf.linalg.eigh()如此慢，即使给出了大批量的大小。即使一个矩阵的对角化无法有效并行化，在数千个矩阵的情况下GPU仍然应该快得多......

可以以某种方式解决这个问题，还是我必须从GPU切换到CPU进行此操作？

代码：

进口

import numpy as np

from matplotlib.ticker import LinearLocator, FormatStrFormatter

import tensorflow as tf

config = tf.ConfigProto(device_count = {'GPU': 1})

sess = tf.Session(config=config)

import time

建造零件

matrix=tf.placeholder(tf.complex64,shape[None,87,4,4],name="matrix")

eigenval,eigenvec=tf.linalg.eigh(tf.linalg.adjoint(matrix))

init = tf.global_variables_initializer()

sess.run(init)

complex_matrix=np.ones((10000,87,4,4))+1j*np.ones((batch_net,path_length,num_orbits,num_orbits))

运行操作并测量时间

t1=time.time()
sess.run(eigenvec,feed_dict={matrix: complex_matrix, eigenvalues_true: eigenvalues })
print(time.time()-t1)

答案

经过一些小实验，我认为在这种情况下，最好将此操作放在CPU上。关键是PCI-GPU通信是这里的瓶颈，所以你根本无法获得良好的GPU利用率。虽然通过在GPU上使用TF op生成随机martix可以减小这种开销

with tf.device('/device:GPU:0'):
    matrix = tf.random.uniform((87000,4,4), minval=0.1, maxval=0.99, dtype=tf.float32)
    eigenval,eigenvec=tf.linalg.eigh(matrix)

它只允许我的系统减少大约40％的计算时间，这仍然比CPU慢。你也可以尝试将张量分成相等的块，执行linalg.eigh和连接结果，但这也几乎没有任何改进

matrix = tf.random.uniform((87000,4,4), minval=0.1, maxval=0.99, dtype=tf.float32)
result = tf.concat([tf.linalg.eigh(x)[1] for x in tf.split(matrix, 1000, axis=0)], axis=0)

我还注意到在CPU上执行的linalg.eigh缩放大致是对数的，而GPU操作看起来似乎是线性的。希望这可以帮助！

一点点更新。看起来像XLA编译器甚至不支持操作SelfAdjointEigV2，所以这段代码

matrix = tf.random.uniform((87000, 4, 4), minval=0.1, maxval=0.99, dtype=tf.float32)
def xla_test(matrix):
    eigenval, eigenvec = tf.linalg.eigh(matrix)
    return eigenvec

y = xla.compile(xla_test, inputs=[matrix])

抛出'检测到不支持的操作'错误