Migliori prestazioni con tf.function

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza l'origine su GitHubScarica quaderno

In TensorFlow 2, l'esecuzione desiderosa è attivata per impostazione predefinita. L'interfaccia utente è intuitiva e flessibile (l'esecuzione di operazioni una tantum è molto più semplice e veloce), ma ciò può andare a scapito delle prestazioni e della distribuzione.

Puoi usare tf.function per creare grafici dai tuoi programmi. È uno strumento di trasformazione che crea grafici del flusso di dati indipendenti da Python dal codice Python. Questo ti aiuterà a creare modelli performanti e portatili ed è necessario utilizzare SavedModel .

Questa guida ti aiuterà a concettualizzare come funziona tf.function sotto il cofano, in modo da poterlo utilizzare in modo efficace.

I principali takeaway e raccomandazioni sono:

  • Esegui il debug in modalità desiderosa, quindi decora con @tf.function .
  • Non fare affidamento sugli effetti collaterali di Python come la mutazione di oggetti o le aggiunte di elenchi.
  • tf.function funziona meglio con TensorFlow ops; Le chiamate NumPy e Python vengono convertite in costanti.

Impostare

import tensorflow as tf

Definisci una funzione di supporto per dimostrare i tipi di errori che potresti riscontrare:

import traceback
import contextlib

# Some helper code to demonstrate the kinds of errors you might encounter.
@contextlib.contextmanager
def assert_raises(error_class):
  try:
    yield
  except error_class as e:
    print('Caught expected exception \n  {}:'.format(error_class))
    traceback.print_exc(limit=2)
  except Exception as e:
    raise e
  else:
    raise Exception('Expected {} to be raised but no error was raised!'.format(
        error_class))

Nozioni di base

Utilizzo

Una Function che definisci (ad esempio applicando il decoratore di @tf.function ) è proprio come un'operazione di base di TensorFlow: puoi eseguirla avidamente; puoi calcolare i gradienti; e così via.

@tf.function  # The decorator converts `add` into a `Function`.
def add(a, b):
  return a + b

add(tf.ones([2, 2]), tf.ones([2, 2]))  #  [[2., 2.], [2., 2.]]

<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[2., 2.],
       [2., 2.]], dtype=float32)>

v = tf.Variable(1.0)
with tf.GradientTape() as tape:
  result = add(v, 1.0)
tape.gradient(result, v)

<tf.Tensor: shape=(), dtype=float32, numpy=1.0>

È possibile utilizzare Function s all'interno di altre Function s.

@tf.function
def dense_layer(x, w, b):
  return add(tf.matmul(x, w), b)

dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))

<tf.Tensor: shape=(3, 2), dtype=float32, numpy=
array([[3., 3.],
       [3., 3.],
       [3., 3.]], dtype=float32)>

Function s può essere più veloce del codice desideroso, specialmente per i grafici con molte piccole operazioni. Ma per i grafici con poche operazioni costose (come le convoluzioni), potresti non vedere molto accelerazione.

import timeit
conv_layer = tf.keras.layers.Conv2D(100, 3)

@tf.function
def conv_fn(image):
  return conv_layer(image)

image = tf.zeros([1, 200, 200, 100])
# Warm up
conv_layer(image); conv_fn(image)
print("Eager conv:", timeit.timeit(lambda: conv_layer(image), number=10))
print("Function conv:", timeit.timeit(lambda: conv_fn(image), number=10))
print("Note how there's not much difference in performance for convolutions")

Eager conv: 0.006058974999177735
Function conv: 0.005791576000774512
Note how there's not much difference in performance for convolutions

Tracciamento

Questa sezione illustra come Function la funzione nascosta, inclusi i dettagli di implementazione che potrebbero cambiare in futuro . Tuttavia, una volta capito perché e quando si verifica il tracciamento, è molto più semplice utilizzare tf.function in modo efficace!

Che cos'è il "tracciamento"?

Una Function esegue il programma in un grafico TensorFlow . Tuttavia, un tf.Graph non può rappresentare tutte le cose che scriveresti in un programma TensorFlow desideroso. Ad esempio, Python supporta il polimorfismo, ma tf.Graph richiede che i suoi input abbiano un tipo di dati e una dimensione specificati. Oppure puoi eseguire attività secondarie come leggere argomenti della riga di comando, generare un errore o lavorare con un oggetto Python più complesso; nessuna di queste cose può essere eseguita in un tf.Graph .

Function colma questa lacuna separando il codice in due fasi:

1) Nella prima fase, denominata " tracing ", Function crea un nuovo tf.Graph . Il codice Python viene eseguito normalmente, ma tutte le operazioni di TensorFlow (come l'aggiunta di due Tensor) sono posticipate : vengono acquisite da tf.Graph e non vengono eseguite.

2) Nella seconda fase viene eseguito un tf.Graph che contiene tutto ciò che è stato differito nella prima fase. Questa fase è molto più veloce della fase di tracciamento.

A seconda dei suoi input, Function non eseguirà sempre la prima fase quando viene chiamata. Vedere "Regole di tracciamento" di seguito per avere un'idea migliore di come si ottiene tale determinazione. Saltare la prima fase ed eseguire solo la seconda fase è ciò che offre le elevate prestazioni di TensorFlow.

Quando Function decide di tracciare, la fase di traccia è immediatamente seguita dalla seconda fase, quindi chiamando Function crea ed esegue sia tf.Graph . Successivamente vedrai come puoi eseguire solo la fase di tracciamento con get_concrete_function .

Quando si passano argomenti di tipi diversi in un Function , vengono eseguite entrambe le fasi:

@tf.function
def double(a):
  print("Tracing with", a)
  return a + a

print(double(tf.constant(1)))
print()
print(double(tf.constant(1.1)))
print()
print(double(tf.constant("a")))
print()

Tracing with Tensor("a:0", shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)

Tracing with Tensor("a:0", shape=(), dtype=float32)
tf.Tensor(2.2, shape=(), dtype=float32)

Tracing with Tensor("a:0", shape=(), dtype=string)
tf.Tensor(b'aa', shape=(), dtype=string)

Si noti che se si chiama ripetutamente una Function con lo stesso tipo di argomento, TensorFlow salterà la fase di traccia e riutilizzerà un grafico tracciato in precedenza, poiché il grafico generato sarebbe identico.

# This doesn't print 'Tracing with ...'
print(double(tf.constant("b")))

tf.Tensor(b'bb', shape=(), dtype=string)

Puoi usare pretty_printed_concrete_signatures() per vedere tutte le tracce disponibili:

print(double.pretty_printed_concrete_signatures())

double(a)
  Args:
    a: int32 Tensor, shape=()
  Returns:
    int32 Tensor, shape=()

double(a)
  Args:
    a: float32 Tensor, shape=()
  Returns:
    float32 Tensor, shape=()

double(a)
  Args:
    a: string Tensor, shape=()
  Returns:
    string Tensor, shape=()

Finora, hai visto che tf.function crea un livello di invio dinamico memorizzato nella cache sulla logica di tracciamento del grafico di TensorFlow. Per essere più precisi sulla terminologia:

  • Un tf.Graph è la rappresentazione grezza, indipendente dal linguaggio e portatile di un calcolo TensorFlow.
  • Una ConcreteFunction avvolge un tf.Graph .
  • Una Function gestisce una cache di ConcreteFunction e seleziona quella giusta per i tuoi input.
  • tf.function wrapping di una funzione Python, restituendo un oggetto Function .
  • Tracing crea un tf.Graph e lo avvolge in una ConcreteFunction , nota anche come traccia.

Regole di tracciamento

Una Function determina se riutilizzare una ConcreteFunction tracciata calcolando una chiave cache da args e kwargs di un input. Una chiave cache è una chiave che identifica una ConcreteFunction in base agli input args e kwargs della Function chiamata, secondo le seguenti regole (che possono cambiare):

  • La chiave generata per un tf.Tensor è la sua forma e dtype.
  • La chiave generata per un tf.Variable è un ID variabile univoco.
  • La chiave generata per una primitiva Python (come int , float , str ) è il suo valore.
  • La chiave generata per nidificati dict s, list s, tuple s, namedtuple s e attr s è la tupla appiattita di leaf-keys (vedi nest.flatten ). (Come risultato di questo appiattimento, la chiamata di una funzione concreta con una struttura di nidificazione diversa da quella utilizzata durante la traccia risulterà in un TypeError).
  • Per tutti gli altri tipi di Python la chiave è univoca per l'oggetto. In questo modo una funzione o un metodo viene tracciato indipendentemente per ogni istanza con cui viene chiamato.

Controllo del ritracciamento

Il ritracciamento, ovvero quando la tua Function crea più di una traccia, aiuta a garantire che TensorFlow generi grafici corretti per ogni set di input. Tuttavia, il tracciamento è un'operazione costosa! Se la tua Function ritraccia un nuovo grafico per ogni chiamata, scoprirai che il tuo codice viene eseguito più lentamente che se non avessi usato tf.function .

Per controllare il comportamento di traccia, puoi utilizzare le seguenti tecniche:

  • Specifica input_signature in tf.function per limitare la traccia.
@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def next_collatz(x):
  print("Tracing with", x)
  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)

print(next_collatz(tf.constant([1, 2])))
# You specified a 1-D tensor in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([[1, 2], [3, 4]]))

# You specified an int32 dtype in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([1.0, 2.0]))

Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([4 1], shape=(2,), dtype=int32)
Caught expected exception 
  <class 'ValueError'>:
Caught expected exception 
  <class 'ValueError'>:
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/1851403433.py", line 9, in <module>
    next_collatz(tf.constant([[1, 2], [3, 4]]))
ValueError: Python inputs incompatible with input_signature:
  inputs: (
    tf.Tensor(
[[1 2]
 [3 4]], shape=(2, 2), dtype=int32))
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.int32, name=None)).
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/1851403433.py", line 13, in <module>
    next_collatz(tf.constant([1.0, 2.0]))
ValueError: Python inputs incompatible with input_signature:
  inputs: (
    tf.Tensor([1. 2.], shape=(2,), dtype=float32))
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.int32, name=None)).

  • Specificare una dimensione [Nessuno] in tf.TensorSpec per consentire flessibilità nel riutilizzo delle tracce.

    Poiché TensorFlow abbina i tensori in base alla loro forma, l'utilizzo di una dimensione None come carattere jolly consentirà alle Function di riutilizzare le tracce per input di dimensioni variabili. L'input di dimensioni variabili può verificarsi se si dispone di sequenze di lunghezza diversa o immagini di dimensioni diverse per ciascun batch (consultare ad esempio i tutorial Transformer e Deep Dream ).

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def g(x):
  print('Tracing with', x)
  return x

# No retrace!
print(g(tf.constant([1, 2, 3])))
print(g(tf.constant([1, 2, 3, 4, 5])))

Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([1 2 3], shape=(3,), dtype=int32)
tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int32)

  • Trasmetti argomenti Python a Tensors per ridurre il ritracciamento.

    Spesso, gli argomenti Python vengono utilizzati per controllare gli iperparametri e le costruzioni di grafi, ad esempio num_layers=10 o training=True o nonlinearity='relu' . Quindi, se l'argomento Python cambia, ha senso che dovresti ripercorrere il grafico.

    Tuttavia, è possibile che un argomento Python non venga utilizzato per controllare la costruzione del grafico. In questi casi, una modifica del valore Python può attivare un inutile ritracciamento. Prendi, ad esempio, questo ciclo di addestramento, che AutoGraph srotolerà dinamicamente. Nonostante le tracce multiple, il grafico generato è in realtà identico, quindi il ritracciamento non è necessario.

def train_one_step():
  pass

@tf.function
def train(num_steps):
  print("Tracing with num_steps = ", num_steps)
  tf.print("Executing with num_steps = ", num_steps)
  for _ in tf.range(num_steps):
    train_one_step()

print("Retracing occurs for different Python arguments.")
train(num_steps=10)
train(num_steps=20)

print()
print("Traces are reused for Tensor arguments.")
train(num_steps=tf.constant(10))
train(num_steps=tf.constant(20))

Retracing occurs for different Python arguments.
Tracing with num_steps =  10
Executing with num_steps =  10
Tracing with num_steps =  20
Executing with num_steps =  20

Traces are reused for Tensor arguments.
Tracing with num_steps =  Tensor("num_steps:0", shape=(), dtype=int32)
Executing with num_steps =  10
Executing with num_steps =  20