5

I am using eigendecomposition in Tensorflow and find that it is extremely slow. Here's the code to show Tensorflow's speed vs numpy and 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')

This gives the following output:

sp: 0.09 s
np: 0.08 s
tf: 5.04 s

Any ideas of what's up here?

Seanslice
  • 153
  • 6
  • For better performance, try wrapping `tf.linalg.eig` in a `@tf.function`. For more details please refer [this](https://www.tensorflow.org/guide/function). Thanks! –  Oct 15 '20 at 14:29

1 Answers1

1

Try wrapping tf.linalg.eig in a @tf.function and you can observe improvement in speed. This is because it converted to graph mode and there will be some optimizations can be done.

Incase of eager mode these may not preformed and it is default behavior in TF 2.x.

You can wrap your code as shown below

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

Please refer comparison w.r.t speed in below

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)

Output:

sp: 0.32 s
np: 0.04 s
tf: 3.62 s
tff: 0.01 s

For more information please refer @tf.function, seeing the speed up and Should I use @tf.function for all functions? by today.