Converting Tensorflow 1.5 to Tensorflow 2

Viewed 2038

How would you convert this Tensorflow 1.5 code to Tensorflow 2?

import tensorflow as tf
try:
    Session = tf.Session
except AttributeError:
    Session = tf.compat.v1.Session
A = random_normal([10000,10000])
B = random_normal([10000,10000])
with Session() as sess:
    print(sess.run(tf.reduce_sum(tf.matmul(A,B))))

The main problem is that the Session class has been removed in Tensorflow 2, and the version exposed in the compat.v1 layer doesn't actually appear to be compatible. When I run this code with Tensorflow 2, it now throws the exception:

RuntimeError: Attempting to capture an EagerTensor without building a function.

If I drop the use of Session entirely, is that still functionally equivalent? If I run:

import tensorflow as tf
A = random_normal([10000,10000])
B = random_normal([10000,10000])
with Session() as sess:
    print(tf.reduce_sum(tf.matmul(A,B)))

it runs significantly faster (0.005sec vs 30sec) in Tensoflow 1.16 with AVX2 support, whereas stock Tensorflow 2 installed from pip (without AVX2 support) also runs a bit faster (30sec vs 60sec).

Why would the use of Session slow down Tensorflow 1.16 by 6000x?

2 Answers

You certainly should make use of the advantages of TF 2.x, including Eager Execution. It's not only very convenient, but also more efficient.

import tensorflow as tf

def get_values():
  A = tf.random.normal([10_000,10_000])
  B = tf.random.normal([10_000,10_000])
  return A,B

@tf.function
def compute():
  A,B = get_values()
  return tf.reduce_sum(tf.matmul(A,B))

print(compute())

You (mostly) don't need any sessions anymore in TF 2.x, Auto Graph does that automatically for you.

Simply annotate the "main" function with @tf.function (there's no need to annotate further ones like get_values, that happens automatically as well).

Regarding your first question, I was able to get this to run in Colab, with the addition of the "disable_eager_execution()" call - It appears that the "eager execution" mode is on by default in TF 2.0:

# Install TensorFlow
try:
  # %tensorflow_version only exists in Colab.
    %tensorflow_version 2.x
except Exception:
    pass

import numpy as np
import tensorflow as tf
from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution()
print(tf.executing_eagerly())
print(tf.__version__)
matdim = 1000
try:
    Session = tf.Session
except AttributeError:
    Session = tf.compat.v1.Session
A = tf.random.normal([matdim,matdim])
B = tf.random.normal([matdim,matdim])
with Session() as sess:
    print(sess.run(tf.reduce_sum(tf.matmul(A,B))))

I hope this helps.

Related