使用 tf.function 的 Tensorflow 2.0 模型非常慢，每次训练计数更改时都会重新编译.Eager

Question

我有从未编译的 keras 代码构建的模型，并尝试通过自定义训练循环运行它们。

TF 2.0 急切(默认情况下)代码在 CPU(笔记本电脑)上运行大约 30 秒。当我使用包装的 tf.function 调用方法创建 keras 模型时，它的运行速度要慢得多，并且似乎需要很长时间才能启动，尤其是"第一次"时间。

例如，在 tf.function 代码中，10 个样本的初始训练需要 40 秒，而 10 个样本的后续训练需要 2 秒。

在 20 个样本上，初始需要 50 秒，随访需要 4 秒。

1 个样本的第一次训练需要 2 秒，跟进需要 200 毫秒。

所以看起来每次火车调用都在创建一个新图表，其中复杂性随着火车计数而缩放！？

我只是在做这样的事情：

@tf.function
def train(n=10):
step = 0
loss = 0.0
accuracy = 0.0
for i in range(n):
step += 1
d, dd, l = train_one_step(model, opt, data)
tf.print(dd)
with tf.name_scope('train'):
for k in dd:
tf.summary.scalar(k, dd[k], step=step)
if tf.equal(step % 10, 0):
tf.print(dd)
d.update(dd)
return d

其中模型是用@tf.functionkeras.model.Model的，根据示例装饰call方法。

Answer 1

我在这里分析了@tf.function的这种行为，使用Python本机类型。

简而言之：tf.function的设计不会自动将 Python 本机类型装箱以tf.Tensor具有明确定义的dtype的对象。

如果您的函数接受tf.Tensor对象，则在第一次调用函数时，将分析函数，构建图形并与该函数关联。在每个非首次调用中，如果tf.Tensor对象的dtype匹配，则重用图形。

但是在使用 Python 本机类型的情况下，每次使用不同的值调用函数时都会构建 graphg。

简而言之：如果您打算使用@tf.function，请将您的代码设计为在任何地方使用tf.Tensor而不是 Python 变量。

tf.function不是一个包装器，它神奇地加速了一个在 eager 模式下运行良好的函数;它是一个包装器，它需要设计 eager 函数(主体、输入参数、dytpe)，了解一旦创建图形会发生什么，以获得真正的加速。