TensorFlow 2.0: вам нужен декоратор @ tf.function поверх каждой функции?

Question

В TensorFlow 2.0 (все еще альфа-версия прямо сейчас) я знаю, что вы можете использовать декоратор @tf.function, чтобы превратить простой код Python в граф. Нужно ли ставить @tf.function поверх каждой функции каждый раз, когда я этого хочу? И @tf.function рассматривает только следующий функциональный блок?

Answer 1

Хотя decorator @ 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)