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 di VM locale o TPU. La shell della VM TPU è consigliata perché è preinstallata gli strumenti a riga di comando necessari per il convertitore. Recupera un file SavedModel esportato ed esegue i seguenti passaggi:

1. Conversione TPU: aggiunge TPUPartitionedCall e altre operazioni TPU al modello per renderlo pubblicabile sulla TPU. Per impostazione predefinita, un modello esportato per l'inferenza non dispone di queste operazioni e non può essere pubblicato sulla TPU, anche se è stato addestrato sulla TPU.
2. Batch: aggiunge operazioni di batch al modello per abilitare questa funzionalità nel grafico e migliorare la velocità effettiva.
3. Conversione BFloat16: converte il formato dei dati del modello da float32 a bfloat16 per migliorare le prestazioni di calcolo e ridurre l'utilizzo della memoria HBM (High Bandwidth Memory) sulla TPU.
4. Ottimizzazione forma di I/O: ottimizza le forme dei tensori per i dati trasferiti tra CPU e TPU per migliorare l'utilizzo della larghezza di banda.

Durante l'esportazione di un modello, gli utenti creano alias di funzione per tutte le funzioni che vogliono eseguire sulla TPU. Passano queste funzioni al convertitore, che le posiziona nella TPU e le ottimizza.

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

Tempo stimato per completare i passaggi illustrati sopra: ~20 min - 30 min

Prerequisiti

1. Il modello deve essere un modello TF2 ed esportato nel formato SavedModel.
2. Il modello deve avere un alias di funzione per la funzione TPU. Per informazioni su come eseguire questa operazione, vedi l'esempio di codice. I seguenti esempi utilizzano tpu_func come alias della funzione TPU.
3. Assicurati che la CPU della tua macchina supporti le istruzioni AVX (Advanced Vector eXtensions), poiché la libreria Tensorflow (la dipendenza del convertitore di inferenza Cloud TPU) viene compilata in modo da utilizzare le istruzioni AVX. La maggior parte delle CPU supporta AVX.
   1. Puoi eseguire lscpu | grep avx per verificare se il set di istruzioni AVX è supportato.

Prima di iniziare

Prima di iniziare la configurazione, segui questi passaggi:

• Crea un nuovo progetto: nella pagina del selettore dei 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 una VM TPU esistente per eseguire l'inferenza con il modello convertito 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 il convertitore di inferenza Cloud TPU. Puoi installare Google Cloud SDK e Docker in locale oppure utilizzare una VM TPU in cui questo software è installato per impostazione predefinita. Questi strumenti 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 questi comandi per consentire l'utilizzo di Docker non root:

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

• Inizializza gli helper delle 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 delle credenziali Docker:

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

• Esegui il pull dell'immagine Docker del convertitore di inferenza:

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

Immagine convertitore

L'immagine serve a effettuare conversioni del modello una tantum. Imposta i percorsi del modello e regola le opzioni del convertitore in base alle tue esigenze. La sezione Esempi di utilizzo fornisce diversi casi d'uso comuni.

Nota: nel codice seguente, è possibile utilizzare un bucket Cloud Storage per input_model_dir e output_model_dir senza specificare mount nel comando docker.

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 esegua il wrapping di tutto ciò che vuoi eseguire sulla TPU. Se @tf.function non esiste, aggiungilo.
2. Quando salvi il modello, fornisci SaveOptions come indicato di seguito per assegnare a model.tpu_func un alias func_on_tpu.
3. Puoi passare questo alias di funzione all'utente che ha completato la conversione.

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 inserire più funzioni sulla TPU. È sufficiente creare più alias di funzione e trasmetterli in converter_options_string all'autore della conversione.

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

Quantizzazione

Importante: questa funzionalità è sperimentale e potrebbe essere soggetta a modifiche senza preavviso. Gli utenti devono valutare attentamente le prestazioni del modello quantizzato prima di utilizzarlo in produzione. Il supporto di TPU v5e per la quantizzazione è fornito in TensorFlow.

La quantizzazione è una tecnica che riduce la precisione dei numeri utilizzati per rappresentare i parametri di un modello. Le dimensioni del modello sono quindi ridotte e il calcolo più veloce. Un modello quantizzato consente di aumentare la velocità effettiva di inferenza, nonché di ridurre l'utilizzo della memoria e le dimensioni di archiviazione, a scapito di piccoli cali di accuratezza.

La nuova funzionalità di quantizzazione post-addestramento in TensorFlow che ha come target TPU è sviluppata a partire dall'analoga funzionalità esistente di TensorFlow Lite che viene utilizzata per il targeting di dispositivi mobili ed periferici. Per saperne di più sulla quantizzazione in generale, dai un'occhiata al documento di TensorFlow Lite.

Concetti di quantizzazione

Questa sezione definisce i concetti specificamente correlati alla quantizzazione con il convertitore di inferenza.

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

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

• Calibrazione: la fase di calibrazione per la quantizzazione è il processo di raccolta delle statistiche sull'intervallo di valori assunto dalle ponderazioni e dalle attivazioni di un modello di rete neurale. Queste informazioni vengono utilizzate per determinare i parametri di quantizzazione per il modello, che sono i valori che verranno utilizzati per convertire le ponderazioni in virgola mobile e le attivazioni 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 assunti dalle ponderazioni e dalle attivazioni del modello. Il set di dati rappresentativo deve soddisfare le seguenti proprietà:

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

  Il set di dati rappresentativo può essere un sottoinsieme del set di dati di addestramento oppure un set di dati separato progettato specificamente per rappresentare gli input reali del modello. La scelta del set di dati da utilizzare dipende dall'applicazione specifica.

• Static Range Quantization (SRQ): SRQ determina l'intervallo di valori per le ponderazioni e le attivazioni di un modello di rete neurale una sola volta, durante la fase di calibrazione. Ciò significa che viene utilizzato lo stesso intervallo di valori per tutti gli input al modello. Questa può essere meno precisa rispetto alla quantizzazione dell'intervallo dinamico, soprattutto per i modelli con un'ampia gamma di valori di input. Tuttavia, la quantizzazione dell'intervallo statico richiede meno calcolo in fase di runtime rispetto alla quantizzazione dell'intervallo dinamico.

• Dynamic Range Quantization (DRQ): DRQ determina l'intervallo di valori per le ponderazioni e le attivazioni di un modello di rete neurale per ogni input. Questo consente al modello di adattarsi all'intervallo di valori dei dati di input, migliorando l'accuratezza. Tuttavia, la quantizzazione dell'intervallo dinamico richiede più calcolo in fase di esecuzione rispetto alla quantizzazione dell'intervallo statico.

  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	Ulteriori informazioni di calcolo

• Quantizzazione di solo peso: la quantizzazione di solo peso è un tipo di quantizzazione che quantifica solo le ponderazioni 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 preservare l'accuratezza del modello.

Come utilizzare la quantizzazione

La quantizzazione può essere applicata configurando e impostando QuantizationOptions per le opzioni del convertitore. Le opzioni degne di nota sono:

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

Creazione del set di dati rappresentativo

Un set di dati rappresentativo è essenzialmente un'iterazione di campioni. Dove un esempio è 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 file 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 quantizza il modello con la chiave di firma di serving_default e l'alias funzione di tpu_func:

Nota: nel codice seguente, imposta CONVERTER_IMAGE come us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:nightly per utilizzare la build notturna.

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 batch a un modello. Per una descrizione delle opzioni di batch che possono essere ottimizzate, consulta la sezione Definizione delle opzioni di batch.

Per impostazione predefinita, il convertitore raggruppa tutte le funzioni TPU nel modello. Può anche eseguire il batch di firme e funzioni fornite dall'utente, per migliorare ulteriormente le prestazioni. Qualsiasi funzione TPU, funzione fornita dall'utente o firma in batch deve soddisfare i requisiti di forma rigorosi dell'operazione di batch.

Il convertitore può anche aggiornare le opzioni di raggruppamento in batch esistenti. Di seguito è riportato un esempio di come aggiungere il batch a un modello. Per ulteriori informazioni sul raggruppamento, consulta la sezione 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
}

Disabilita le ottimizzazioni della forma bfloat16 e IO

Le ottimizzazioni delle forme BFloat16 e IO sono abilitate per impostazione predefinita. Se non funzionano bene con il tuo modello, puoi disattivarli.

# 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 di conversione nel log dopo aver eseguito il convertitore di inferenza. 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 di calcolo del modello di output su CPU e TPU e ripartisce ulteriormente il costo della TPU per ogni funzione, che dovrebbe riflettere la selezione delle funzioni TPU nelle opzioni del convertitore.

Se vuoi utilizzare meglio la TPU, ti consigliamo di sperimentare la struttura del modello e modificare le opzioni del convertitore.

Nota: nell'attuale report sulle conversioni, tutte le operazioni nelle funzioni TPU sono conteggiate come costo TPU.

Domande frequenti

Quali funzioni devo inserire nella TPU?

È preferibile inserire il maggior numero possibile del modello sulla TPU, poiché la maggior parte delle operazioni viene eseguita più velocemente sulla TPU.

Se il modello non contiene op, stringhe o tensori di sparse non compatibili con TPU, l'inserimento dell'intero modello sulla TPU è in genere la strategia migliore. Puoi farlo trovando o creando una funzione che esegua il wrapping dell'intero modello, creando un alias di funzione per quest'ultimo e passandolo al convertitore.

Se il modello contiene parti che non possono funzionare sulla TPU (ad esempio operazioni, stringhe o tensori sparsi non compatibili con TPU), la scelta delle funzioni TPU dipende da dove si trova la parte incompatibile.

• Se il modello si trova all'inizio o alla fine, puoi eseguire il refactoring del modello per mantenerlo sulla CPU. Esempi sono le fasi di pre- e post-elaborazione delle stringhe. Per ulteriori informazioni sullo spostamento del codice sulla CPU, consulta la sezione "Come si sposta una parte del modello sulla CPU?" Mostra un modo tipico di refactoring del modello.
• Se si trova nel mezzo del modello, è meglio suddividerlo in tre parti, contenere nella parte centrale tutte le operazioni non compatibili con TPU e poi eseguirlo sulla CPU.
• Se si tratta di un tensore sparso, valuta la possibilità di chiamare tf.sparse.to_dense sulla CPU e di passare il tensore denso risultante alla parte TPU del modello.

Un altro fattore da considerare è l'utilizzo di HBM. Le tabelle di incorporamento possono utilizzare molta HBM. Se superano i limiti hardware della TPU, devono essere inseriti nella CPU, insieme alle operazioni di ricerca.

Quando possibile, deve esistere una sola funzione TPU sotto una firma. Se la struttura del tuo modello richiede la chiamata di più funzioni TPU per ogni richiesta di inferenza in arrivo, devi essere consapevole della maggiore latenza derivante dall'invio di tensori tra CPU e TPU.

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

Come posso spostare una parte del modello sulla CPU?

Se il modello contiene parti che non possono essere gestite sulla TPU, devi eseguire il refactoring del modello per spostarle nella CPU. Ecco un esempio di giocattolo. Il modello è un modello linguistico con una fase di pre-elaborazione. Il codice per le definizioni e le funzioni dei livelli viene omesso 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 fornito direttamente sulla TPU per due motivi. Il primo è una stringa. In secondo luogo, la funzione preprocess può contenere molte operazioni di stringa. Entrambi non sono compatibili con TPU.

Per il refactoring di questo modello, puoi creare un'altra funzione chiamata tpu_func per ospitare bert_layer ad alta intensità di calcolo. Quindi crea un alias di funzione per tpu_func e passalo al convertitore. In questo modo, tutto ciò che si trova all'interno di tpu_func verrà eseguito sulla TPU e tutto il resto in model_func verrà eseguito sulla 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 maggior parte delle operazioni TensorFlow standard è supportata sulla TPU, ma alcune, tra cui stringhe e tensori sparsi, non sono supportate. Il convertitore non controlla la presenza di operazioni non compatibili con TPU. Un modello che contiene queste operazioni può passare la conversione. Tuttavia, quando viene eseguita per l'inferenza, si verificano errori come quelli indicati di seguito.

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

Se il modello dispone di operazioni non compatibili con TPU, queste devono essere esterne alla funzione TPU. Inoltre, la stringa è un formato dati non supportato nella TPU. Di conseguenza, le variabili di tipo stringa non devono essere inserite nella funzione TPU. Inoltre, i parametri e i valori restituiti della funzione TPU non devono essere di tipo stringa. Allo stesso modo, evita di inserire tensori sparsi nella funzione TPU, inclusi i parametri e i valori restituiti.

Solitamente non è difficile eseguire il refactoring della parte incompatibile del modello e spostarla sulla 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 perché la libreria op dell'operazione personalizzata, che contiene la definizione completa dell'operazione, non è collegata al convertitore.

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

Cosa devo fare se ho 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 quando si esegue il mio modello?

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

Come posso convertire un modello già esportato senza un alias funzione?

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

Gli alias di funzione sono una mappatura da una stringa definita dall'utente a un nome di funzione concreto. Consentono di fare riferimento più facilmente a una funzione specifica nel modello. Il convertitore accetta sia alias di funzione sia nomi di funzioni non elaborati.

Puoi trovare i nomi delle funzioni concrete esaminando saved_model.pb.

L'esempio seguente mostra come inserire una funzione concreta denominata __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 posso risolvere un errore di vincolo della costante di tempo di compilazione?

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

Esistono due modi per risolvere questo problema.

• La soluzione migliore è aggiornare gli input dell'operazione in modo che abbiano una forma nota in modo statico quando XLA compila il programma TPU. Questa compilazione viene eseguita subito prima dell'esecuzione della parte TPU del modello. Ciò significa che la forma deve essere nota in modo statico prima dell'esecuzione di TpuFunction.
• Un'altra opzione è modificare TpuFunction per non includere più l'operazione problematica.

Perché viene visualizzato un errore di raggruppamento delle forme?

Il raggruppamento in batch ha requisiti di forma rigorosi che consentono di raggruppare le richieste in entrata lungo la 0a dimensione (nota anche come dimensione di batch). Questi requisiti di forma derivano dall'operazione di batch di TensorFlow e non possono essere rilassati.

Il mancato rispetto di questi requisiti comporterà errori quali:

1. I tensori di input di raggruppamento devono avere almeno una dimensione.
2. Le dimensioni degli input devono corrispondere.
3. I tensori di input di batch forniti in una determinata chiamata di un'operazione devono avere dimensioni uguali a 0.
4. La dimensione 0 del tensore di output in batch non corrisponde alla somma delle dimensioni 0a dei tensori di input.

Per soddisfare questi requisiti, valuta la possibilità di fornire una funzione o una firma diversa per il batch. Per soddisfare questi requisiti, potrebbe anche essere necessario modificare le funzioni esistenti.

Se una funzione è in fase di raggruppamento, assicurati che tutte le forme input_signature di @tf.function abbiano Nessuna nella 0a dimensione. Se una firma viene creata in batch, assicurati che tutti i relativi input abbiano -1 nella 0a dimensione.

Per una spiegazione completa sul motivo per cui si verificano questi errori e su come risolverli, consulta la sezione Batch di approfondimenti.

Problemi noti

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

Sebbene il convertitore sia in grado di gestire la maggior parte degli scenari di chiamata di funzione attraverso il confine CPU-TPU, esiste un raro caso limite in cui l'operazione non riesce. È quando una funzione TPU chiama indirettamente un'altra funzione TPU.

Questo perché il convertitore modifica il chiamante diretto di una funzione TPU dalla chiamata della funzione TPU stessa alla chiamata di uno stub di chiamata TPU. Lo stub della chiamata contiene operazioni che possono funzionare solo sulla CPU. Quando una funzione TPU chiama una funzione che alla fine chiama il chiamante diretto, le operazioni della CPU potrebbero essere spostate sulla TPU per l'esecuzione, il che genererà errori del 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 viene visualizzato il seguente errore, significa che il modello ha raggiunto 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 è il refactoring del modello per evitare uno scenario che richiama una funzione. Se ritieni che sia difficile, contatta il team dell'Assistenza Google per parlarne.

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 batch

Il batch viene utilizzato per migliorare la velocità effettiva e l'utilizzo di TPU. Consente di elaborare più richieste contemporaneamente. Durante l'addestramento, è possibile eseguire il raggruppamento in batch utilizzando tf.data. Durante l'inferenza, in genere viene eseguita aggiungendo un'operazione nel grafico che raggruppa le richieste in entrata. L'operazione attende che abbia un numero sufficiente di richieste o che venga raggiunto un timeout prima di generare un batch di grandi dimensioni dalle singole richieste. Per ulteriori informazioni sulle diverse opzioni di raggruppamento che possono essere ottimizzate, incluse dimensioni e timeout dei batch, consulta la sezione Definizione delle opzioni di raggruppamento.

Per impostazione predefinita, il convertitore inserisce l'operazione di batch direttamente prima del calcolo della TPU. Aggrega le funzioni TPU fornite dall'utente e qualsiasi calcolo TPU preesistente nel modello con operazioni di batch. È possibile eseguire l'override di questo comportamento predefinito indicando al convertitore quali funzioni e/o firme devono essere raggruppate.

L'esempio seguente mostra come aggiungere il raggruppamento predefinito.

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
}

Raggruppamento delle firme

Nota: questa funzionalità è ancora in fase sperimentale. Dopo il test, il campo verrà spostato fuori dal wrapper sperimentale.

Il raggruppamento delle firme in batch raggruppa l'intero modello partendo dagli input della firma e andando agli output della firma. A differenza del comportamento di batching predefinito del convertitore, il raggruppamento delle firme in batch associa sia il calcolo TPU sia quello della CPU. Questo offre un aumento delle prestazioni del 10-20% durante l'inferenza su alcuni modelli.

Come tutti i batch, anche quelli delle firme hanno requisiti rigorosi per la forma. Per garantire che vengano soddisfatti questi requisiti di forma, gli input delle firme devono avere forme che abbiano almeno due dimensioni. La prima è la dimensione batch e deve avere -1. Ad esempio, (-1, 4), (-1) o (-1, 128, 4, 10) sono tutte forme di input valide. Se questo non è possibile, valuta la possibilità di utilizzare il comportamento predefinito di raggruppamento o il batch delle funzioni.

Per utilizzare il raggruppamento delle firme in batch, 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"
  }
}

Batch di funzioni

Nota: questa funzionalità è ancora in fase sperimentale. Dopo il test, il campo verrà spostato fuori dal wrapper sperimentale.

Il raggruppamento delle funzioni può essere utilizzato per indicare al convertitore le funzioni da eseguire in batch. Per impostazione predefinita, il convertitore raggruppa tutte le funzioni TPU. Il batch di funzioni sostituisce questo comportamento predefinito.

Il batch di funzioni può essere utilizzato per il calcolo in batch della CPU. Molti modelli registrano un miglioramento delle prestazioni quando il calcolo della CPU viene suddiviso in batch. Il modo migliore per eseguire in batch il calcolo della CPU è utilizzare il batch delle firme, ma potrebbe non funzionare con alcuni modelli. In questi casi, è possibile utilizzare il batch di funzioni per eseguire il batch di parte del calcolo della CPU oltre al calcolo della TPU. Tieni presente che l'operazione di raggruppamento non può essere eseguita sulla TPU, pertanto qualsiasi funzione di raggruppamento fornita deve essere chiamata nella CPU.

Il batch di funzioni può essere utilizzato anche per soddisfare i requisiti di forma rigorosi imposti dall'operazione di raggruppamento. Nei casi in cui le funzioni TPU non soddisfino i requisiti di forma dell'operazione di batch, è possibile utilizzare il batch delle funzioni per indicare all'autore della conversione di raggruppare diverse funzioni.

Per utilizzarlo, genera un valore function_alias per la funzione da creare in batch. Puoi farlo trovando o creando una funzione nel modello che esegue il wrapping di tutto ciò che vuoi in batch. Assicurati che questa funzione soddisfi i requisiti di forma rigida imposti dall'operazione di raggruppamento. Aggiungi @tf.function se non ne ha già uno. È importante fornire input_signature a @tf.function. La 0a dimensione deve essere None perché è la dimensione batch, quindi non può avere una dimensione fissa. Ad esempio, [None, 4], [None] o [None, 128, 4, 10] sono tutte forme di input valide. Quando salvi il modello, fornisci SaveOptions come quelli mostrati di seguito per assegnare a model.batch_func l'alias "batch_func". Dopodiché puoi passare l'alias di questa funzione all'utente che ha completato una conversione.

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)

Quindi, 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

• num_batch_threads: (numero intero) numero di thread di pianificazione per l'elaborazione di batch di lavoro. Determina il numero di batch elaborati in parallelo. Questo dovrebbe essere approssimativamente in linea con il numero di core TPU disponibili.
• max_batch_size: (numero intero) dimensione massima del batch consentita. Può essere superiore a allowed_batch_sizes per utilizzare la suddivisione in batch di grandi dimensioni.
• batch_timeout_micros: (numero intero) numero massimo di microsecondi da attendere prima di generare un batch incompleto.
• allowed_batch_sizes: (elenco di numeri interi) se l'elenco non è vuoto, i gruppi verranno compressi fino alla dimensione più vicina nell'elenco. L'elenco deve essere in aumento monotonico e l'elemento finale deve essere inferiore o uguale a max_batch_size.
• max_enqueued_batches: (numero intero) numero massimo di batch accodati per l'elaborazione prima che le richieste non vadano a buon fine.

Aggiornamento delle opzioni di raggruppamento esistenti

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

Nota: non puoi aggiornare le opzioni di raggruppamento impostando tpu_functions insieme utilizzando --converter_options_string.

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 le forme di raggruppamento

I batch vengono creati concatenando i tensori di input tra le richieste lungo la dimensione batch (0a). I tensori di output sono suddivisi lungo la loro dimensione 0. Per eseguire queste operazioni, l'operazione di raggruppamento ha requisiti di forma rigorosi per i propri input e output.

Procedura dettagliata

Per comprendere questi requisiti, è utile comprendere prima come viene eseguito il batch. Nell'esempio riportato di seguito, raggruppiamo una semplice operazione tf.matmul.

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 gli input A e B con le forme (2, 3, 2) e (2, 2, 4).

È stato raggiunto il timeout del raggruppamento. Il modello supporta una dimensione batch pari a 3, pertanto le richieste di inferenza 1 e 2 vengono raggruppate senza spaziatura interna. I tensori in batch vengono formati concatenando le richieste 1 e 2 lungo la dimensione batch (0). Poiché A del numero 1 ha la forma (1, 3, 2) e la A del numero 2 ha la forma (2, 3, 2), quando sono concatenate lungo la dimensione batch (0), la forma risultante è (3, 3, 2).

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

L'output di tf.matmul è in batch, quindi deve essere suddiviso nuovamente in richieste separate. L'operazione di raggruppamento crea la suddivisione in base alla dimensione batch (0a) di ogni tensore di output. Decide come suddividere la 0a dimensione in base alla forma degli input originali. Poiché le forme della richiesta n. 1 hanno una dimensione 0 pari a 1, il relativo output ha una dimensione 0 pari a 1 per una forma di (1, 3, 4). Poiché le forme della richiesta n. 2 hanno una dimensione 0 pari a 2, il relativo output ha una dimensione 0 pari a 2 per una forma di (2, 3, 4).

Requisiti delle forme

Per eseguire la concatenazione degli 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 concatenarsi lungo la dimensione 0, i tensori devono avere almeno due dimensioni.

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

   Il mancato rispetto di questo requisito causerà un errore come: Batching input tensors must have at least one dimension. Una semplice correzione di questo errore è rendere lo scalare un vettore.

2. In richieste di inferenza diverse (ad esempio chiamate di esecuzione di sessioni diverse), i tensori di input con lo stesso nome hanno le stesse dimensioni per ogni dimensione, tranne la 0a. Ciò consente di concatenare gli input in modo pulito nella 0a dimensione.

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

   Il mancato rispetto di questo requisito causerà un errore come: Dimensions of inputs should match.

3. Per una determinata richiesta di inferenza, tutti gli input devono avere la stessa dimensione 0. Se diversi tensori di input dell'operazione di raggruppamento hanno dimensioni 0 diverse, il team operativo di raggruppamento non sa come suddividere i tensori di output.

   Nella procedura dettagliata precedente, i tensori della richiesta n. 1 hanno tutti una dimensione di 0a dimensione pari a 1. Ciò consente all'operazione di raggruppamento di sapere che l'output deve avere una dimensione 0 pari a 1. Allo stesso modo, i tensori della richiesta 2 hanno una dimensione di 0 a 2, quindi il suo output avrà una dimensione di dimensione 0 pari a 2. Quando l'operazione di batch suddivide la forma finale di (3, 3, 4), produce (1, 3, 4) per la richiesta 1 e (2, 3, 4) per la richiesta 2.

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

4. La 0a dimensione di dimensione della forma di ogni tensore di output deve essere la somma di tutte le dimensioni di 0a dimensione di tutti i tensori di input (oltre a qualsiasi spaziatura interna introdotta dall'operazione di raggruppamento per soddisfare il valore allowed_batch_size più grande successivo). In questo modo l'operazione di raggruppamento può suddividere i tensori di output lungo la loro 0a dimensione in base alla 0a dimensione dei tensori di input.

   Nella procedura dettagliata precedente, i tensori di input hanno una dimensione 0 pari a 1 dalla richiesta 1 e 2 dalla richiesta 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) soddisfa questo requisito. Se 3 non fosse una dimensione batch valida, ma 4 lo fosse, l'operazione di batch avrebbe dovuto occupare la 0a dimensione degli input da 3 a 4. In questo caso, ogni tensore di output dovrebbe avere una dimensione 0 pari a 4.

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

Risoluzione degli errori dei requisiti delle forme

Per soddisfare questi requisiti, valuta la possibilità di fornire una funzione o una firma diversa per il batch. Per soddisfare questi requisiti, potrebbe anche essere necessario modificare le funzioni esistenti.

Se una funzione viene eseguita in batch, assicurati che tutte le forme input_signature di @tf.function abbiano None nella 0a dimensione (ovvero la dimensione batch). Se una firma viene eseguita in batch, 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 da tre tensori separati che possono avere dimensioni di 0 dimensioni diverse.