点积和矩阵乘法的Tensorflow慢

Question

我目前正在尝试在 Tensorflow 中编写一些线性代数代码，并将性能与 numpy 实现进行比较。我遇到的问题是，对于重复矩阵向量乘法或重复矩阵矩阵乘法，Tensorflow 代码要慢几个数量级。我目前正在使用 CPU 测试没有 GPU 的 Tensorflow。我的目的是在 Tensorflow 中实现我的项目，使其在 CPU 和 GPU 上都可以快速运行，具体取决于用户可用的资源。矩阵大小可以在大约 10x10 到非常大的矩阵之间变化。当使用 tensorflow 来拥有一个通用代码库时，我愿意接受小矩阵的一点性能损失，只要它不像下面那样剧烈。

这里有两个来自 IPython 的计时示例：

%%timeit
dim = 10000
m1 = np.random.uniform(dim)
m2 = np.random.uniform(dim,dim)
tmp = np.dot(m2,m1)
for i in range(200):
    tmp = np.dot(m2,tmp)

这需要 151 µs ± 6.56 µs。

现在使用 TensorFlow：

%%timeit
a1 = tf.random.uniform([dim,dim])
a2 = tf.random.uniform([dim])
tmp = tf.tensordot(a1,a2,1)
for i in range(200):
    tmp = tf.tensordot(a1,tmp,1)

仅需要 1.07 秒 ± 128 毫秒。然后再次评估需要 7.5 秒。

sess = tf.Session()
with sess.as_default():
    out = tmp.eval()

我在使用 matmul 时观察到了同样的趋势。

非常感谢您的任何提示！

彼得

Answer 1

您的代码中有几个问题。

1. np.random.uniform(dim) 和 np.random.uniform(dim,dim) 返回两个浮点数而不是两个矩阵。您可能想使用np.random.rand(dim) 和np.random.rand(dim,dim)

2. 你timeit创建graph，以及生成随机矩阵。我建议只使用 timeit np.dot 和 sess.run()。

以下是tf.matmul和np.dot的基准代码，改编自here

import tensorflow as tf
import numpy as np
import os
import time


n = 10000
dtype = tf.float32
with tf.device("/cpu:0"):
    matrix1 = tf.Variable(tf.ones((n, n), dtype=dtype))
    matrix2 = tf.Variable(tf.ones((n, n), dtype=dtype))
    product = tf.matmul(matrix1, matrix2)

# avoid optimizing away redundant nodes
config = tf.ConfigProto(graph_options=tf.GraphOptions(optimizer_options=tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0)))
sess = tf.Session(config=config)

sess.run(tf.global_variables_initializer())
iters = 10

# pre-warming
sess.run(product.op)

start = time.time()
for i in range(iters):
    sess.run(product.op)
end = time.time()
elapsed = (end - start)
print('tensorflow cpu: %d x %d matmul took: %.2f sec' % (n, n, elapsed/iters))

########### numpy test ###########
np_m1 = np.ones((n,n),dtype=np.float32)
np_m2 = np.ones((n,n),dtype=np.float32)
start = time.time()
for i in range(iters):
    np.dot(np_m1,np_m2)
end = time.time()
elapsed = (end - start)
print('numpy: %d x %d matmul took: %.2f sec' % (n, n, elapsed/iters))

在我的机器上，输出是

# tensorflow cpu: 10000 x 10000 matmul took: 2.79 sec
# numpy: 10000 x 10000 matmul took: 4.36 sec