Question

我需要对角化大量N x N正定矩阵，这些矩阵以TensorFlow张量组织，大小为batch_size x N x N

我一直认为GPU由于具有较高的并行度，因此进行此类计算的速度更快。我显然是错的。

我尝试了一个简单的基准测试，将AMD Ryzen 7 1700 8核与不错的GPU，Nvidia 1060 3GB，Cuda 9.1的结果进行了比较。所有结果都是在Ubuntu 16.04 x64上使用最新的tensorflow 1.12完成的。 MWE基准如下：

import tensorflow as tf
import time

N = 512
batch_size = 2**8

with tf.device('/device:CPU:0'):
    P = tf.random_uniform([batch_size,N,N],dtype=tf.float32, minval=0.0, maxval=1.0)
    P = (tf.linalg.transpose(P) + P) / 2.0
    P = tf.matrix_set_diag(P, tf.constant([0.0], shape=[batch_size,N]))
    L = tf.matrix_diag(tf.reduce_sum(P,axis=1)) - P
    l = tf.linalg.eigvalsh(L)
    # Initialize the global variables
    init = tf.global_variables_initializer()

    with tf.Session() as sess:
        start = time.time()
        print(sess.run(l).shape)
        end = time.time()
        print('Elapsed time = ', end - start)

在我的CPU上，大约需要0.8秒才能解决2 ** 8个独立的特征值问题，而如果我通过设置with tf.device('/device/GPU:0)'切换到GPU，结果将很尴尬：将近16秒。

据我所知，tensorflow依赖于CUDA Blas库的特征值例程。他们这么慢正常吗？

P.S我可以在至强8核处理器与GTX 1070Ti上大致复制相同的结果。

在Tensorflow CPU与GPU中对特征值计算进行基准测试：CPU赢了吗？

0 个答案: