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

Question

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

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

这在 8 核 CPU 上花费了不到 0.04 秒，而 GPU 需要 19 秒 - NumPy 大约花费了 0.4 秒。

所以我的问题是:为什么 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)

Answer 1

经过一些实验，我认为在这种情况下最好将此操作放在 CPU 上。关键是 PCI-GPU 通信是这里的瓶颈，因此您根本无法获得良好的 GPU 利用率。尽管可以通过在 GPU 上使用 TF 运算生成随机 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 慢得多。 您也可以尝试将张量分割成相等的 block ，执行 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])

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