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

Question

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

Answer 1

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

Answer 2

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

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