Должен ли я использовать @tf.function для всех функций?

Question

Официальный учебник по@tf.function говорит:

Чтобы получить максимальную производительность и сделать вашу модель доступной для развертывания где угодно, используйте tf.function для построения графиков из ваших программ. Благодаря AutoGraph удивительное количество кода Python работает только с tf.function, но все же есть подводные камни, которых следует опасаться.

Основные выводы и рекомендации:

• Не полагайтесь на побочные эффекты Python, такие как мутация объекта или добавление списка.
• tf.function лучше всего работает с операциями TensorFlow, а не с примитивами NumPy или Python.
• В случае сомнений используйте идиому for x in y.

В нем только упоминается, как реализовать@tf.functionаннотированные функции, но не когда их использовать.

Есть ли эвристика о том, как решить, стоит ли мне хотя бы попытаться аннотировать функцию с помощью tf.function? Кажется, что нет причин не делать этого, если только мне не лень убрать побочные эффекты или изменить некоторые вещи, напримерrange()-> tf.range(). Но если я хочу это сделать...

Есть ли причина не использовать @tf.function для всех функций?

Answer 1

TL; DR: это зависит от вашей функции и от того, находитесь ли вы в производстве или в разработке. Не использоватьtf.functionесли вы хотите иметь возможность легко отлаживать свою функцию или если она подпадает под ограничения совместимости кода AutoGraph или tf.v1. Я настоятельно рекомендую посмотреть, как Inside TensorFlow рассказывает об AutoGraph и функциях, а не о сеансах.

Ниже я приведу причины, которые взяты из информации, размещенной в Интернете через Google.

В целом tf.functionдекоратор заставляет функцию компилироваться как вызываемый объект, который выполняет график TensorFlow. Это влечет за собой:

• При необходимости преобразование кода через AutoGraph (включая любые функции, вызываемые из аннотированной функции)
• Отслеживание и выполнение сгенерированного кода графа

Доступна подробная информация о дизайнерских идеях, стоящих за этим.

Преимущества украшения функции tf.function

Общие преимущества

• Более быстрое выполнение, особенно если функция состоит из множества небольших операций (Источник)

Для функций с кодом Python / Использование AutoGraph через tf.function украшение

Если вы хотите использовать AutoGraph, используя tf.functionнастоятельно рекомендуется вместо прямого вызова AutoGraph. Причины для этого включают: зависимости автоматического управления, это требуется для некоторых API, больше кэширования и помощников по исключениям (Источник).

Недостатки декорирования функции tf.function

Общие недостатки

• Если функция состоит только из нескольких дорогостоящих операций, большого ускорения не будет (Источник)

Для функций с кодом Python / Использование AutoGraph через tf.function украшение

• Отсутствие перехвата исключений (должно выполняться в режиме ожидания; вне декорированной функции) (Источник)
• Отладка намного сложнее
• Ограничения из-за скрытых побочных эффектов и потока управления TF

Доступна подробная информация об ограничениях AutoGraph.

Для функций с кодом tf.v1

• Не допускается создание переменных более одного раза в tf.function, но это может быть изменено по мере прекращения использования кода tf.v1 (Источник)

Для функций с кодом tf.v2

• Особых недостатков нет

Примеры ограничений

Создание переменных более одного раза

Не разрешается создавать переменные более одного раза, например v в следующем примере:

@tf.function
def f(x):
    v = tf.Variable(1)
    return tf.add(x, v)

f(tf.constant(2))

# => ValueError: tf.function-decorated function tried to create variables on non-first call.

В следующем коде это смягчается, если убедиться, что self.v создается только один раз:

class C(object):
    def __init__(self):
        self.v = None
    @tf.function
    def f(self, x):
        if self.v is None:
            self.v = tf.Variable(1)
        return tf.add(x, self.v)

c = C()
print(c.f(tf.constant(2)))

# => tf.Tensor(3, shape=(), dtype=int32)

Скрытые побочные эффекты, не фиксируемые AutoGraph

Такие изменения, как self.aв этом примере невозможно скрыть, что приводит к ошибке, поскольку кросс-функциональный анализ еще не выполнен (Источник):

class C(object):
    def change_state(self):
        self.a += 1

    @tf.function
    def f(self):
        self.a = tf.constant(0)
        if tf.constant(True):
            self.change_state() # Mutation of self.a is hidden
        tf.print(self.a)

x = C()
x.f()

# => InaccessibleTensorError: The tensor 'Tensor("add:0", shape=(), dtype=int32)' cannot be accessed here: it is defined in another function or code block. Use return values, explicit Python locals or TensorFlow collections to access it. Defined in: FuncGraph(name=cond_true_5, id=5477800528); accessed from: FuncGraph(name=f, id=5476093776).

Изменения на виду не проблема:

class C(object):
    @tf.function
    def f(self):
        self.a = tf.constant(0)
        if tf.constant(True):
            self.a += 1 # Mutation of self.a is in plain sight
        tf.print(self.a)

x = C()
x.f()

# => 1

Пример ограничения из-за потока управления TF

Этот оператор if приводит к ошибке, потому что для потока управления TF необходимо определить значение else:

@tf.function
def f(a, b):
    if tf.greater(a, b):
        return tf.constant(1)

# If a <= b would return None
x = f(tf.constant(3), tf.constant(2))   

# => ValueError: A value must also be returned from the else branch. If a value is returned from one branch of a conditional a value must be returned from all branches.

Answer 2

Tf .function полезен при создании и использовании вычислительных графиков, их следует использовать при обучении и развертывании, однако он не нужен для большинства ваших функций.

Допустим, мы создаем специальный слой, который будет отдельно от более крупной модели. Мы не хотели бы иметь декоратор tf.function над функцией, которая создает этот слой, потому что это просто определение того, как будет выглядеть слой.

С другой стороны, допустим, что мы собираемся либо сделать прогноз, либо продолжить обучение, используя некоторую функцию. Мы хотели бы иметь декоратор tf.function, потому что на самом деле мы используем вычислительный граф для получения некоторого значения.

Отличным примером может быть построение модели кодировщика-декодера. НЕ помещайте декоратор вокруг функции создания кодировщика или декодера или любого слоя, это только определение того, что он будет делать. ОБЯЗАТЕЛЬНО поместите декоратор вокруг метода "поезда" или "прогнозирования", потому что они фактически будут использовать вычислительный граф для вычислений.

Answer 3

В моем понимании и согласно документации, используя tf.function настоятельно рекомендуется в основном для ускорения вашего кода, поскольку код, заключенный в tf.function будет преобразован в граф, и поэтому есть место для некоторых оптимизаций (например, обрезки операций, сворачивания и т. д.), которые могут не выполняться, когда тот же самый код запускается с нетерпением.

Однако есть также несколько случаев, когда использование tf.functionможет вызвать дополнительные накладные расходы или не привести к заметному увеличению скорости. Один примечательный случай - это когда обернутая функция мала и используется в вашем коде только несколько раз, поэтому накладные расходы на вызов графика могут быть относительно большими. Другой случай - это когда большая часть вычислений уже выполняется на устройстве-ускорителе (например, GPU, TPU), и поэтому ускорение, полученное при вычислении графа, может быть незначительным.

В документации также есть раздел, в котором ускорение обсуждается в различных сценариях, и в начале этого раздела были упомянуты два вышеупомянутых случая:

Просто оберните функцию, использующую тензор, в tf.functionне ускоряет ваш код автоматически. Для небольших функций, вызываемых несколько раз на одной машине, накладные расходы на вызов графа или фрагмента графа могут доминировать во время выполнения. Кроме того, если большая часть вычислений уже выполнялась на ускорителе, например, в стеках сверток с высокой нагрузкой на графический процессор, ускорение графа не будет большим.

Для сложных вычислений графики могут обеспечить значительное ускорение. Это связано с тем, что графики сокращают обмен данными между Python и устройством и обеспечивают некоторое ускорение.

Но в конце концов, если это применимо к вашему рабочему процессу, я думаю, что лучший способ определить это для вашего конкретного варианта использования и среды - профилировать ваш код, когда он запускается в активном режиме (т. Е. Без использования tf.function) по сравнению с тем, когда он выполняется в графическом режиме (т. е. с использованием tf.function широко).