Présentation du convertisseur d'inférence Cloud TPU v5e

Introduction

Le convertisseur d'inférences Cloud TPU prépare et optimise un modèle TensorFlow 2 (TF2) pour l'inférence TPU. Le convertisseur s'exécute dans une interface système de VM locale ou TPU. L'interface système de la VM TPU est recommandée, car les outils de ligne de commande nécessaires au convertisseur y sont préinstallés. Il utilise un SavedModel exporté et effectue les étapes suivantes:

  1. Conversion de TPU: elle ajoute TPUPartitionedCall et d'autres opérations TPU au modèle pour le rendre diffusable sur le TPU. Par défaut, un modèle exporté pour l'inférence ne comporte pas ce type d'opérations et ne peut pas être diffusé sur le TPU, même s'il a été entraîné sur ce dernier.
  2. Traitement par lot: des opérations de traitement par lot sont ajoutées au modèle pour activer le traitement par lot dans le graphique et améliorer le débit.
  3. Conversion BFloat16: elle convertit le format de données du modèle float32 en bfloat16 pour améliorer les performances de calcul et réduire l'utilisation de la mémoire à haut débit (HBM) sur le TPU.
  4. Optimisation de la forme d'E/S: optimise les formes de Tensor pour les données transférées entre le processeur et le TPU, afin d'améliorer l'utilisation de la bande passante.

Lors de l'exportation d'un modèle, les utilisateurs créent des alias pour toutes les fonctions qu'ils souhaitent exécuter sur le TPU. Ils transmettent ces fonctions au convertisseur, qui les place sur le TPU et les optimise.

Le convertisseur d'inférences Cloud TPU est disponible sous la forme d'une image Docker pouvant être exécutée dans n'importe quel environnement sur lequel Docker est installé.

Temps estimé pour effectuer les étapes ci-dessus: environ 20 min à 30 min

Prérequis

  1. Le modèle doit être un modèle TF2 et exporté au format SavedModel.
  2. Le modèle doit avoir un alias de fonction pour la fonction TPU. Consultez l'exemple de code pour savoir comment procéder. Les exemples suivants utilisent tpu_func comme alias de fonction TPU.
  3. Assurez-vous que le processeur de votre machine est compatible avec les instructions AVX (Advanced Vector eXtensions), car la bibliothèque TensorFlow (la dépendance du convertisseur d'inférences Cloud TPU) est compilée de manière à utiliser les instructions AVX. La plupart des processeurs sont compatibles avec AVX.
    1. Vous pouvez exécuter lscpu | grep avx pour vérifier si l'ensemble d'instructions AVX est compatible.

Avant de commencer

Avant de commencer la configuration, procédez comme suit:

  • Créer un projet : dans la console Google Cloud, sur la page du sélecteur de projet, sélectionnez ou créez un projet Cloud.

  • Configurer une VM TPU : créez une VM TPU à l'aide de la console Google Cloud ou de gcloud, ou utilisez une VM TPU existante pour exécuter des inférences avec le modèle converti sur la VM TPU.

    • Assurez-vous que l'image de VM TPU est basée sur TensorFlow. Par exemple, --version=tpu-vm-tf-2.11.0.
    • Le modèle converti sera chargé et diffusé sur cette VM TPU.
  • Assurez-vous de disposer des outils de ligne de commande nécessaires pour utiliser le convertisseur d'inférence Cloud TPU. Vous pouvez installer Google Cloud SDK et Docker localement ou utiliser une VM TPU sur laquelle ce logiciel est installé par défaut. Ces outils vous permettent d'interagir avec l'image du convertisseur.

    Connectez-vous à l'instance avec SSH à l'aide de la commande suivante:

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

Configuration de l'environnement

Configurez votre environnement à partir de l'interface système de votre VM TPU ou de votre interface système locale.

Shell de VM TPU

  • Dans l'interface système de la VM TPU, exécutez les commandes suivantes pour autoriser l'utilisation de fichiers Docker non racines:

    sudo usermod -a -G docker ${USER}
    newgrp docker
    
  • Initialisez vos assistants d'identification Docker:

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

Interface système locale

Dans votre interface système locale, configurez l'environnement en procédant comme suit:

  • Installez le SDK Cloud, qui inclut l'outil de ligne de commande gcloud.

  • Installez Docker:

  • Autorisez l'utilisation de Docker non racine:

    sudo usermod -a -G docker ${USER}
    newgrp docker
    
  • Connectez-vous à votre environnement:

    gcloud auth login
    
  • Initialisez vos assistants d'identification Docker:

    gcloud auth configure-docker \
        us-docker.pkg.dev
    
  • Extrayez l'image Docker du convertisseur d'inférences:

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

Image du convertisseur

L'image sert à effectuer des conversions de modèles ponctuelles. Définissez les chemins d'accès au modèle et ajustez les options de conversion en fonction de vos besoins. La section Exemples d'utilisation présente plusieurs cas d'utilisation courants.

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
    }
'

Inférence avec le modèle converti dans une 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))

Exemples d'utilisation

Ajouter un alias de fonction pour la fonction TPU

  1. Recherchez ou créez dans votre modèle une fonction qui encapsule tout ce que vous souhaitez exécuter sur le TPU. Si @tf.function n'existe pas, ajoutez-le.
  2. Lors de l'enregistrement du modèle, indiquez SaveOptions comme ci-dessous pour attribuer un alias func_on_tpu à model.tpu_func.
  3. Vous pouvez transmettre cet alias de fonction au convertisseur.
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)

Convertir un modèle avec plusieurs fonctions TPU

Vous pouvez placer plusieurs fonctions sur le TPU. Il vous suffit de créer plusieurs alias de fonction et de les transmettre dans converter_options_string au convertisseur.

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

Quantification

La quantification est une technique qui réduit la précision des nombres utilisés pour représenter les paramètres d'un modèle. Le modèle réduit ainsi la taille et les calculs sont plus rapides. Un modèle quantifié permet d'améliorer le débit d'inférence ainsi qu'une utilisation de la mémoire et une taille de stockage réduites, au détriment de la justesse.

La nouvelle fonctionnalité de quantification post-entraînement de TensorFlow, qui cible TPU, est développée à partir de la fonctionnalité similaire existante de TensorFlow Lite, qui est utilisée pour cibler les appareils mobiles et de périphérie. Pour en savoir plus sur la quantification en général, consultez le document de TensorFlow Lite.

Concepts de quantification

Cette section définit des concepts spécifiquement liés à la quantification avec le convertisseur d'inférence.

Les concepts liés à d'autres configurations de TPU (par exemple, les tranches, les hôtes, les puces et les TensorCore) sont décrits sur la page Architecture du système TPU.

  • Quantification post-entraînement (PTQ): la technique du PTQ réduit la taille et la complexité de calcul d'un modèle de réseau de neurones sans affecter de manière significative sa justesse. La méthode PTQ convertit les pondérations et activations à virgule flottante d'un modèle entraîné en entiers de précision inférieure, tels que des entiers de 8 bits ou 16 bits. Cela peut réduire considérablement la taille du modèle et la latence d'inférence, tout en n'entraînant qu'une légère perte de précision.

  • Calibrage: l'étape d'étalonnage pour la quantification consiste à collecter des statistiques sur la plage de valeurs prises en compte par les pondérations et les activations d'un modèle de réseau de neurones. Ces informations permettent de déterminer les paramètres de quantification du modèle, qui sont les valeurs qui seront utilisées pour convertir les pondérations et les activations à virgule flottante en entiers.

  • Ensemble de données représentatif: un ensemble de données représentatif pour la quantification est un petit ensemble de données qui représente les données d'entrée réelles du modèle. Elle est utilisée lors de l'étape de calibrage de la quantification pour recueillir des statistiques sur la plage de valeurs prise en compte par les pondérations et les activations du modèle. L'ensemble de données représentatif doit satisfaire aux propriétés suivantes:

    • Il doit représenter correctement les entrées réelles du modèle lors de l'inférence. Cela signifie qu'il doit couvrir la plage de valeurs que le modèle est susceptible de voir dans le monde réel.
    • Il doit traverser collectivement chaque branche d'expressions conditionnelles (par exemple, tf.cond), le cas échéant. Ce point est important, car le processus de quantification doit pouvoir gérer toutes les entrées possibles du modèle, même si elles ne sont pas explicitement représentées dans l'ensemble de données représentatif.
    • Il doit être suffisamment grand pour collecter suffisamment de statistiques et réduire les erreurs. En règle générale, il est recommandé d'utiliser plus de 200 échantillons représentatifs.

    L'ensemble de données représentatif peut être un sous-ensemble de l'ensemble de données d'entraînement ou un ensemble de données distinct spécialement conçu pour être représentatif des entrées réelles du modèle. Le choix de l'ensemble de données à utiliser dépend de l'application spécifique.

  • Quadrant de la plage statique (SRC): la SRC détermine la plage de valeurs pour les pondérations et les activations d'un modèle de réseau de neurones une seule fois, au cours de l'étape d'étalonnage. Cela signifie que la même plage de valeurs est utilisée pour toutes les entrées du modèle. Cette méthode peut s'avérer moins précise que la quantification de plage dynamique, en particulier pour les modèles comportant un large éventail de valeurs d'entrée. Cependant, la quantification de plage statique nécessite moins de calculs au moment de l'exécution que la quantification de plage dynamique.

  • Quadrant de la plage dynamique (DRC): il détermine la plage de valeurs pour les pondérations et les activations d'un modèle de réseau de neurones pour chaque entrée. Cela permet au modèle de s'adapter à la plage de valeurs des données d'entrée, ce qui peut améliorer la précision. Cependant, la quantification de plage dynamique nécessite plus de calculs au moment de l'exécution que la quantification de plage statique.

    Fonctionnalité Quantification de la plage statique Quantification de la plage dynamique
    Plage de valeurs Déterminé une fois, lors de l'étalonnage Déterminé pour chaque entrée
    Justesse Peut être moins précis, en particulier pour les modèles comportant un large éventail de valeurs d'entrée peut être plus précis, en particulier pour les modèles comportant un large éventail de valeurs d'entrée ;
    de la complexité Plus simple Plus complexe
    Calcul au moment de l'exécution Moins de calculs Autres calculs
  • Quantification pondérée uniquement: la quantification pondérée uniquement est un type de quantification qui ne quantifie que les pondérations d'un modèle de réseau de neurones, tout en laissant les activations en virgule flottante. Cette option peut s'avérer intéressante pour les modèles sensibles à la précision, car elle permet de préserver la justesse du modèle.

Utiliser la quantification

Vous pouvez appliquer la quantification en configurant et en définissant QuantizationOptions sur les options du convertisseur. Les options notables sont les suivantes:

  • tags: collection de balises identifiant la MetaGraphDef dans la SavedModel à quantifier. Il n'est pas nécessaire de préciser si vous n'avez qu'un seul MetaGraphDef.
  • signature_keys: séquence de clés identifiant SignatureDef contenant des entrées et des sorties. Si aucune valeur n'est spécifiée, ["serving_default"] est utilisé.
  • quantization_method: méthode de quantification à appliquer. Si aucune valeur n'est spécifiée, la quantification STATIC_RANGE sera appliquée.
  • op_set: doit être conservé au format XLA. Il s'agit actuellement de l'option par défaut. Vous n'avez pas besoin de la spécifier.
  • représenter_datasets: spécifier l'ensemble de données utilisé pour calibrer les paramètres de quantification.

Créer l'ensemble de données représentatif

Un ensemble de données représentatif est essentiellement un itérable d'échantillons. Où un échantillon est une carte de: {input_key: input_value}. Exemple :

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

Les ensembles de données représentatifs doivent être enregistrés sous forme de fichiers TFRecord à l'aide de la classe TfRecordRepresentativeDatasetSaver actuellement disponible dans le package pip tf-nightly. Exemple :

# 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})

Examples

L'exemple suivant quantifie le modèle avec la clé de signature serving_default et l'alias de fonction 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}" \
          } \
        } \
      } \
    } '

Ajouter un traitement par lot

L'utilitaire de conversion permet d'ajouter le traitement par lot à un modèle. Pour obtenir une description des options de traitement par lot pouvant être ajustées, consultez la section Définition des options de traitement par lot.

Par défaut, le convertisseur regroupe toutes les fonctions TPU du modèle. Il peut également regrouper les signatures et fonctions fournies par l'utilisateur par lot, ce qui peut améliorer davantage les performances. Toute fonction TPU, une fonction fournie par l'utilisateur ou une signature par lot, doit répondre aux exigences strictes concernant la forme de l'opération de traitement par lot.

Le convertisseur peut également mettre à jour les options de traitement par lot existantes. Voici un exemple d'ajout du traitement par lot à un modèle. Pour en savoir plus sur le traitement par lot, consultez la page Présentation détaillée du traitement par lot.

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
}

Désactiver les optimisations de forme bfloat16 et d'E/S

Les optimisations de forme BFloat16 et E/S sont activées par défaut. Si elles ne fonctionnent pas bien avec votre modèle, vous pouvez les désactiver.

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

Rapport sur les conversions

Vous pouvez trouver ce rapport de conversion à partir du journal après avoir exécuté Inference Converter. Vous trouverez un exemple ci-dessous.

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

Ce rapport estime le coût de calcul du modèle de sortie sur le processeur et le TPU, et détaille le coût du TPU pour chaque fonction, ce qui doit refléter votre sélection des fonctions TPU dans les options du convertisseur.

Si vous souhaitez mieux utiliser le TPU, vous pouvez tester la structure du modèle et ajuster les options du convertisseur.

Questions fréquentes

Quelle(s) fonction(s) dois-je placer sur le TPU ?

Il est préférable de placer autant de modèles que possible sur le TPU, car la grande majorité des opérations s'exécutent plus rapidement sur ce TPU.

Si votre modèle ne contient aucune opération, chaîne ou Tensor creux incompatible avec le TPU, la meilleure stratégie consiste généralement à placer le modèle entier sur le TPU. Pour ce faire, recherchez ou créez une fonction qui encapsule l'intégralité du modèle, puis créez un alias de fonction pour cette fonction et transmettez-le au convertisseur.

Si votre modèle contient des parties qui ne peuvent pas fonctionner sur le TPU (par exemple, des opérations incompatibles avec le TPU, des chaînes ou des Tensors creux), le choix des fonctions TPU dépend de l'emplacement de la partie incompatible.

  • Si elle se trouve au début ou à la fin du modèle, vous pouvez le refactoriser pour le conserver sur le processeur. Il s'agit par exemple des étapes de prétraitement et de post-traitement des chaînes. Pour en savoir plus sur le déplacement du code vers le processeur, consultez la section Comment déplacer une partie du modèle vers le processeur. Il montre une méthode classique de refactorisation du modèle.
  • Si elle se trouve au milieu du modèle, il est préférable de le diviser en trois parties, de contenir toutes les opérations incompatibles avec le TPU dans la partie intermédiaire, puis de l'exécuter sur le processeur.
  • S'il s'agit d'un Tensor creux, envisagez d'appeler tf.sparse.to_dense sur le processeur et de transmettre le Tensor dense obtenu à la partie TPU du modèle.

Un autre facteur à prendre en compte est l'utilisation du HBM. La représentation vectorielle continue de tableaux peut nécessiter beaucoup d'outils HBM. Si elles dépassent les limites matérielles du TPU, elles doivent être placées sur le processeur, en même temps que les opérations de recherche.

Dans la mesure du possible, une seule fonction TPU doit exister sous une même signature. Si la structure de votre modèle nécessite d'appeler plusieurs fonctions TPU par requête d'inférence entrante, vous devez tenir compte de la latence supplémentaire liée à l'envoi de Tensors entre le processeur et le TPU.

Un bon moyen d'évaluer la sélection des fonctions TPU consiste à consulter le rapport sur les conversions. Elle indique le pourcentage de calculs effectués sur le TPU, ainsi qu'une répartition du coût de chaque fonction TPU.

Comment déplacer une partie du modèle vers le processeur ?

Si votre modèle contient des parties qui ne peuvent pas être diffusées sur le TPU, vous devez refactoriser le modèle pour les déplacer vers le processeur. Voici un exemple de jouet. Il s'agit d'un modèle de langage comportant une étape de prétraitement. Par souci de simplicité, le code des définitions et des fonctions des couches est omis.

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)

Ce modèle ne peut pas être diffusé directement sur le TPU pour deux raisons. Tout d'abord, le paramètre est une chaîne. Deuxièmement, la fonction preprocess peut contenir de nombreuses opérations de chaîne. Les deux ne sont pas compatibles avec le TPU.

Pour refactoriser ce modèle, vous pouvez créer une autre fonction appelée tpu_func afin d'héberger la fonction bert_layer qui utilise beaucoup de ressources de calcul. Créez ensuite un alias de fonction pour tpu_func et transmettez-le au convertisseur. De cette manière, tout ce qui se trouve dans tpu_func s'exécutera sur le TPU, et tout ce qui reste dans model_func s'exécutera sur le processeur.

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)

Que dois-je faire si le modèle comporte des opérations, des chaînes ou des Tensors creux incompatibles avec le TPU ?

La plupart des opérations TensorFlow standards sont compatibles avec le TPU, mais quelques-unes, y compris les Tensors creux et les chaînes, ne le sont pas. Le convertisseur ne recherche pas les opérations incompatibles avec le TPU. Un modèle contenant ces opérations peut donc transmettre la conversion. Toutefois, lors de son exécution à des fins d'inférence, des erreurs telles que celles indiquées ci-dessous se produiront.

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

Si votre modèle comporte des opérations incompatibles avec le TPU, celles-ci doivent être placées en dehors de la fonction TPU. De plus, le format de chaîne n'est pas pris en charge sur le TPU. Les variables de type chaîne ne doivent donc pas être placées dans la fonction TPU. De plus, les paramètres et les valeurs de retour de la fonction TPU ne doivent pas être de type chaîne. De même, évitez de placer des Tensors creux dans la fonction TPU, y compris dans ses paramètres et ses valeurs de retour.

Il n'est généralement pas difficile de refactoriser la partie incompatible du modèle et de la déplacer vers le processeur. Voici un exemple.

Comment prendre en charge les opérations personnalisées dans le modèle ?

Si des opérations personnalisées sont utilisées dans votre modèle, le convertisseur risque de ne pas les reconnaître et de ne pas parvenir à convertir le modèle. En effet, la bibliothèque d'opérations de l'opération personnalisée, qui contient la définition complète de l'opération, n'est pas liée au convertisseur.

Étant donné que le code du convertisseur n'est pas encore en Open Source, il ne peut pas être créé avec une opération personnalisée.

Que dois-je faire si j'ai un modèle TensorFlow 1 ?

Le convertisseur n'est pas compatible avec les modèles TensorFlow 1. Les modèles TensorFlow 1 doivent être migrés vers TensorFlow 2.

Dois-je activer le pont MLIR lors de l'exécution de mon modèle ?

La plupart des modèles convertis peuvent être exécutés avec le pont MLIR TF2XLA le plus récent ou le pont TF2XLA d'origine.

Comment convertir un modèle déjà exporté sans alias de fonction ?

Si un modèle a été exporté sans alias de fonction, le moyen le plus simple consiste à l'exporter à nouveau et à créer un alias de fonction. Si la réexportation n'est pas possible, il est toujours possible de convertir le modèle en fournissant un concrete_function_name. Cependant, l'identification du bon concrete_function_name nécessite un travail de détection.

Les alias de fonction correspondent à un mappage entre une chaîne définie par l'utilisateur et un nom de fonction concret. Ils facilitent la référence à une fonction spécifique du modèle. Le convertisseur accepte les alias de fonction et les noms de fonction bruts concrets.

Vous pouvez trouver des noms de fonction concrets en examinant saved_model.pb.

L'exemple suivant montre comment placer une fonction concrète appelée __inference_serve_24 sur le 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"
    }'

Comment résoudre une erreur de contrainte de constante de temps de compilation ?

Pour l'entraînement et l'inférence, XLA nécessite que les entrées de certaines opérations aient une forme connue au moment de la compilation du TPU. Cela signifie que lorsque XLA compile la partie TPU du programme, les entrées de ces opérations doivent avoir une forme connue statiquement.

Il existe deux façons de résoudre ce problème.

  • La meilleure option consiste à mettre à jour les entrées de l'opération afin qu'elles aient une forme connue de manière statique au moment où XLA compile le programme TPU. Cette compilation est effectuée juste avant l'exécution de la partie TPU du modèle. Cela signifie que la forme doit être connue de manière statique au moment où TpuFunction est sur le point de s'exécuter.
  • Une autre option consiste à modifier TpuFunction pour ne plus inclure l'opération problématique.

Pourquoi une erreur de traitement par lot s'affiche-t-elle ?

Le traitement par lot présente des exigences strictes en matière de forme qui permettent aux requêtes entrantes d'être regroupées selon leur 0e dimension (également appelée "dimension de traitement par lot"). Ces exigences de forme proviennent de l'opération de traitement par lot TensorFlow et ne peuvent pas être assouplies.

Si vous ne respectez pas ces exigences, des erreurs se produiront. Par exemple:

  1. Le traitement par lot des Tensors d'entrée doit avoir au moins une dimension.
  2. Les dimensions des entrées doivent correspondre.
  3. Le traitement par lot des Tensors d'entrée fournis dans un appel d'opération donné doit avoir la même taille de 0 dimension.
  4. La dimension 0 du Tensor de sortie par lot n'est pas égale à la somme des tailles 0 des Tensors d'entrée.

Pour répondre à ces exigences, envisagez de fournir une fonction ou une signature différente pour le lot. Il peut également être nécessaire de modifier des fonctions existantes pour répondre à ces exigences.

Si une fonction est par lot, assurez-vous que les formes de la @tf.function input_signature ont toutes la valeur None dans la 0e dimension. Si une signature est par lot, assurez-vous que toutes ses entrées présentent la valeur -1 dans la dimension 0.

Pour obtenir une explication complète des causes de ces erreurs et de leur résolution, consultez la page Présentation détaillée du traitement par lot.

Problèmes connus

La fonction TPU ne peut pas appeler indirectement une autre fonction TPU

Bien que le convertisseur puisse gérer la plupart des scénarios d'appel de fonction au-delà de la limite CPU-TPU, il existe un cas limite rare qui échouerait. C'est lorsqu'une fonction TPU appelle indirectement une autre fonction TPU.

En effet, le convertisseur modifie l'appelant direct d'une fonction TPU en appelant la fonction TPU elle-même en appelant un bouchon d'appel TPU. Le bouchon d'appel contient des opérations qui ne peuvent fonctionner que sur le processeur. Lorsqu'une fonction TPU appelle une fonction qui appelle finalement l'appelant direct, ces opérations de processeur peuvent être transférées sur le TPU pour être exécutées, ce qui génère des erreurs de noyau manquantes. Notez que ce cas est différent d'une fonction TPU qui appelle directement une autre fonction TPU. Dans ce cas, le convertisseur ne modifie aucune des fonctions pour appeler le bouchon d'appel. Ainsi, elle peut fonctionner.

Dans le convertisseur, nous avons implémenté la détection de ce scénario. Si l'erreur suivante s'affiche, cela signifie que votre modèle a atteint ce cas 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 solution générale consiste à refactoriser le modèle pour éviter un tel scénario d'appel de fonction. Si vous trouvez cela difficile, contactez l'équipe d'assistance Google pour en savoir plus.

Reference

Options du convertisseur au format 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;
}

Présentation détaillée du traitement par lot

Le traitement par lot permet d'améliorer le débit et l'utilisation des TPU. Elle permet de traiter plusieurs requêtes en même temps. Pendant l'entraînement, le traitement par lot peut être effectué à l'aide de tf.data. Lors de l'inférence, cette opération consiste généralement à ajouter une opération dans le graphe qui regroupe les requêtes entrantes. L'opération attend qu'elle ait suffisamment de requêtes ou qu'un délai d'inactivité soit atteint avant de générer un lot volumineux à partir des requêtes individuelles. Consultez la section Définition des options de traitement par lot pour en savoir plus sur les différentes options de traitement par lot pouvant être ajustées, y compris les tailles de lot et les délais d'inactivité.

traitement par lot dans le graphique

Par défaut, le convertisseur insère l'opération de traitement par lot directement avant le calcul TPU. Elle encapsule la ou les fonctions TPU fournies par l'utilisateur et tout calcul TPU préexistant dans le modèle avec des opérations de traitement par lot. Il est possible d'ignorer ce comportement par défaut en indiquant au convertisseur quelles fonctions et/ou signatures doivent être regroupées.

L'exemple suivant montre comment ajouter le traitement par lot par défaut.

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
}

Traitement par lot des signatures

Le traitement par lot de la signature regroupe l'ensemble du modèle en commençant par les entrées de la signature et en allant jusqu'à ses sorties. Contrairement au comportement de traitement par lot par défaut du convertisseur, le traitement par lot des signatures regroupe à la fois les calculs TPU et les calculs du processeur. Cela permet un gain de performances de 10 à 20 % lors de l'inférence sur certains modèles.

Comme tout traitement par lot, le traitement par lot Signature présente des exigences strictes concernant les formes. Pour vous assurer que ces exigences sont remplies, les entrées de signature doivent avoir des formes présentant au moins deux dimensions. La première dimension est la taille de lot et doit avoir une taille de -1. Par exemple, (-1, 4), (-1) ou (-1, 128, 4, 10) sont toutes des formes d'entrée valides. Si cela n'est pas possible, envisagez d'utiliser le comportement de traitement par lot par défaut ou le traitement par lot des fonctions.

Pour utiliser le traitement par lot des signatures, indiquez le ou les noms de signature sous forme de signature_name à l'aide de 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"
  }
}

Traitement par lot des fonctions

Le traitement par lot des fonctions peut être utilisé pour indiquer au convertisseur quelles fonctions doivent être mises en lot. Par défaut, le convertisseur regroupe toutes les fonctions du TPU. Le traitement par lot des fonctions remplace ce comportement par défaut.

Le traitement par lot des fonctions peut être utilisé pour traiter les processeurs par lot. De nombreux modèles constatent une amélioration des performances lorsque leur calcul du processeur est effectué par lot. Le meilleur moyen de regrouper les calculs de processeur par lot consiste à utiliser le traitement par lot des signatures. Toutefois, cette méthode peut ne pas fonctionner pour certains modèles. Dans ce cas, le traitement par lot des fonctions peut être utilisé pour traiter une partie du calcul du processeur en plus du calcul TPU. Notez que l'opération de traitement par lot ne peut pas être exécutée sur le TPU. Par conséquent, toute fonction de traitement par lot fournie doit être appelée sur le processeur.

Le traitement par lot des fonctions peut également être utilisé pour répondre aux exigences strictes concernant les formes imposées par l'opération de traitement par lot. Dans les cas où les fonctions TPU ne répondent pas aux exigences de forme de l'opération de traitement par lot, le traitement par lot des fonctions peut être utilisé pour indiquer au convertisseur de regrouper différentes fonctions.

Pour l'utiliser, générez un function_alias pour la fonction qui doit être mise en lot. Pour ce faire, recherchez ou créez dans votre modèle une fonction qui encapsule tout ce que vous souhaitez regrouper. Assurez-vous que cette fonction répond aux exigences de forme strictes imposées par l'opération de traitement par lot. Ajoutez @tf.function si elle n'en a pas déjà une. Il est important de fournir le input_signature au @tf.function. La 0e dimension doit être None, car il s'agit de la dimension du lot et ne peut donc pas avoir une taille fixe. Par exemple, [None, 4], [None] ou [None, 128, 4, 10] sont toutes des formes d'entrée valides. Lors de l'enregistrement du modèle, fournissez des SaveOptions comme ceux indiqués ci-dessous pour attribuer à model.batch_func un alias "batch_func". Vous pouvez ensuite transmettre cet alias de fonction au convertisseur.

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)

Transmettez ensuite le ou les function_alias à l'aide de 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"
  }
}

Définition des options de traitement par lot

  • num_batch_threads: (entier) nombre de threads de planification pour le traitement des lots de travail. Détermine le nombre de lots traités en parallèle. Ce nombre correspond à peu près au nombre de cœurs de TPU disponibles.
  • max_batch_size: (entier) taille de lot maximale autorisée. Peut être supérieur à allowed_batch_sizes pour utiliser un fractionnement par lot volumineux.
  • batch_timeout_micros: (entier) nombre maximal de microsecondes à attendre avant de générer un lot incomplet.
  • allowed_batch_sizes: (liste d'entiers) Si la liste n'est pas vide, les lots sont complétés jusqu'à la taille la plus proche dans la liste. La liste doit être croissante de manière monotone, et l'élément final doit être inférieur ou égal à max_batch_size.
  • max_enqueued_batches: (entier) nombre maximal de lots mis en file d'attente pour traitement avant l'échec rapide des requêtes.

Mettre à jour les options de traitement par lot existantes

Vous pouvez ajouter ou mettre à jour des options de traitement par lot en exécutant l'image Docker en spécifiant "batch_options" et en définissant disable_default_optimizations sur "true" à l'aide de l'option --converter_options_string. Les options de traitement par lot sont appliquées à chaque fonction TPU ou opération de traitement par lot préexistante.

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

Exigences concernant les formes de traitement par lot

Les lots sont créés en concaténant les Tensors d'entrée des requêtes selon leur dimension de lot (0). Les Tensors de sortie sont divisés selon leur dimension 0. Pour effectuer ces opérations, l'opération de traitement par lot présente des exigences strictes en termes de forme pour ses entrées et ses sorties.

Tutoriel

Pour comprendre ces exigences, il est utile de comprendre d'abord comment le traitement par lot est effectué. L'exemple ci-dessous regroupe une opération tf.matmul simple.

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

La première requête d'inférence produit les entrées A et B avec, respectivement, les formes (1, 3, 2) et (1, 2, 4). La deuxième requête d'inférence produit les entrées A et B avec les formes (2, 3, 2) et (2, 2, 4).

requête d'inférence 1

Le délai avant expiration du traitement par lot est atteint. Comme le modèle accepte une taille de lot de 3, les requêtes d'inférence n° 1 et n° 2 sont regroupées sans remplissage. Les Tensors par lot sont formés en concaténant les requêtes n° 1 et n° 2 avec la dimension de lot (0). Comme la forme A de n° 1 a une forme de (1, 3, 2) et que celle de n° 2 a une forme de (2, 3, 2), lorsqu'elles sont concaténées le long de la dimension de lot (0), la forme résultante est (3, 3, 2).

requête par lot

Le tf.matmul est exécuté et génère une sortie au format (3, 3, 4).

requête matmul par lot

La sortie de tf.matmul est regroupée et doit donc être divisée en requêtes distinctes. Pour ce faire, l'opération de traitement par lot effectue une division en fonction de la dimension de lot (0) de chaque Tensor de sortie. Il décide comment diviser la 0e dimension en fonction de la forme des entrées d'origine. Étant donné que les formes de la requête n° 1 ont une dimension 0 de 1, la sortie a une dimension 0 de 1 pour une forme de (1, 3, 4). Étant donné que les formes de la requête n° 2 ont une dimension 0 de 2, sa sortie présente une 0e dimension de 2 pour une forme de (2, 3, 4).

résultats des requêtes d'inférence

Exigences concernant les formes

Pour effectuer la concaténation des entrées et la division de la sortie décrites ci-dessus, l'opération de traitement par lot présente les exigences de forme suivantes:

  1. Les entrées du traitement par lot ne peuvent pas être scalaires. Pour être concaténés selon la dimension 0, les Tensors doivent avoir au moins deux dimensions.

    Dans le tutoriel ci-dessus. Ni A, ni B ne sont des scalaires.

    Si vous ne respectez pas cette exigence, une erreur telle que "Batching input tensors must have at least one dimension" sera renvoyée. Une solution simple à cette erreur consiste à transformer le scalaire en vecteur.

  2. Dans différentes requêtes d'inférence (par exemple, différents appels d'exécution de session), les Tensors d'entrée portant le même nom ont la même taille pour chaque dimension, à l'exception de la dimension 0. Cela permet de concaténer correctement les entrées le long de leur 0e dimension.

    Dans le tutoriel ci-dessus, la forme A de la requête n° 1 est (1, 3, 2). Cela signifie que toute requête ultérieure doit produire une forme avec le modèle (X, 3, 2). La requête 2 répond à cette exigence avec (2, 3, 2). De même, la forme B de la requête n° 1 est (1, 2, 4). Par conséquent, toutes les requêtes futures doivent produire une forme présentant le modèle (X, 2, 4).

    Si vous ne respectez pas cette exigence, une erreur telle que "Dimensions of inputs should match" sera renvoyée.

  3. Pour une requête d'inférence donnée, toutes les entrées doivent avoir la même taille de 0e dimension. Si différents Tensors d'entrée de l'opération de traitement par lot ont des dimensions 0 différentes, l'opération de traitement par lot ne sait pas comment diviser les Tensors de sortie.

    Dans le tutoriel ci-dessus, les Tensors de la requête n° 1 ont tous une taille de dimension 0 de 1. Cela permet à l'opération de traitement par lot de savoir que sa sortie doit avoir une taille de 0e dimension de 1. De même, les Tensors de la requête n° 2 ont une taille de dimension 0 de 2. Sa sortie aura donc une taille de dimension 0 de 2. Lorsque l'opération de traitement par lot divise la forme finale de (3, 3, 4), elle génère (1, 3, 4) pour la requête n° 1 et (2, 3, 4) pour la requête n° 2.

    Si vous ne remplissez pas cette exigence, des erreurs telles que celles indiquées ci-dessous seront générées: Batching input tensors supplied in a given op invocation must have equal 0th-dimension size.

  4. La taille 0 de la forme de chaque Tensor de sortie doit être la somme de la taille 0 de tous les Tensors d'entrée (plus toute marge intérieure introduite par l'opération de traitement par lot pour atteindre la valeur allowed_batch_size la plus élevée). Cela permet à l'opération de traitement par lot de fractionner les Tensors de sortie selon leur 0e dimension en fonction de la dimension 0 des Tensors d'entrée.

    Dans le tutoriel ci-dessus, les Tensors d'entrée ont une dimension 0 de 1 pour la requête n° 1 et de 2 pour la requête n° 2. Par conséquent, chaque Tensor de sortie doit avoir une dimension 0 de 3, car 1+2=3. Le Tensor de sortie (3, 3, 4) répond à cette exigence. Si la taille de lot 3 n'était pas valide, mais que la valeur 4 l'était, l'opération de traitement par lot aurait dû compléter la 0e dimension des entrées de 3 à 4. Dans ce cas, chaque Tensor de sortie doit avoir une taille de dimension 0 de 4.

    Si vous ne remplissez pas cette exigence, une erreur telle que celle-ci s'affichera: Batched output tensor's 0th dimension does not equal the sum of the 0th dimension sizes of the input tensors.

Résoudre les erreurs liées aux exigences de forme

Pour répondre à ces exigences, envisagez de fournir une fonction ou une signature différente pour le lot. Il peut également être nécessaire de modifier des fonctions existantes pour répondre à ces exigences.

Si une fonction est mise par lot, assurez-vous que les formes de sa @tf.function de la signature d'entrée (input_signature) ont toutes None dans la dimension 0 (également appelée dimension de lot). Si une signature est traitée par lot, assurez-vous que toutes ses entrées présentent la valeur -1 dans la dimension 0.

L'opération BatchFunction n'est pas compatible avec SparseTensors en tant qu'entrées ni comme sorties. En interne, chaque Tensor creux est représenté par trois Tensors distincts pouvant avoir des tailles de 0 dimension différentes.