Tensorflow特征值分解极慢答案

【问题标题】：Tensorflow eigenvalue decomposition is extremely slowTensorflow特征值分解极慢
【发布时间】：2020-09-04 18:34:31
【问题描述】：

我在 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 中。更多详情请参考this。谢谢！

标签： python tensorflow eigenvalue eigenvector

【解决方案1】：

尝试将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?。

【讨论】：