tensorflow - 如何使用 tf.matmul 执行高效的稀疏矩阵乘法？

Question

我正在尝试使用 tf.matmul() 执行稀疏矩阵乘法。

但是，推理速度比密集矩阵乘法慢得多。

根据 tf.sparse_matmul() 中的描述:

• 在一个平台上使用此乘法与密集矩阵相乘的盈亏平衡点是稀疏矩阵中有 30% 的零值。

因此，我创建了具有 7/8 个零值的稀疏矩阵。

这是我的代码:

import tensorflow as tf
import numpy as np
import time
a = tf.Variable(np.arange(1000).reshape(250,4) ,dtype=tf.float32) #dense matrix
b = tf.Variable(np.array([0,0,0,0,0,0,0,1],dtype=np.float32).reshape(4,2),dtype=tf.float32) # sparse matrix
c = tf.matmul(a,b,b_is_sparse=True) # do the sparse matrix multiplication

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    num_iteration = 5000
    num_burnin = 50
    duration = 0

    for i in range(num_iteration+num_burnin):
        startTime  = time.time()
        result = sess.run(c)
        endTime = time.time()
        if i > num_burnin :
            duration+= endTime-startTime

   print(" Average Inference Time = %.3f  ms"%(duration*1000/num_iteration))

我设置“b_is_sparse=True”来执行稀疏矩阵乘法，在我的 GeForce GTX 960M 上大约需要 0.380 毫秒。

但是，如果我设置“b_is_sparse=False”来执行密集矩阵乘法，则大约需要 0.280 毫秒。

我试过用tf.sparse_tensor_dense_matmul和tf.embedding_lookup_sparse来做稀疏矩阵乘法，但是推理速度还是比稠密矩阵乘法慢。

我的代码或执行稀疏矩阵乘法的其他方式有问题吗？

任何建议将不胜感激!!

Answer 1

相对性能取决于许多因素。稀疏乘法可以比密集矩阵的密集乘法更快(希望如此)，但你是对的，它也可能更慢。

一方面，这取决于矩阵的大小。

这是两个方阵相乘的结果，一个是随机的，一个是用零填充，并记录了密集和备用乘法的计算时间。

如您所见，即使矩阵完全为零，对于较小的矩阵大小，稀疏乘法也可能比密集乘法慢——实际上，对于 120x120 的矩阵，速度几乎要慢三倍。在我的计算机上进行的这个实验中，稀疏矩阵乘法开始接管大约 700x700 的大小，最终速度快了大约 2 倍。当然 YMMV 取决于您的配置。