Question

我正在尝试找到一种方法，让 TF 模型在 TF2.3 的图形模式下快速运行。看起来我获得的速度比在 Eager 模式下达到的速度要慢得多，这是出乎意料的。下面是一个简单的例子（我的实际模型比较复杂，但效果差不多）。

您可以将以下内容粘贴到脚本中，并以“2000 1”或“2000 0”作为参数运行（分别表示在 Eager 或 Graph 模式下）。我在图形模式下的执行速度至少慢了 10 倍。正如您所看到的，我使用相同的参数运行了该函数两次，所以大概所有的初步工作（跟踪和上帝知道还有什么）应该只在第一次运行时发生。

谁能告诉我我做错了什么？例如，有没有更快的方法？谢谢！！！

import tensorflow as tf
import datetime
import random
import sys

random.seed(123)

size = int(sys.argv[1])
eager = int(sys.argv[2])

print ("Running TF timing scrupt with size =", size, "eager mode", ("on" if eager else "off"))

if not eager: tf.compat.v1.disable_eager_execution()

a = tf.constant([[random.random() for i in range(size)] for j in range(size)])
b = tf.constant([[random.random() for i in range(size)] for j in range(size)])

@tf.function
def mulX(a, b):
    c = a
    for j in range(40):
        c = c * b
    return tf.math.reduce_sum(c)

if not eager:
    sess = tf.compat.v1.Session()

print ("before 1st run", datetime.datetime.now())
res = tf.math.reduce_sum(mulX(a, b))
if not eager: res = res.eval(session = sess)
print (res)
print ("after 1st run", datetime.datetime.now())

print ("before 2nd run", datetime.datetime.now())
res = tf.math.reduce_sum(mulX(a, b))
if not eager: res = res.eval(session = sess)
print (res)
print ("after 2nd run", datetime.datetime.now())

TF2，图形模式对急切模式非常慢

0 个答案: