keras - 使用 tf.function 的 Tensorflow 2.0 模型非常慢，并且每次列车计数发生变化时都会重新编译。 Eager 运行速度提高了约 4 倍

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

型号在哪里 keras.model.Model与 @tf.function装修 call方法按照例子。

Answer 1

我分析了@tf.function的这种行为这里Using a Python native type .

简而言之:tf.function的设计不会自动将 Python native 类型装箱到 tf.Tensor具有明确定义的对象 dtype .

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

但是在使用 Python 原生类型的情况下，graphg 正在构建 每次使用不同的值调用该函数时 .

简而言之:设计您的代码以使用 tf.Tensor如果您打算使用 @tf.function，则可以在任何地方而不是 Python 变量.
tf.function不是一个可以神奇地加速在 Eager 模式下运行良好的函数的包装器；是一个包装器，需要设计热切函数(主体、输入参数、dytpes)以了解创建图形后会发生什么，以获得真正的加速。