TensorFlow 2.0：在每个函数之上都需要一个@ tf.function装饰器吗？

Question

我知道在TensorFlow 2.0（现在仍是Alpha版本）中，您可以使用装饰器@tf.function将纯Python代码转换为图形。我@tf.function每次需要时都必须在每个功能之上放吗？并且是否@tf.function仅考虑以下功能块？

Answer 1

虽然装饰器@tf.function 适用于紧随其后的功能块，但它调用的任何函数也将以图形模式执行。请参阅Effective TF2 指南，其中指出：

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

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