TensorFlow 2.0：你需要在每个函数之上添加一个 @tf.function 装饰器吗？答案

【问题标题】：TensorFlow 2.0: do you need a @tf.function decorator on top of each function?TensorFlow 2.0：你需要在每个函数之上添加一个 @tf.function 装饰器吗？
【发布时间】：2019-08-04 13:14:00
【问题描述】：

在 TensorFlow 2.0（目前仍为 alpha 版本）中，我知道您可以使用装饰器 @tf.function 将纯 Python 代码转换为图形。每次我想要的时候，我是否必须将@tf.function 放在每个函数的顶部？ @tf.function 是否只考虑以下功能块？

【问题讨论】：

看一看。 github.com/tensorflow/addons/issues/13

标签： python tensorflow tensorflow2.0

【解决方案1】：

虽然装饰器@tf.function 应用于紧随其后的功能块，但任何被它调用的函数也将在图形模式下执行。请参阅Effective TF2 guide 声明：

在 TensorFlow 2.0 中，用户应该将他们的代码重构为更小的函数，这些函数可以根据需要进行调用。一般来说，没有必要用 tf.function 来装饰这些较小的函数；仅使用 tf.function 装饰高级计算 - 例如，训练的一步，或模型的前向传递。

【讨论】：

【解决方案2】：

@tf.function 将 Python 函数转换为其图形表示。

要遵循的模式是定义训练步骤函数，这是计算量最大的函数，并用@tf.function 装饰它。

通常，代码如下：

#model,loss, and optimizer defined previously

@tf.function
def train_step(features, labels):
   with tf.GradientTape() as tape:
        predictions = model(features)
        loss_value = loss(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss_value

for features, labels in dataset:
    lv = train_step(features, label)
    print("loss: ", lv)

【讨论】：