TensorFlow 特征值分解非常缓慢。

Question

我正在使用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

这里有什么想法吗？

- Seanslice

---

为了获得更好的性能，请尝试在tf.linalg.eig中使用@tf.function进行包装。有关更多详细信息，请参见此链接。谢谢！ - user11530462

Answer 1

尝试在@tf.function中包装tf.linalg.eig，您可以观察到速度的提升。这是因为它转换为图模式，并且可以进行一些优化。

在急切执行模式下，这些可能不会被执行，并且在TF 2.x中是默认行为。

您可以像下面展示的那样包装您的代码。

@tf.function
def oper(A_tf):
    d, v = tf.linalg.eig(A_tf)

请参考下面的速度比较。

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，查看加速和我应该为所有函数使用@tf.function吗？

- user11530462