问题描述
我在Tensorflow中使用特征分解,发现它非常慢。这是显示Tensorflow的速度与numpy和scipy的关系的代码:
import numpy as np
import scipy as sp
import tensorflow as tf
from time import time
A = np.random.randn(400,400)
A_tf = tf.constant(A)
cur = time()
d,v = sp.linalg.eig(A)
print(f'sp: {time() - cur:4.2f} s')
cur = time()
d,v = np.linalg.eig(A)
print(f'np: {time() - cur:4.2f} s')
cur = time()
d,v = tf.linalg.eig(A_tf)
print(f'tf: {time() - cur:4.2f} s')
这将提供以下输出:
sp: 0.09 s
np: 0.08 s
tf: 5.04 s
这里有什么想法吗?
解决方法
尝试将tf.linalg.eig
包装在@tf.function
中,您会发现速度有所提高。
这是因为它已转换为图形模式,并且可以进行一些优化。
在紧急模式下,这些功能可能未执行,这是TF 2.x中的默认行为。
您可以如下所示包装代码
@tf.function
def oper(A_tf):
d,v = tf.linalg.eig(A_tf)
请参阅下面的比较速度w.r.t
import numpy as np
import scipy as sp
import tensorflow as tf
from time import time
A = np.random.randn(400,400)
A_tf = tf.constant(A)
cur = time()
d,v = sp.linalg.eig(A)
print(f'sp: {time() - cur:4.2f} s')
cur = time()
d,v = np.linalg.eigh(A)
print(f'np: {time() - cur:4.2f} s')
d,v = tf.linalg.eig(A_tf)
print(f'tf: {time() - cur:4.2f} s')
@tf.function
def oper(A_tf):
cur = time()
d,v = tf.linalg.eig(A_tf)
print(f'tff: {time() - cur:4.2f} s')
oper(A_tf)
输出:
sp: 0.32 s
np: 0.04 s
tf: 3.62 s
tff: 0.01 s
有关更多信息,请在今天之前参考@tf.function,seeing the speed up和Should I use @tf.function for all functions?。