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

时间:2019-04-12 15:42:48

标签: python tensorflow

因此,我刚发现导致GPU上的代码运行速度变慢的罪魁祸首:tf.linalg.eigh()

这个想法很简单:我创建-假设-87.000 4x4 Hermitian矩阵,并且想要获得它们的特征值和特征向量。为此,我有一个形状为[87.000,4,4]的占位符matrix,我将其放入tf.linalg.eigh(matrix)中。我运行Session并将这些矩阵作为输入(矩阵的数据类型为complex64)作为输入,并希望将特征值作为输出。

这使用了一个少于0.04s的8核CPU,而GPU需要19s-在NumPy上花费了大约0.4s。

所以我的问题是:为什么tf.linalg.eigh()在GPU上这么慢,即使一个批处理量很大。即使不能有效地并行化一个矩阵的对角化,但在成千上万个矩阵的情况下,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

tf部分的建造

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)

1 个答案:

答案 0 :(得分:0)

经过一些试验,我认为在这种情况下最好将此操作放置在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])

引发“检测到不受支持的操作”错误