Questa pagina è stata tradotta dall'API Cloud Translation.
Switch to English

XLA: ottimizzatore del compilatore per l'apprendimento automatico

XLA (Accelerated Linear Algebra) è un compilatore specifico del dominio per l'algebra lineare che può accelerare i modelli TensorFlow senza potenzialmente modifiche al codice sorgente.

I risultati sono miglioramenti della velocità e dell'utilizzo della memoria: la maggior parte dei benchmark interni esegue ~ 1,15 volte più veloce dopo l'attivazione di XLA. Il set di dati di seguito viene valutato su una singola GPU NVidia V100:

introduzione

Quando viene eseguito un programma TensorFlow, tutte le operazioni vengono eseguite singolarmente dall'esecutore TensorFlow. Ogni operazione TensorFlow ha un'implementazione del kernel GPU precompilata a cui invia l'esecutore.

XLA offre una modalità alternativa di modelli in esecuzione: compila il grafico TensorFlow in una sequenza di kernel di calcolo generati appositamente per il modello dato. Poiché questi kernel sono univoci per il modello, possono sfruttare le informazioni specifiche del modello per l'ottimizzazione. Ad esempio, diamo un'occhiata a un'ottimizzazione che XLA fa nel contesto di un semplice calcolo TensorFlow: 

 def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

Esegui senza XLA, il grafico lancia tre kernel: uno per la moltiplicazione, uno per l'aggiunta e uno per la riduzione. Tuttavia, XLA può ottimizzare il grafico in modo che calcoli il risultato in un singolo lancio del kernel. Lo fa "fondendo" l'aggiunta, la moltiplicazione e la riduzione in un singolo kernel GPU. Inoltre, questa operazione fusa non scrive i valori intermedi prodotti da y*z e x+y*z nella memoria; invece "trasmette" i risultati di questi calcoli intermedi direttamente ai loro utenti, mantenendoli interamente nei registri GPU. Fusion è l'ottimizzazione singola più importante di XLA. La larghezza di banda della memoria è in genere la risorsa più scarsa per gli acceleratori hardware, quindi rimuovere le operazioni di memoria è uno dei modi migliori per migliorare le prestazioni.

Abilita XLA per i modelli TensorFlow

Auto-raggruppamento

Un modo più semplice per iniziare a utilizzare XLA nei modelli TensorFlow è abilitare il clustering automatico, che trova automaticamente i cluster (sottografi collegati) all'interno del grafico TensorFlow che può essere compilato ed eseguito utilizzando XLA. Il clustering automatico su GPU può essere abilitato impostando la variabile di ambiente TF_XLA_FLAGS : 

 $ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Il clustering automatico è attualmente ottimizzato per i carichi di lavoro della GPU, ma può anche essere abilitato sulla CPU utilizzando inoltre il flag --tf_xla_cpu_global_jit : 

 $ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

Per un esempio di utilizzo dettagliato, consultare il tutorial colab del clustering automatico .

Compilazione esplicita con tf.function

Il clustering automatico è un ottimo strumento per rendere il modello più veloce senza alcuna modifica al codice, ma può essere difficile capire quali modifiche sono state eseguite.

L'API di compilazione esplicita offre un controllo più accurato per la scelta delle funzioni da compilare. Ad esempio, la seguente funzione TensorFlow che esegue l'allenamento MNIST è compilata con XLA: 

 @tf.function(experimental_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))

L'API experimental_compile ha una semantica da compilare : l'intera funzione viene compilata con XLA o viene generata un'eccezione errors.InvalidArgumentError . Al momento XLA non può compilare funzioni in cui le dimensioni non sono inferibili : in altre parole , se non è possibile inferire le dimensioni di tutti i tensori senza eseguire l'intero calcolo. Ad esempio, la seguente funzione non verrà compilata: 

 @tf.function
def not_compilable(x):
  return tf.unique(x)

Le forme possono variare tra le varie corse però: 

 @tf.function(experimental_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))

Vedi il tutorial colab per un esempio di utilizzo più dettagliato.

Compilazione AOT (Ahead-of-time) per CPU con tfcompile

È inoltre possibile utilizzare uno strumento tfcompile autonomo, che converte il grafico TensorFlow in codice eseguibile (solo per CPU x86-64).

Ispeziona i programmi compilati

XLA offre servizi di introspezione che consentono di ispezionare i programmi generati. Per scaricare i programmi generati, utilizzare la variabile di ambiente XLA_FLAGS : 

 $ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

Dopo aver eseguito il dumping, è possibile trovare i seguenti file in /tmp/generated :

  • module_XXXX.*_optimizations.txt Programmi XLA generati, uno per ogni cluster compilato. Allegare quelli quando si inviano segnalazioni di bug XLA è estremamente utile!

  • module_XXXX.ir-*.ll File generati nella rappresentazione intermedia LLVM , con elementi intrinseci NVPTX .

  • module_XXXX.ptx File PTX generati.

Puoi anche scaricare il grafico visualizzando l'incorporamento dei cluster XLA all'interno del grafico TensorFlow con: 

 $ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

Segnalazioni di bug riproducibili

Una segnalazione di bug è molto più semplice da riprodurre se include dump per i programmi XLA generati e l'incorporamento automatico del clustering utilizzato. Per generarli per un programma TensorFlow in esecuzione con il clustering automatico, avviare: 

 $ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"

Quando si archiviano i bug, allegare il contenuto della /tmp/generated (di cui sopra).

Se possibile, prova a isolare un bug in un singolo programma XLA usando replay_computation ed eseguendolo iterativamente su programmi generati.

Frontend XLA

Oltre a TensorFlow, i programmi XLA possono essere generati da:

  • JAX : trasformazioni componibili di programmi Python + NumPy
  • Julia : la lingua Julia per il calcolo scientifico
  • PyTorch : framework PyTorch

Ulteriori letture