Introduzione al convertitore di inferenza Cloud TPU v5e

Introduzione

Il convertitore di inferenza Cloud TPU prepara e ottimizza un modello TensorFlow 2 (TF2) per l'inferenza TPU. Il convertitore viene eseguito in una shell VM locale o TPU. La shell della VM TPU è consigliata perché è preinstallata con il comando gli strumenti a riga di comando necessari per il convertitore. È necessario esportare SavedModel ed esegue le seguenti operazioni:

  1. Conversione TPU: aggiunge TPUPartitionedCall e altre operazioni TPU al per renderlo pubblicabile sulla TPU. Per impostazione predefinita, un modello esportato di inferenza non dispone di tali operazioni e non può essere fornita sulla TPU, anche se è stato addestrato sulla TPU.
  2. Batching: aggiunge operazioni in batch al modello per abilitare il raggruppamento in batch nel grafico. per una velocità effettiva migliore.
  3. Conversione BFloat16: consente di convertire il formato dei dati del modello da float32 a bfloat16 per migliorare le prestazioni di calcolo e ridurre l'utilizzo della memoria ad alta larghezza di banda (HBM) sulla TPU.
  4. Ottimizzazione forma IO: ottimizza le forme dei tensori per i dati trasferiti tra CPU e TPU per migliorare l'utilizzo della larghezza di banda.

Quando esportano un modello, gli utenti creano alias di funzione per tutte le funzioni che vogliono eseguire sulla TPU. Passano queste funzioni al convertitore e alla Il convertitore le posiziona sulla TPU e le ottimizza.

Il convertitore di inferenza di Cloud TPU è disponibile come immagine Docker, che può essere eseguite in qualsiasi ambiente in cui Docker è installato.

Tempo stimato per completare i passaggi mostrati sopra: circa 20-30 minuti

Prerequisiti

  1. Il modello deve essere un modello TF2 ed esportato nel SavedModel formato.
  2. Il modello deve avere un alias funzione per la funzione TPU. Per sapere come fare, consulta l'esempio di codice. Gli esempi seguenti utilizzano tpu_func come TPU alias di funzione.
  3. Assicurati che la CPU del computer supporti Advanced Vector eXtensions (AVX) di Cloud TPU, la libreria TensorFlow (la dipendenza Inference Converter) viene compilato per utilizzare le istruzioni AVX. La maggior parte delle CPU il supporto AVX.
    1. Puoi eseguire lscpu | grep avx per verificare se AVX di istruzioni è supportato.

Prima di iniziare

Prima di iniziare la configurazione, segui questi passaggi:

  • Creare un nuovo progetto: Nella pagina del selettore progetti della console Google Cloud, seleziona o crea un progetto cloud.

  • Configura una VM TPU: Crea una nuova VM TPU utilizzando la console Google Cloud o gcloud oppure utilizza un per eseguire l'inferenza con la VM TPU sulla VM TPU.

    • Assicurati che l'immagine VM TPU sia basata su TensorFlow. Ad esempio, --version=tpu-vm-tf-2.11.0.
    • Il modello convertito verrà caricato e pubblicato su questa VM TPU.

  • Assicurati di avere gli strumenti a riga di comando necessari per utilizzare Cloud TPU Convertitore di inferenza. Puoi installare Google Cloud SDK e Docker in locale o utilizzare una VM TPU su cui questo software è installato per impostazione predefinita. Questi strumenti ti consentono di interagire con l'immagine del convertitore.

    Connettiti all'istanza con SSH utilizzando il seguente comando: 

    gcloud compute tpus tpu-vm ssh ${tpu-name} --zone ${zone} --project ${project-id}

Configurazione dell'ambiente

Configura il tuo ambiente dalla shell della VM TPU o dalla shell locale.

Shell VM TPU

  • Nella shell della VM TPU, esegui i seguenti comandi per consentire l'utilizzo di Docker senza privilegi di root:

    sudo usermod -a -G docker ${USER}
newgrp docker

  • Inizializza gli aiuti per le credenziali Docker:

    gcloud auth configure-docker \
  us-docker.pkg.dev

Shell locale

Nella shell locale, configura l'ambiente seguendo questi passaggi:

  • Installa Cloud SDK, che include lo strumento a riga di comando gcloud.

  • Installa Docker:

  • Consenti l'utilizzo di Docker non root:

    sudo usermod -a -G docker ${USER}
newgrp docker

  • Accedi al tuo ambiente:

    gcloud auth login

  • Inizializza gli helper per le credenziali Docker:

    gcloud auth configure-docker \
    us-docker.pkg.dev

  • Estrai l'immagine Docker di Inference Converter:

      CONVERTER_IMAGE=us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
  docker pull ${CONVERTER_IMAGE}

Immagine del convertitore

L'immagine serve a effettuare conversioni del modello una tantum. Impostare i percorsi del modello e regolare opzioni del convertitore per soddisfare le tue esigenze. La sezione Esempi di utilizzo fornisce diversi casi d'uso comuni.

docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      function_alias: "tpu_func"
    }
    batch_options {
      num_batch_threads: 2
      max_batch_size: 8
      batch_timeout_micros: 5000
      allowed_batch_sizes: 2
      allowed_batch_sizes: 4
      allowed_batch_sizes: 8
      max_enqueued_batches: 10
    }
'

Inferenza con il modello convertito nella VM TPU

# Initialize the TPU
resolver = tf.distribute.cluster_resolver.TPUClusterResolver("local")
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

# Load the model
model = tf.saved_model.load(${CONVERTED_MODEL_PATH})

# Find the signature function for serving
serving_signature = 'serving_default' # Change the serving signature if needed
serving_fn = model.signatures[serving_signature]
# Run the inference using requests.
results = serving_fn(**inputs)
logging.info("Serving results: %s", str(results))

Esempi di utilizzo

Aggiungi un alias di funzione per la funzione TPU

  1. Trova o crea una funzione nel modello che aggrega tutto ciò che vuoi per l'esecuzione sulla TPU. Se @tf.function non esiste, aggiungilo.
  2. Quando salvi il modello, fornisci SaveOptions come segue per ottenere model.tpu_func un alias func_on_tpu.
  3. Puoi passare questo alias di funzione al convertitore.
class ToyModel(tf.keras.Model):
  @tf.function(
      input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
  def tpu_func(self, x):
    return x * 1.0

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'func_on_tpu': model.tpu_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

Converti un modello con più funzioni TPU

Puoi aggiungere più funzioni alla TPU. Basta creare più funzioni alias e passarli in converter_options_string all'utente che ha completato una conversione.

tpu_functions {
  function_alias: "tpu_func_1"
}
tpu_functions {
  function_alias: "tpu_func_2"
}

Quantizzazione

La quantizzazione è una tecnica che riduce la precisione dei numeri utilizzato per rappresentare i parametri di un modello. Ciò si traduce in dimensioni del modello inferiori e calcoli più rapidi. Un modello quantizzato offre guadagni in termini di throughput dell'inferenza, nonché un utilizzo della memoria e dimensioni di archiviazione inferiori, a fronte di piccoli cali di precisione.

La nuova funzionalità di quantizzazione post-addestramento in TensorFlow che ha come target la TPU è stata sviluppata a partire dalla funzionalità esistente simile in TensorFlow Lite che viene utilizzata per scegliere come target i dispositivi mobili e edge. Per scoprire di più sulla quantizzazione in generale, puoi consultare la documentazione di TensorFlow Lite.

Concetti di quantizzazione

Questa sezione definisce concetti specificamente correlati alla quantizzazione con Inference Converter.

I concetti relativi ad altre configurazioni TPU (ad esempio, slice, host, chip e TensorCore) sono descritti nella pagina Architettura del sistema TPU.

  • Quantizzazione post-addestramento (PTQ): la PTQ è una tecnica che riduce le dimensioni e la complessità di calcolo di un modello di rete neurale influenzarne significativamente la precisione. La PTQ funziona convertendo i coefficienti e le attivazioni in virgola mobile di un modello addestrato in numeri interi di precisione inferiore, ad esempio interi a 8 o 16 bit. Ciò può causare una significativa riduzione delle dimensioni del modello e della latenza di inferenza, mentre subirà una piccola perdita di accuratezza.

  • Calibrazione: il passaggio di calibrazione per la quantizzazione è il processo di raccogliere statistiche sull'intervallo di valori ponderati e delle attivazioni di un modello di rete neurale. Queste informazioni vengono utilizzate per determinare i parametri di quantizzazione del modello, ovvero i valori che verranno utilizzati per convertire le attivazioni e i pesi a virgola mobile in numeri interi.

  • Set di dati rappresentativo: un set di dati rappresentativo per la quantizzazione. è un piccolo set di dati che rappresenta i dati di input effettivi per il modello. Viene utilizzato durante la fase di calibrazione della quantizzazione per raccogliere statistiche sull'intervallo di valori di cui le ponderazioni e le attivazioni del modello. Il set di dati rappresentativo deve soddisfare seguenti proprietà:

    • Deve rappresentare correttamente gli input effettivi per il modello durante l'inferenza. Ciò significa che deve coprire l'intervallo di valori è probabile che il modello veda nel mondo reale.
    • Dovrebbe scorrere collettivamente ogni ramo dei condizionali (ad es. tf.cond), se presenti. Questo è importante perché di quantizzazione deve essere in grado di gestire tutti i possibili input al modello, anche se non sono rappresentati esplicitamente nel e rappresentativo.
    • Deve essere sufficientemente grande da raccogliere statistiche sufficienti e ridurre gli errori. Come regola generale, ti consigliamo di utilizzare più di 200 campioni rappresentativi.

    Il set di dati rappresentativo può essere un sottoinsieme del set di dati di addestramento o un set di dati separato progettato specificamente per essere rappresentativo degli input reali del modello. La scelta di il set di dati da usare dipende dall'applicazione specifica.

  • Quantizzazione dell'intervallo statico (SRQ): la quantizzazione dell'intervallo statico determina una volta sola l'intervallo di valori per le ponderazioni e le attivazioni di un modello di rete neurale, durante il passaggio di calibrazione. Ciò significa che viene utilizzato lo stesso intervallo di valori per tutti gli input del modello. Questo valore può essere meno preciso rispetto all'intervallo dinamico la quantizzazione, soprattutto per i modelli con un'ampia gamma di valori di input. Tuttavia, la quantizzazione dell'intervallo statico richiede meno calcoli in fase di esecuzione rispetto alla quantizzazione dell'intervallo dinamico.

  • Dynamic Range Quantization (DRQ): DRQ determina l'intervallo di per le ponderazioni e le attivazioni di un modello di rete neurale ogni input. In questo modo, il modello si adatta all'intervallo di valori dei dati di input, il che può migliorare l'accuratezza. Tuttavia, la quantizzazione dinamica richiede più calcoli in fase di esecuzione rispetto alla quantizzazione statica.

    Funzionalità Quantizzazione dell'intervallo statico Quantizzazione dell'intervallo dinamico
    Intervallo di valori Determinato una volta, durante la calibrazione Determinato per ogni input
    Accuratezza Può essere meno preciso, soprattutto per i modelli con un'ampia gamma di valori di input Può essere più preciso, soprattutto per i modelli con un'ampia gamma di valori di input
    complessità Più semplice Più complesso
    Calcolo in fase di esecuzione Meno calcoli Più calcolo

  • Quantizzazione solo dei pesi: la quantizzazione solo dei pesi è un tipo di quantizzazione che esegue la quantizzazione solo dei pesi di un modello di rete neurale, lasciando le attivazioni in virgola mobile. Questa può essere una buona opzione per i modelli sensibili all'accuratezza, in quanto può contribuire a preservarne l'accuratezza.

Come utilizzare la quantizzazione

La quantizzazione può essere applicata configurando e impostando QuantizationOptions sulle opzioni del convertitore. Le opzioni più rilevanti sono:

  • tags: raccolta di tag che identificano i MetaGraphDef all'interno del SavedModel da quantizzare. Non è necessario specificare se hai un solo MetaGraphDef.
  • firma_keys: sequenza di chiavi che identificano SignatureDef contenenti input e output. Se non specificato, viene utilizzato ["serving_default"].
  • quantization_method: metodo di quantizzazione da applicare. Se non specificato, verrà applicata la quantizzazione STATIC_RANGE.
  • op_set: deve essere mantenuto come XLA. Attualmente è l'opzione predefinita, non è necessario specificarla.
  • representative_datasets: specifica il set di dati utilizzato per la calibrazione dei parametri di quantizzazione.

Creazione del set di dati rappresentativo

Un insieme di dati rappresentativo è essenzialmente un insieme di campioni. Dove un campione è una mappa di: {input_key: input_value}. Ad esempio:

representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]

I set di dati rappresentativi devono essere salvati come TFRecord utilizzando la classe TfRecordRepresentativeDatasetSaver attualmente disponibile nel pacchetto tf-nightly pip. Ad esempio:

# Assumed tf-nightly installed.
import tensorflow as tf
representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]
tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
       path_map={'serving_default': '/tmp/representative_dataset_path'}
    ).save({'serving_default': representative_dataset})

Esempi

L'esempio seguente quantifica il modello con la chiave di firma di serving_default e alias di funzione di tpu_func:

docker run \
  --mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
  --mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
  ${CONVERTER_IMAGE} \
  --input_model_dir=/tmp/input \
  --output_model_dir=/tmp/output \
  --converter_options_string=' \
    tpu_functions { \
      function_alias: "tpu_func" \
    } \
    external_feature_configs { \
      quantization_options { \
        signature_keys: "serving_default" \
        representative_datasets: { \
          key: "serving_default" \
          value: { \
            tfrecord_file_path: "${TF_RECORD_FILE}" \
          } \
        } \
      } \
    } '

Aggiungi raggruppamento

Il convertitore può essere utilizzato per aggiungere la gestione in batch a un modello. Per una descrizione di opzioni di raggruppamento in batch, vedi Definizione delle opzioni di raggruppamento in batch.

Per impostazione predefinita, il convertitore eseguirà il batch di tutte le funzioni TPU nel modello. Inoltre, può elaborare in batch le firme e le funzioni fornite dall'utente, il che può migliorare ulteriormente il rendimento. Qualsiasi funzione TPU, funzione fornita dall'utente o una firma in batch, devono soddisfare i requisiti rigidi requisiti di forma.

L'utente che ha completato una conversione può inoltre aggiorna le opzioni di batch esistenti. Di seguito è riportato un esempio di come aggiungere il raggruppamento a un modello. Per maggiori informazioni informazioni sui processi in batch, vedi Approfondimento sul batch.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

Disattivare le ottimizzazioni di bfloat16 e della forma IO

Le ottimizzazioni della forma di BFloat16 e IO sono abilitate per impostazione predefinita. Se non funzionano bene con il tuo modello, possono essere disattivati.

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

Report sulle conversioni

Puoi trovare questo report sulle conversioni nel log dopo aver eseguito l'Inferenza Convertitore. Di seguito è riportato un esempio.

-------- Conversion Report --------
TPU cost of the model: 96.67% (2034/2104)
CPU cost of the model:  3.33% (70/2104)

Cost breakdown
================================
%         Cost    Name
--------------------------------
3.33      70      [CPU cost]
48.34     1017    tpu_func_1
48.34     1017    tpu_func_2
--------------------------------

Questo report stima il costo computazionale del modello di output su CPU e TPU e suddivide ulteriormente il costo della TPU per ogni funzione, che dovrebbe riflettere la tua selezione delle funzioni TPU nelle opzioni del convertitore.

Se vuoi utilizzare meglio la TPU, ti consigliamo di fare esperimenti con il modello strutturare e regolare le opzioni del convertitore.

Domande frequenti

Quali funzioni devo inserire nella TPU?

È meglio inserire il maggior numero possibile di operazioni sul modello TPU, perché la maggior parte delle operazioni viene eseguita più velocemente su TPU.

Se il modello non contiene operazioni, stringhe o tensori sparsi incompatibili con la TPU, in genere la strategia migliore è eseguire l'intero modello sulla TPU. e puoi trovare o creare una funzione che aggrega l'intero modello, creando un alias di funzione e passandolo al convertitore.

Se il modello contiene parti che non possono funzionare sulla TPU (ad es. operazioni, stringhe o tensori sparsi non compatibili con la TPU), la scelta delle funzioni TPU dipende dalla posizione della parte incompatibile.

  • Se si trova all'inizio o alla fine del modello, puoi eseguire il refactoring del modello per mantenerlo sulla CPU. Alcuni esempi sono le fasi di pre- e post-elaborazione delle stringhe. Per ulteriori informazioni sullo spostamento del codice nella CPU, consulta: "Come posso spostare una parte del modello nella CPU?" Mostra un modo tipico per eseguire il refactoring del modello.
  • Se si trova a metà del modello, è meglio dividere il modello in tre parti e contenere tutte le operazioni incompatibili con TPU nella parte centrale e farla funzionare sulla CPU.
  • Se è un tensore sparso, valuta la possibilità di chiamare tf.sparse.to_dense sulla CPU e passaggio del tensore denso risultante alla porzione TPU del modello.

Un altro fattore da considerare è l'utilizzo di HBM. L'incorporamento delle tabelle può richiedere HBM. Se superano i limiti hardware della TPU, devono essere messi sulla CPU, insieme alle operazioni di ricerca.

Se possibile, deve esistere una sola funzione TPU sotto una firma. Se la struttura del tuo modello richiede di chiamare più funzioni TPU per richiesta di inferenza in entrata, è necessario essere consapevoli della latenza aggiuntiva dell'invio tensori tra CPU e TPU.

Un buon modo per valutare la selezione delle funzioni TPU è controllare il report sulle conversioni. Mostra la percentuale di calcolo posizionato sulla TPU e del costo di ogni funzione TPU.

Come faccio a spostare una parte del modello sulla CPU?

Se il modello contiene parti che non possono essere eseguite sulla TPU, devi eseguire il refactoring del modello per spostarle sulla CPU. Ecco un esempio di giocattolo. Il modello è un modello linguistico con una fase di pre-elaborazione. Il codice per le definizioni dei livelli e vengono omesse per semplicità.

class LanguageModel(tf.keras.Model):
  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.bert_layer(word_ids)

Questo modello non può essere pubblicato direttamente sulla TPU per due motivi. Innanzitutto, il parametro è una stringa. In secondo luogo, la funzione preprocess può contenere molte stringhe operazioni. Entrambi non sono compatibili con TPU.

Per eseguire il refactoring di questo modello, puoi creare un'altra funzione denominata tpu_func in ospitare bert_layer che richiede un'alta intensità di calcolo. Quindi crea un alias di funzione per tpu_func e passarlo all'utente che ha completato una conversione. In questo modo, tutto ciò che si trova all'interno tpu_func verrà eseguito sulla TPU e tutto ciò che rimane in model_func verrà eseguito il per la CPU.

class LanguageModel(tf.keras.Model):
  @tf.function
  def tpu_func(self, word_ids):
    return self.bert_layer(word_ids)

  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.tpu_func(word_ids)

Cosa devo fare se il modello ha operazioni, stringhe o tensori sparsi non compatibili con TPU?

La TPU è supportata per la maggior parte delle operazioni standard di TensorFlow, ma alcune inclusi tensori e stringhe sparsi non sono supportati. Il convertitore non controlla la presenza di operazioni incompatibili con TPU. Pertanto, un modello contenente queste operazioni può passare la conversione. Tuttavia, quando la esegui per l'inferenza, si verificheranno errori come quelli riportati di seguito.

'tf.StringToNumber' op isn't compilable for TPU device.

Se il modello contiene operazioni incompatibili con TPU, queste devono essere posizionate al di fuori della funzione TPU. Inoltre, la stringa è un formato di dati non supportato sulla TPU. Quindi... le variabili di tipo stringa non devono essere inserite nella funzione TPU. E i e i valori restituiti della funzione TPU non devono essere di tipo stringa come beh. Analogamente, evita di inserire tensori sparsi nella funzione TPU, inclusi i suoi parametri e valori restituiti.

Di solito non è difficile eseguire il refactoring della parte incompatibile del modello nella CPU. Ecco un esempio.

Come supportare le operazioni personalizzate nel modello?

Se nel modello vengono utilizzate operazioni personalizzate, il convertitore potrebbe non riconoscerle e non riuscire a convertire il modello. Questo accade perché la libreria dell'operatore personalizzato, che contiene la definizione completa dell'operatore, non è collegata al convertitore.

Poiché al momento il codice del convertitore non è ancora open source, non è possibile creato con un'operazione personalizzata.

Cosa devo fare se dispongo di un modello TensorFlow 1?

Il convertitore non supporta i modelli TensorFlow 1. È necessario eseguire la migrazione dei modelli TensorFlow 1 a TensorFlow 2.

Devo abilitare il bridge MLIR durante l'esecuzione del mio modello?

La maggior parte dei modelli convertiti può essere eseguita con il più recente bridge MLIR TF2XLA il bridge TF2XLA originale.

Come faccio a convertire un modello già esportato senza un alias di funzione?

Se un modello è stato esportato senza un alias di funzione, il modo più semplice è esportarlo di nuovo e creare un alias di funzione. Se la riesportazione non è disponibile, è comunque possibile convertire il modello fornendo un concrete_function_name. Tuttavia, identificare il corretto concrete_function_name richiede un po' di lavoro di detective.

Gli alias di funzione sono una mappatura di una stringa definita dall'utente a una funzione concreta nome. Facilitano il riferimento a una funzione specifica nel modello. Il convertitore accetta sia gli alias di funzione sia i nomi di funzioni concrete non elaborati.

I nomi delle funzioni specifiche sono disponibili esaminando saved_model.pb.

L'esempio seguente mostra come inserire una funzione concreta chiamata __inference_serve_24 sulla TPU.

sudo docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      concrete_function_name: "__inference_serve_24"
    }'

Come faccio a risolvere un errore di vincolo costante in fase di compilazione?

Sia per l'addestramento che per l'inferenza, XLA richiede che gli input per determinate operazioni abbiano un nota al momento di compilazione della TPU. Ciò significa che quando XLA compila la parte del programma relativa alla TPU, gli input di queste operazioni devono avere una forma nota in modo statico.

Esistono due modi per risolvere il problema.

  • La soluzione migliore è aggiornare gli input dell'operazione in modo che abbiano una per il momento in cui XLA compila il programma TPU. Questa compilazione avviene appena prima dell'esecuzione della parte TPU del modello. Ciò significa che la forma deve essere nota in modo statico al momento dell'esecuzione di TpuFunction.
  • Un'altra opzione è modificare TpuFunction in modo che non includa più il parametro operazione problematica

Perché ricevo un errore relativo alla forma del batching?

Il raggruppamento ha requisiti rigorosi per la forma che consentono di raggruppare le richieste in arrivo in base alla dimensione 0 (ovvero la dimensione di raggruppamento). Questi requisiti relativi alla forma provengono dall'operazione di accumulo di TensorFlow e non possono essere allentati.

Il mancato rispetto di questi requisiti comporterà errori quali:

  1. I tensori di input per il raggruppamento devono avere almeno una dimensione.
  2. Le dimensioni degli input devono corrispondere.
  3. I tensori di input per il raggruppamento forniti in una determinata chiamata all'operatore devono avere dimensioni uguali per la dimensione 0.
  4. La dimensione 0 del tensore di output raggruppato non corrisponde alla somma delle dimensioni 0 del tensore di input.

Per soddisfare questi requisiti, valuta la possibilità di fornire un'altra funzione o firma per creare un batch. Potrebbe anche essere necessario modificare le funzioni esistenti per soddisfare questi requisiti.

Se una funzione viene distribuito in batch, assicurati che le forme di input_signature di @tf.function siano tutte Nessuna è presente nella dimensione 0°. Se una firma viene raggruppata, assicurati che tutti i relativi input abbiano -1 nella dimensione 0.

Per una spiegazione completa del motivo per cui si verificano questi errori e di come risolverli vedere: Approfondimento sul batch.

Problemi noti

La funzione TPU non può chiamare indirettamente un'altra funzione TPU

Sebbene il convertitore possa gestire la maggior parte degli scenari di chiamata di funzioni oltre il confine tra CPU e TPU, esiste un caso limite raro in cui non riesce. È quando una TPU chiama indirettamente un'altra funzione TPU.

Questo perché il convertitore modifica il chiamante diretto di una funzione TPU in modo che chiami la funzione TPU stessa anziché uno stub di chiamata TPU. Lo stub della chiamata contiene operazioni che possono funzionare solo sulla CPU. Quando una funzione TPU chiama una qualsiasi che alla fine chiama il chiamante diretto, queste operazioni della CPU potrebbero sulla TPU da eseguire, generando errori di kernel mancanti. Tieni presente che questo caso è diverso da una funzione TPU che chiama direttamente un'altra funzione TPU. In questo caso, il convertitore non modifica nessuna funzione per chiamare lo stub di chiamata, quindi può funzionare.

Nel convertitore abbiamo implementato il rilevamento di questo scenario. Se vedi il seguente errore, significa che il modello ha riscontrato questo caso limite:

Unable to place both "__inference_tpu_func_2_46" and "__inference_tpu_func_4_68"
on the TPU because "__inference_tpu_func_2_46" indirectly calls
"__inference_tpu_func_4_68". This behavior is unsupported because it can cause
invalid graphs to be generated.

La soluzione generale è eseguire il refactoring del modello per evitare questo scenario di chiamata di funzione. Se non riesci a farlo, contatta il team di assistenza Google per ulteriori informazioni.

Riferimento

Opzioni di conversione in formato Protobuf

message ConverterOptions {
  // TPU conversion options.
  repeated TpuFunction tpu_functions = 1;

  // The state of an optimization.
  enum State {
    // When state is set to default, the optimization will perform its
    // default behavior. For some optimizations this is disabled and for others
    // it is enabled. To check a specific optimization, read the optimization's
    // description.
    DEFAULT = 0;
    // Enabled.
    ENABLED = 1;
    // Disabled.
    DISABLED = 2;
  }

  // Batch options to apply to the TPU Subgraph.
  //
  // At the moment, only one batch option is supported. This field will be
  // expanded to support batching on a per function and/or per signature basis.
  //
  //
  // If not specified, no batching will be done.
  repeated BatchOptions batch_options = 100;

  // Global flag to disable all optimizations that are enabled by default.
  // When enabled, all optimizations that run by default are disabled. If a
  // default optimization is explicitly enabled, this flag will have no affect
  // on that optimization.
  //
  // This flag defaults to false.
  bool disable_default_optimizations = 202;

  // If enabled, apply an optimization that reshapes the tensors going into
  // and out of the TPU. This reshape operation improves performance by reducing
  // the transfer time to and from the TPU.
  //
  // This optimization is incompatible with input_shape_opt which is disabled.
  // by default. If input_shape_opt is enabled, this option should be
  // disabled.
  //
  // This optimization defaults to enabled.
  State io_shape_optimization = 200;

  // If enabled, apply an optimization that updates float variables and float
  // ops on the TPU to bfloat16. This optimization improves performance and
  // throughtput by reducing HBM usage and taking advantage of TPU support for
  // bfloat16.
  //
  // This optimization may cause a loss of accuracy for some models. If an
  // unacceptable loss of accuracy is detected, disable this optimization.
  //
  // This optimization defaults to enabled.
  State bfloat16_optimization = 201;

  BFloat16OptimizationOptions bfloat16_optimization_options = 203;

  // The settings for XLA sharding. If set, XLA sharding is enabled.
  XlaShardingOptions xla_sharding_options = 204;
}

message TpuFunction {
  // The function(s) that should be placed on the TPU. Only provide a given
  // function once. Duplicates will result in errors. For example, if
  // you provide a specific function using function_alias don't also provide the
  // same function via concrete_function_name or jit_compile_functions.
  oneof name {
    // The name of the function alias associated with the function that
    // should be placed on the TPU. Function aliases are created during model
    // export using the tf.saved_model.SaveOptions.
    //
    // This is a recommended way to specify which function should be placed
    // on the TPU.
    string function_alias = 1;

    // The name of the concrete function that should be placed on the TPU. This
    // is the name of the function as it found in the GraphDef and the
    // FunctionDefLibrary.
    //
    // This is NOT the recommended way to specify which function should be
    // placed on the TPU because concrete function names change every time a
    // model is exported.
    string concrete_function_name = 3;

    // The name of the signature to be placed on the TPU. The user must make
    // sure there is no TPU-incompatible op under the entire signature.
    string signature_name = 5;

    // When jit_compile_functions is set to True, all jit compiled functions
    // are placed on the TPU.
    //
    // To use this option, decorate the relevant function(s) with
    // @tf.function(jit_compile=True), before exporting. Then set this flag to
    // True. The converter will find all functions that were tagged with
    // jit_compile=True and place them on the TPU.
    //
    // When using this option, all other settings for the TpuFunction
    // will apply to all functions tagged with
    // jit_compile=True.
    //
    // This option will place all jit_compile=True functions on the TPU.
    // If only some jit_compile=True functions should be placed on the TPU,
    // use function_alias or concrete_function_name.
    bool jit_compile_functions = 4;
  }

}

message BatchOptions {
  // Number of scheduling threads for processing batches of work. Determines
  // the number of batches processed in parallel. This should be roughly in line
  // with the number of TPU cores available.
  int32 num_batch_threads = 1;

  // The maximum allowed batch size.
  int32 max_batch_size = 2;

  // Maximum number of microseconds to wait before outputting an incomplete
  // batch.
  int32 batch_timeout_micros = 3;

  // Optional list of allowed batch sizes. If left empty,
  // does nothing. Otherwise, supplies a list of batch sizes, causing the op
  // to pad batches up to one of those sizes. The entries must increase
  // monotonically, and the final entry must equal max_batch_size.
  repeated int32 allowed_batch_sizes = 4;

  // Maximum number of batches enqueued for processing before requests are
  // failed fast.
  int32 max_enqueued_batches = 5;

  // If set, disables large batch splitting which is an efficiency improvement
  // on batching to reduce padding inefficiency.
  bool disable_large_batch_splitting = 6;

  // Experimental features of batching. Everything inside is subject to change.
  message Experimental {
    // The component to be batched.
    // 1. Unset if it's for all TPU subgraphs.
    // 2. Set function_alias or concrete_function_name if it's for a function.
    // 3. Set signature_name if it's for a signature.
    oneof batch_component {
      // The function alias associated with the function. Function alias is
      // created during model export using the tf.saved_model.SaveOptions, and is
      // the recommended way to specify functions.
      string function_alias = 1;

      // The concreate name of the function. This is the name of the function as
      // it found in the GraphDef and the FunctionDefLibrary. This is NOT the
      // recommended way to specify functions, because concrete function names
      // change every time a model is exported.
      string concrete_function_name = 2;

      // The name of the signature.
      string signature_name = 3;
    }
  }

  Experimental experimental = 7;
}

message BFloat16OptimizationOptions {
  // Indicates where the BFloat16 optimization should be applied.
  enum Scope {
    // The scope currently defaults to TPU.
    DEFAULT = 0;
    // Apply the bfloat16 optimization to TPU computation.
    TPU = 1;
    // Apply the bfloat16 optimization to the entire model including CPU
    // computations.
    ALL = 2;
  }

  // This field indicates where the bfloat16 optimization should be applied.
  //
  // The scope defaults to TPU.
  Scope scope = 1;

  // If set, the normal safety checks are skipped. For example, if the model
  // already contains bfloat16 ops, the bfloat16 optimization will error because
  // pre-existing bfloat16 ops can cause issues with the optimization. By
  // setting this flag, the bfloat16 optimization will skip the check.
  //
  // This is an advanced feature and not recommended for almost all models.
  //
  // This flag is off by default.
  bool skip_safety_checks = 2;

  // Ops that should not be converted to bfloat16.
  // Inputs into these ops will be cast to float32, and outputs from these ops
  // will be cast back to bfloat16.
  repeated string filterlist = 3;
}

message XlaShardingOptions {
  // num_cores_per_replica for TPUReplicateMetadata.
  //
  // This is the number of cores you wish to split your model into using XLA
  // SPMD.
  int32 num_cores_per_replica = 1;

  // (optional) device_assignment for TPUReplicateMetadata.
  //
  // This is in a flattened [x, y, z, core] format (for
  // example, core 1 of the chip
  // located in 2,3,0 will be stored as [2,3,0,1]).
  //
  // If this is not specified, then the device assignments will utilize the same
  // topology as specified in the topology attribute.
  repeated int32 device_assignment = 2;

  // A serialized string of tensorflow.tpu.TopologyProto objects, used for
  // the topology attribute in TPUReplicateMetadata.
  //
  // You must specify the mesh_shape and device_coordinates attributes in
  // the topology object.
  //
  // This option is required for num_cores_per_replica > 1 cases due to
  // ambiguity of num_cores_per_replica, for example,
  // pf_1x2x1 with megacore and df_1x1
  // both have num_cores_per_replica = 2, but topology is (1,2,1,1) for pf and
  // (1,1,1,2) for df.
  // - For pf_1x2x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,2,1,1]
  //   device_coordinates=flatten([0,0,0,0], [0,1,0,0])
  // - For df_1x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,1,1,2]
  //   device_coordinates=flatten([0,0,0,0], [0,0,0,1])
  // - For df_2x2, mesh shape and device_coordinates looks like:
  //   mesh_shape = [2,2,1,2]
  //   device_coordinates=flatten(
  //    [0,0,0,0],[0,0,0,1],[0,1,0,0],[0,1,0,1]
  //    [1,0,0,0],[1,0,0,1],[1,1,0,0],[1,1,0,1])
  bytes topology = 3;
}

Approfondimento sul batching

Il batch è utilizzato per migliorare la velocità effettiva e l'utilizzo della TPU. Consente di elaborare più richieste contemporaneamente. Durante l'addestramento, può essere eseguita utilizzando tf.data. Durante l'inferenza, in genere viene eseguita aggiungendo un'operazione nel grafico che raggruppa le richieste in arrivo. L'operazione attende che sia sufficiente richieste o un timeout prima che generi un batch di grandi dimensioni richieste individuali. Consulta la sezione Definizione delle opzioni di raggruppamento per ulteriori informazioni sulle diverse opzioni di raggruppamento che possono essere ottimizzate, inclusi i timeout e le dimensioni dei batch.

batch nei grafici

Per impostazione predefinita, il convertitore inserisce l'operazione di raggruppamento direttamente prima del calcolo TPU. Aggrega le funzioni TPU fornite dall'utente e qualsiasi TPU preesistente il calcolo nel modello con operazioni di batch. È possibile sostituire questo valore il comportamento predefinito indicando all'utente che ha completato una conversione funzioni e/o firme devono essere raggruppati in batch.

L'esempio seguente mostra come aggiungere la modalità batch predefinita.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

Batch di firme

La creazione in batch della firma raggruppa l'intero modello a partire dagli input della firma e andare agli output della firma. A differenza del comportamento predefinito del Converter, il raggruppamento delle firme raggruppa sia il calcolo della TPU sia il calcolo della CPU. Ciò permette di ottenere un miglioramento delle prestazioni dal 10% al 20% durante l'inferenza su alcune di grandi dimensioni.

Come tutti i tipi di batch, il raggruppamento in batch di Signature rigidi requisiti di forma. Per garantire che questi requisiti siano soddisfatti, l'inserimento della firma deve avere forme con almeno due dimensioni. La prima dimensione è la dimensione del batch e deve avere una dimensione pari a -1. Ad esempio, (-1, 4), (-1) o (-1, 128, 4, 10) sono tutte forme di input valide. Se non è possibile, valuta la possibilità di utilizzare il comportamento di raggruppamento predefinito o il raggruppamento delle funzioni.

Per utilizzare la creazione in batch di firme, fornisci i nomi delle firme come signature_name utilizzando BatchOptions.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    signature_name: "serving_default"
  }
}

Raggruppamento di funzioni

Il raggruppamento in batch di funzioni può essere utilizzato per indicare all'utente che ha completato una conversione quali funzioni devono essere in batch. Per impostazione predefinita, il convertitore eseguirà il batch di tutte le funzioni TPU. Il bagging delle funzioni sostituisce questo comportamento predefinito.

È possibile utilizzare il raggruppamento in batch di funzioni per il calcolo in batch della CPU. Molti modelli registrano un miglioramento del rendimento quando il calcolo della CPU viene eseguito in batch. Il modo migliore per il calcolo in batch della CPU prevede l'esecuzione in batch di firme, ma potrebbe non funzionare di grandi dimensioni. In questi casi, il batching delle funzioni può essere utilizzato per eseguire il batch di parte del calcolo della CPU oltre a quello della TPU. Tieni presente che l'operazione di raggruppamento non può vengono eseguite sulla TPU, quindi qualsiasi funzione di batch fornita deve essere chiamata per la CPU.

Il raggruppamento delle funzioni può essere utilizzato anche per soddisfare i rigorosi requisiti di forma imposti dall'operazione di raggruppamento. Nei casi in cui le funzioni TPU non soddisfino i requisiti di forma dell'operazione di raggruppamento, il raggruppamento delle funzioni può essere utilizzato per indicare al convertitore di raggruppare funzioni diverse.

Per utilizzarlo, genera un function_alias per la funzione che dovrebbe essere in batch. Puoi farlo trovando o creando una funzione nel tuo modello che racchiude tutto ciò che vuoi raggruppare. Assicurati che questa funzione soddisfi le requisiti di forma rigida imposti dall'operazione di raggruppamento in batch. Aggiungi @tf.function se non ne ha già uno. È importante fornire il input_signature al @tf.function. La dimensione 0 deve essere None perché è la dimensione del batch, pertanto non può avere dimensioni fisse. Ad esempio, [None, 4], [None] o [None, 128, 4, 10] sono tutti forme di input valide. Quando salvi il modello, fornisci SaveOptions come quelli mostrati di seguito per assegnare a model.batch_func un alias "batch_func". Quindi puoi passare questo alias di funzione al convertitore.

class ToyModel(tf.keras.Model):
  @tf.function(input_signature=[tf.TensorSpec(shape=[None, 10],
                                              dtype=tf.float32)])
  def batch_func(self, x):
    return x * 1.0

  ...

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'batch_func': model.batch_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

A questo punto, passa i function_alias utilizzando BatchOptions.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    function_alias: "batch_func"
  }
}

Definizione delle opzioni di raggruppamento in batch

  • num_batch_threads: (intero) Numero di thread di pianificazione per l'elaborazione di batch di lavoro. Determina il numero di batch elaborati in parallelo. Dovrebbe essere più o meno in linea con il numero di core TPU disponibili.
  • max_batch_size: dimensione massima consentita del batch (numero intero). Può essere più grande rispetto a allowed_batch_sizes per utilizzare la suddivisione del batch di grandi dimensioni.
  • batch_timeout_micros: numero massimo di microsecondi da attendere (numero intero) prima di generare un batch incompleto.
  • allowed_batch_sizes: (elenco di numeri interi) se l'elenco non è vuoto, i batch verranno aumentati fino alla dimensione più vicina nell'elenco. L'elenco deve essere aumentando monotonicamente e l'elemento finale deve essere minore o uguale a max_batch_size.
  • max_enqueued_batches: (intero) Numero massimo di batch in coda per l'elaborazione prima che le richieste non vengano completate rapidamente.

Aggiornamento delle opzioni di raggruppamento esistenti

Puoi aggiungere o aggiornare le opzioni di raggruppamento eseguendo l'immagine Docker specificando batch_options e impostando disable_default_optimizations su true utilizzando il --converter_options_string flag. Le opzioni batch verranno applicate a ogni funzione TPU o operazione di raggruppamento esistente.

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}
disable_default_optimizations=True

Requisiti per la forma del batching

I batch vengono creati concatenando tensori di input tra le richieste lungo il relativo batch (0°). I tensori di output vengono suddivisi in base alla dimensione 0. Nell'ordine per eseguire queste operazioni, l'operazione di raggruppamento in batch ha requisiti di forma rigorosi i suoi input e output.

Procedura dettagliata

Per comprendere questi requisiti, è utile prima capire come viene eseguito il raggruppamento. Nell'esempio seguente, stiamo raggruppando una semplice tf.matmul op.

def my_func(A, B)
    return tf.matmul(A, B)

La prima richiesta di inferenza produce gli input A e B con le forme (1, 3, 2) e (1, 2, 4) rispettivamente. La seconda richiesta di inferenza produce input A e B con le forme (2, 3, 2) e (2, 2, 4).

richiesta di inferenza 1

È stato raggiunto il timeout per il batch. Il modello supporta una dimensione del batch di 3 le richieste di inferenza n. 1 e n. 2 vengono raggruppate insieme senza spaziatura interna. La i tensori batch sono formati concatenando le richieste #1 e #2 lungo il batch (0°) . Poiché A di riga 1 ha una forma (1, 3, 2) e A di riga 2 ha una forma (2, 3, 2), quando vengono concatenati lungo la dimensione del batch (0), la forma risultante è (3, 3, 2).

richiesta batch

tf.matmul viene eseguito e produce un output con la forma (3, 3, 4).

richiesta matmul in batch

L'output di tf.matmul è raggruppato, pertanto deve essere suddiviso nuovamente in richieste separate. L'operazione di raggruppamento esegue questa operazione dividendo in base alla dimensione del batch (0) di ogni tensore di output. Decide come suddividere la dimensione 0 in base alla forma degli input originali. Poiché le forme della richiesta 1 hanno una dimensione 0 di 1, l'output ha una dimensione 0 di 1 per una forma di (1, 3, 4). Poiché le forme della richiesta 2 hanno una dimensione 0 pari a 2, l'output ha una dimensione 0 pari a 2 per una forma (2, 3, 4).

Risultati della richiesta di inferenza

Requisiti per le forme

Per eseguire l'operazione di concatenazione dell'input e la suddivisione dell'output descritta sopra, l'operazione di raggruppamento ha i seguenti requisiti di forma:

  1. Gli input per il raggruppamento non possono essere scalari. Per concatenare Dimensione 0°, i tensori devono avere almeno due dimensioni.

    Nella procedura dettagliata riportata sopra. Né A né B sono scalari.

    La mancata soddisfazione di questo requisito causerà un errore come: Batching input tensors must have at least one dimension. Una semplice correzione per questo errore è tradurre il vettore in un vettore.

  2. In diverse richieste di inferenza (ad esempio, diverse invocazioni di Esecuzione sessione), i tensori di input con lo stesso nome hanno le stesse dimensioni per ogni dimensione, tranne per la dimensione 0. Ciò consente di gestire gli input in modo concatenate nella loro 0a dimensione.

    Nella procedura dettagliata riportata sopra, la richiesta 1 ha la forma (1, 3, 2). Ciò significa che qualsiasi richiesta futura deve produrre una forma con il pattern(X, 3, 2). La richiesta 2 soddisfa questo requisito con (2, 3, 2). Allo stesso modo, la richiesta B della richiesta n. 1 ha una forma di (1, 2, 4), quindi tutte le richieste future devono produrre una forma con il motivo (X, 2, 4).

    La mancata soddisfazione di questo requisito causerà un errore come: Dimensions of inputs should match.

  3. Per una determinata richiesta di inferenza, tutti gli input devono avere lo stesso 0° dimensioni. Se diversi tensori di input all'operazione di raggruppamento hanno dimensioni diverse per la dimensione 0, l'operazione di raggruppamento non sa come suddividere i tensori di output.

    Nella procedura dettagliata precedente, i tensori della richiesta 1 hanno tutti una dimensione 0 di dimensione 1. In questo modo, l'operazione di raggruppamento sa che l'output deve avere una dimensione di 0 di dimensione pari a 1. Analogamente, i tensori della richiesta 2 hanno una dimensione 0 di dimensione 2, quindi l'output avrà una dimensione 0 di dimensione 2. Quando l'operazione di raggruppamento suddivide la forma finale di (3, 3, 4), produce (1, 3, 4) per la richiesta 1 e (2, 3, 4) per la richiesta 2.

    La mancata soddisfazione di questo requisito comporterà errori come: Batching input tensors supplied in a given op invocation must have equal 0th-dimension size.

  4. La dimensione della dimensione 0 della forma di ogni tensore di output deve essere la somma di tutte le dimensioni della dimensione 0 dei tensori di input (più eventuali spaziature introdotte dall'operazione di raggruppamento per soddisfare la dimensione allowed_batch_size successiva più grande). In questo modo, l'operazione di raggruppamento può suddividere i tensori di output in base alla dimensione 0 in base alla dimensione 0 dei tensori di input.

    Nella procedura dettagliata riportata sopra, i tensori di input hanno una 0a dimensione pari a 1 dalla richiesta n. 1 e 2 dalla richiesta n. 2. Pertanto, ogni tensore di output deve avere una dimensione 0 pari a 3 perché 1 + 2 = 3. Il tensore di output (3, 3, 4) che soddisfano questo requisito. Se 3 non fosse stata una dimensione del batch valida, ma 4 sì, l'operazione di raggruppamento avrebbe dovuto aggiungere elementi alla dimensione 0 degli input da 3 a 4. In questo caso, ogni tensore di output deve avere con una dimensione di 4.

    Se non soddisfi questo requisito, verrà visualizzato un errore simile a: Batched output tensor's 0th dimension does not equal the sum of the 0th dimension sizes of the input tensors.

Risoluzione degli errori relativi ai requisiti di forma

Per soddisfare questi requisiti, valuta la possibilità di fornire un'altra funzione o firma per creare un batch. Potrebbe anche essere necessario modificare le funzioni esistenti per soddisfare questi requisiti.

Se una funzione viene eseguita in batch, assicurati che tutte le forme della firma di input di @tf.function abbiano None nella dimensione 0 (ovvero la dimensione batch). Se viene eseguito il raggruppamento di una firma, assicurati che tutti i relativi input abbiano -1 nella dimensione 0.

L'operazione BatchFunction non supporta SparseTensors come input o output. Internamente, ogni tensore sparso è rappresentato come tre tensori distinti che possono avere dimensioni diverse per la dimensione 0.