Devo usar @ tf.function para todas as funções?

Um tutorial oficial sobre @tf.functiondiz:

Para obter o máximo desempenho e tornar seu modelo implantável em qualquer lugar, use tf.function para criar gráficos de seus programas. Graças ao AutoGraph, uma quantidade surpreendente de código Python funciona apenas com o tf.function, mas ainda há algumas armadilhas com as quais desconfiar.

Os principais tópicos e recomendações são:

• Não confie nos efeitos colaterais do Python, como mutação de objeto ou anexos à lista.
• O tf.function funciona melhor com operações TensorFlow, em vez de operações NumPy ou primitivas Python.
• Em caso de dúvida, use o idioma for x no y.

Ele menciona apenas como implementar @tf.functionfunções anotadas, mas não quando usá-las.

Existe uma heurística sobre como decidir se devo pelo menos tentar anotar uma função tf.function? Parece que não há razões para não fazê-lo, a menos que eu tenha preguiça de remover efeitos colaterais ou mudar algumas coisas como range()-> tf.range(). Mas se eu estiver disposto a fazer isso ...

Existe alguma razão para não usar @tf.functionpara todas as funções?

TLDR: Depende da sua função e se você está em produção ou desenvolvimento. Não use tf.functionse você deseja poder depurar sua função facilmente ou se ela se enquadra nas limitações do AutoGraph ou da compatibilidade de código tf.v1. Eu recomendo assistir as palestras do Inside TensorFlow sobre AutoGraph e Funções, não Sessões .

A seguir, detalharei os motivos, todos extraídos de informações disponibilizadas on-line pelo Google.

Em geral, o tf.functiondecorador faz com que uma função seja compilada como uma chamada que executa um gráfico TensorFlow. Isto implica:

• Conversão do código através do AutoGraph, se necessário (incluindo quaisquer funções chamadas de uma função anotada)
• Rastreando e executando o código gráfico gerado

Há informações detalhadas disponíveis sobre as idéias de design por trás disso.

Benefícios de decorar uma função com tf.function

Benefícios gerais

• Execução mais rápida , especialmente se a função consistir em muitas operações pequenas (Origem)

Para funções com código Python / Uso do AutoGraph via tf.functiondecoração

Se você deseja usar o AutoGraph, tf.functioné altamente recomendável usar diretamente o AutoGraph. Os motivos para isso incluem: Dependências de controle automático, são necessárias para algumas APIs, mais cache e auxiliares de exceção (Origem) .

Desvantagens de decorar uma função com tf.function

Desvantagens gerais

• Se a função consistir apenas em poucas operações caras, não haverá muita aceleração (origem)

Para funções com código Python / Uso do AutoGraph via tf.functiondecoração

• Nenhuma captura de exceção (deve ser feita no modo ansioso; fora da função decorada) (Fonte)
• Depurar é muito mais difícil
• Limitações devido a efeitos colaterais ocultos e fluxo de controle de TF

Informações detalhadas sobre as limitações do AutoGraph estão disponíveis.

Para funções com código tf.v1

• Não é permitido criar variáveis ​​mais de uma vez tf.function, mas isso está sujeito a alterações à medida que o código tf.v1 é eliminado gradualmente (Origem)

Para funções com código tf.v2

• Sem desvantagens específicas

Exemplos de limitações

Criando variáveis ​​mais de uma vez

Não é permitido criar variáveis ​​mais de uma vez, como vno exemplo a seguir:

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

No código a seguir, isso é atenuado, garantindo que ele self.vseja criado apenas uma vez:

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)

Efeitos colaterais ocultos não capturados pelo AutoGraph

Alterações como self.aneste exemplo não podem ser ocultadas, o que leva a um erro, pois a análise de função cruzada ainda não foi realizada (ainda) (origem) :

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

Mudanças à vista de todos não são um problema:

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

Exemplo de limitação devido ao fluxo de controle de TF

Esta instrução if leva a um erro porque o valor para else precisa ser definido para o fluxo de controle TF:

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

O tf.function é útil na criação e no uso de gráficos computacionais, eles devem ser usados ​​no treinamento e na implantação, no entanto, não é necessário para a maioria das suas funções.

Digamos que estamos construindo uma camada especial que fará parte de um modelo maior. Não gostaríamos de ter o decorador tf.function acima da função que constrói essa camada, porque é apenas uma definição da aparência da camada.

Por outro lado, digamos que vamos fazer uma previsão ou continuar nosso treinamento usando alguma função. Gostaríamos de ter o decorador tf.function porque, na verdade, estamos usando o gráfico computacional para obter algum valor.

Um ótimo exemplo seria construir um modelo de codificador-decodificador. NÃO coloque o decorador em torno da função, crie o codificador ou decodificador ou qualquer camada, que é apenas uma definição do que ele fará. Coloque o decorador em torno do método "treinar" ou "prever", porque eles realmente vão usar o gráfico computacional para computação.