python - 将 Tensorflow 1.5 转换为 Tensorflow 2

Question

如何将此 Tensorflow 1.5 代码转换为 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))))

主要问题是 Tensorflow 2 中删除了 Session 类，并且 compat.v1 层中公开的版本实际上并不兼容。当我使用 Tensorflow 2 运行此代码时，它现在抛出异常:

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

如果我完全放弃使用 Session，那么功能上仍然等效吗？如果我运行:

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)))

在支持 AVX2 的 Tensoflow 1.16 中，它的运行速度明显更快(0.005 秒 vs 30 秒)，而从 pip 安装的普通 Tensorflow 2(不支持 AVX2)也运行得更快一些(30 秒 vs 60 秒)。

为什么使用 Session 会使 Tensorflow 1.16 减慢 6000 倍？

Answer 1

您当然应该利用 TF 2.x 的优势，包括 Eager Execution。不仅非常方便，而且效率更高。

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())

您(大多数情况下)在 TF 2.x 中不再需要任何 session ，Auto Graph 会自动为您完成此操作。

只需使用 @tf.function 注释“main”函数(无需注释诸如 get_values 之类的其他函数，这也会自动发生)。