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

XLA: ottimizzazione 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 potenzialmente senza modifiche al codice sorgente.

I risultati sono miglioramenti nella velocità e nell'utilizzo della memoria: la maggior parte dei benchmark interni viene eseguita ~ 1,15 volte più velocemente dopo l'abilitazione 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 individualmente dall'esecutore TensorFlow. Ogni operazione TensorFlow ha un'implementazione del kernel GPU precompilata a cui invia l'esecutore.

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

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

Eseguito senza XLA, il grafico lancia tre kernel: uno per la moltiplicazione, uno per l'addizione 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 unico kernel GPU. Inoltre, questa operazione di fusione non scrive in memoria i valori intermedi prodotti da y*z x+y*z ; invece "trasmette" i risultati di questi calcoli intermedi direttamente ai loro utenti mantenendoli interamente nei registri della GPU. Fusion è l'ottimizzazione più importante di XLA. La larghezza di banda della memoria è in genere la risorsa più scarsa sugli acceleratori hardware, quindi rimuovere le operazioni di memoria è uno dei modi migliori per migliorare le prestazioni.

Abilita XLA per i modelli TensorFlow

Auto-clustering

Un modo più semplice per iniziare a utilizzare XLA nei modelli TensorFlow è abilitare il clustering automatico, che trova automaticamente i cluster (sottografi connessi) all'interno del grafico TensorFlow che possono essere compilati ed eseguiti utilizzando XLA. Il clustering automatico sulla GPU può essere abilitato impostando la variabile d'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 GPU, ma può anche essere abilitato sulla CPU utilizzando anche 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, vedere il tutorial sul clustering automatico colab .

Compilazione esplicita con tf.function

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

L'API di compilazione esplicita offre un controllo più dettagliato per la scelta delle funzioni da compilare. Ad esempio, la seguente funzione TensorFlow che esegue l'addestramento 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 oppure viene generata un'eccezione errors.InvalidArgumentError . XLA non può attualmente compilare funzioni in cui le dimensioni non sono inferibili : cioè, 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 piste 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 fornisce servizi di introspezione che consentono di ispezionare i programmi generati. Per eseguire il dump dei programmi generati, utilizzare la variabile d'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 Generato programmi XLA , uno per ogni cluster compilato. Allegarli quando si inviano segnalazioni di bug XLA è estremamente utile!

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

  • module_XXXX.ptx File PTX generati.

Puoi anche eseguire il dump del grafico visualizzando l'incorporamento di 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ù facile da riprodurre se include dump per i programmi XLA generati e l'incorporamento di cluster automatico utilizzato. Per generarli per un programma TensorFlow in esecuzione con il clustering automatico, avvia: 

$ 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 segnalano bug, allegare il contenuto della /tmp/generated (a cui si fa riferimento sopra).

Se possibile, prova a isolare un bug in un singolo programma XLA utilizzando replay_computation ed eseguendolo in modo iterativo sui programmi generati.

Frontend XLA

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

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

Ulteriore lettura