Introducción al convertidor de inferencia de Cloud TPU v5e

Introducción

El convertidor de inferencia de Cloud TPU prepara y optimiza un modelo de TensorFlow 2 (TF2) para la inferencia TPU. El conversor se ejecuta en una shell local o de VM de TPU. Se recomienda el shell de la VM de TPU porque viene preinstalado con el comando de línea de comandos necesarias para el conversor. Toma una solicitud SavedModel y realiza los siguientes pasos:

1. Conversión de TPU: Agrega TPUPartitionedCall y otras operaciones de TPU al para que se pueda publicar en la TPU. Por defecto, un modelo exportado para la inferencia no tiene esas operaciones y no se puede entregar en la TPU, incluso si se entrenó en la TPU.
2. Agrupación en lotes: Agrega operaciones por lotes al modelo para habilitar el procesamiento por lotes en el grafo. para mejorar la capacidad de procesamiento.
3. Conversión de BFloat16: Convierte el formato de datos del modelo a partir del float32 a bfloat16 para mejorar rendimiento computacional y menor memoria de alto ancho de banda (HBM) uso en la TPU.
4. Optimización de forma de IO: Optimiza las formas de tensor para datos que se transfieren entre la CPU y la TPU para mejorar el uso del ancho de banda.

Cuando exportan un modelo, los usuarios crean alias de funciones para cualquier función que que le gustaría ejecutar en la TPU. Pasan estas funciones al convertidor y al El convertidor los coloca en la TPU y los optimiza.

El convertidor de inferencia de Cloud TPU está disponible como una imagen de Docker que se puede que se ejecutan en cualquier entorno con Docker instalado.

Tiempo estimado para completar los pasos anteriores: 20 min a 30 min

Requisitos previos

1. El modelo debe ser de TF2 y exportarse en el SavedModel de un conjunto de datos tengan un formato común.
2. El modelo debe tener un alias de función para la función de TPU. Consulta la ejemplo de código sobre cómo hacerlo. En los siguientes ejemplos, se usa tpu_func como la TPU alias de la función.
3. Asegúrate de que la CPU de tu máquina sea compatible con Advanced Vector eXtensions (AVX). instrucciones, que la biblioteca de TensorFlow (la dependencia de la Cloud TPU) el convertidor de inferencia) se compila para usar instrucciones de AVX. La mayoría de las CPUs tienen compatibilidad con AVX.
   1. Puedes ejecutar lscpu | grep avx para verificar si el complemento AVX de instrucciones.

Antes de comenzar

Antes de comenzar la configuración, haz lo siguiente:

• Crea un proyecto nuevo: En la consola de Google Cloud, en la página del selector de proyectos, selecciona o crea una Cloud.

• Configura una VM de TPU: Crea una nueva VM de TPU con la consola de Google Cloud o gcloud, o usa un VM de TPU existente para ejecutar inferencias con la VM modelo en la VM de TPU.

  • Asegúrate de que la imagen de VM de TPU esté basada en TensorFlow. Por ejemplo, --version=tpu-vm-tf-2.11.0.
  • El modelo convertido se cargará y entregará en esta VM de TPU.

• Asegúrate de tener las herramientas de línea de comandos que necesitas para usar Cloud TPU. Convertidor de inferencia. Puedes instalar el SDK de Google Cloud y Docker de forma local o usar una VM de TPU que tenga este software instalado de forma predeterminada. Usa estas herramientas para interactuar con la imagen del convertidor.

  Conéctate a la instancia con SSH a través del siguiente comando:

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

Configuración del entorno

Configura el entorno desde la shell de la VM de TPU o desde tu shell local.

Shell de VM de TPU

• En la shell de tu VM de TPU, ejecuta los siguientes comandos para Permitir el uso de Docker no raíz:

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

• Inicializa tus auxiliares de credenciales de Docker:

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

Shell local

En tu shell local, configura el entorno mediante los siguientes pasos:

• Instala el SDK de Cloud. que incluye la herramienta de línea de comandos de gcloud.

• Instala Docker:

• Permite el uso de Docker no raíz:

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

• Accede a tu entorno:

  gcloud auth login
  

• Inicializa tus auxiliares de credenciales de Docker:

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

• Extrae la imagen de Docker del convertidor de inferencia:

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

Imagen del convertidor

La imagen sirve para realizar conversiones de modelos únicas. Establece las rutas del modelo y ajustar el opciones de conversión según tus necesidades. El Ejemplos de uso se proporcionan varios casos de uso comunes.

Nota: En el siguiente código, se puede usar un bucket de Cloud Storage. para input_model_dir y output_model_dir sin especificar mount en el 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
    }
'

Inferencia con el modelo convertido en la VM de 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))

Ejemplos de uso

Agrega un alias de función para la función de TPU

1. Encuentra o crea una función en tu modelo que encaje todo lo que deseas. para que se ejecute en la TPU. Si @tf.function no existe, agrégalo.
2. Cuando guardes el modelo, proporciona SaveOptions como el de abajo para dar model.tpu_func un alias func_on_tpu.
3. Puedes pasar este alias de función al conversor.

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)

Convierte un modelo con varias funciones de TPU

Puedes colocar varias funciones en la TPU. Simplemente crea varias funciones y pasarlos en converter_options_string al conversor.

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

Cuantización

Importante: Esta función es experimental y puede estar sujeta a cambios. sin previo aviso. Los usuarios deben evaluar cuidadosamente el rendimiento de la cuantizado antes de usarlo en producción. Compatibilidad con TPU v5e para se proporciona la cuantización en TensorFlow.

La cuantización es una técnica que reduce la precisión de los números usada para representar los parámetros de un modelo. Esto da como resultado un modelo más pequeño y un procesamiento más rápido. Un modelo cuantificado proporciona ganancias en capacidad de procesamiento de inferencia, menor uso de memoria y tamaño de almacenamiento a costa de pequeñas disminuciones en la exactitud.

La nueva función de cuantización posterior al entrenamiento en TensorFlow que se orienta TPU se desarrolló a partir de una función existente similar en TensorFlow Lite que se usa para orientarse a dispositivos móviles y perimetrales. Para obtener más información cuantización en general, puedes observar la Documento de TensorFlow Lite.

Conceptos de cuantización

En esta sección, se definen conceptos relacionados en específico con la cuantización. con el convertidor de inferencia.

Conceptos relacionados con otras opciones de configuración de TPU (por ejemplo, unidades, hosts, chips y TensorCores) se describen en el Arquitectura de sistemas de TPU.

• Cuantización posterior al entrenamiento (PTQ): Es una técnica que reduce el tamaño y la complejidad computacional de un modelo de red neuronal lo que afecta significativamente su precisión. PTQ funciona convirtiendo la ponderaciones de punto flotante y activaciones de un modelo entrenado para números enteros de menor precisión, como los de 8 o 16 bits. Esto puede provocar una reducción significativa en el tamaño del modelo y la latencia de la inferencia, mientras que lo que incurrirá en una pequeña pérdida de exactitud.

• Calibración: El paso de calibración para la cuantización es el proceso. de recopilar estadísticas sobre el rango de valores que las ponderaciones y que realizan las activaciones de un modelo de red neuronal. Esta información se usa para determinar los parámetros de cuantización para el modelo, que son los valores que se usará para convertir los pesos de punto flotante y las activaciones a en números enteros.

• Conjunto de datos representativo: Es un conjunto de datos representativo para la cuantización. es un conjunto de datos pequeño que representa los datos de entrada reales para el modelo. Se usa durante el paso de calibración de la cuantización para recopilar estadísticas sobre el rango de valores que los pesos y las activaciones que tomará el modelo. El conjunto de datos representativo debe cumplir con los siguientes propiedades:

  • Debe representar adecuadamente las entradas reales al modelo durante la inferencia. Esto significa que debe abarcar el rango de valores que que el modelo probablemente vea en el mundo real.
  • Debería fluir colectivamente a través de cada rama de condicionales (como tf.cond), si existe alguno. Esto es importante porque el proceso de cuantización debe poder manejar todas las entradas posibles al modelo, incluso si no están representados de forma explícita en el representativo.
  • Debe ser lo suficientemente grande como para recopilar suficientes estadísticas y reducir . Como regla general, se recomienda usar más de 200 muestras representativas.

  El conjunto de datos representativo puede ser un subconjunto del conjunto de datos de entrenamiento, o puede ser un conjunto de datos separado que está diseñado específicamente para representativos de las entradas reales del modelo. La elección de qué conjunto de datos usar depende de la aplicación específica.

• Cuantización de rango estático (SRQ): SRQ determina el rango de valores. para los pesos y las activaciones de un modelo de red neuronal una vez, durante de calibración. Esto significa que se usa el mismo rango de valores para todos entradas al modelo. Puede ser menos preciso que el rango dinámico en especial para modelos con un amplio rango de valores de entrada. Sin embargo, la cuantización del rango estático requiere menos procesamiento en la ejecución. tiempo que la cuantización del rango dinámico.

• Cuantización de rango dinámico (DRQ): DRQ determina el rango valores para los pesos y las activaciones de un modelo de red neuronal para cada entrada. Esto permite que el modelo se adapte al rango de valores de los datos de entrada, lo que puede mejorar su exactitud. Sin embargo, el rango dinámico la cuantización requiere más procesamiento en el tiempo de ejecución que el rango estático mediante la cuantización.

  Función	Cuantización de rango estático	Cuantización de rango dinámico
  Rango de valores	Se determina una vez durante la calibración	Determinado para cada entrada
  Precisión	Puede ser menos preciso, en especial para modelos con una amplia gama de valores de entrada.	Puede ser más preciso, en especial para modelos con una amplia gama de valores de entrada.
  Complejidad	Es más simple	Es más complejo
  Procesamiento en tiempo de ejecución	Menos procesamiento	Más procesamiento

• Cuantización solo de peso: La cuantización solo de peso es un tipo de que solo cuantiza los pesos de un modelo de red neuronal, y dejar las activaciones en punto flotante. Esta puede ser una buena opción para los modelos que son sensibles a la exactitud, ya que puede ayudar para preservar la exactitud del modelo.

Cómo usar la cuantización

La cuantización se puede aplicar mediante la configuración y los QuantizationOptions a las opciones del conversor. Las opciones destacadas son las siguientes:

• etiquetas: Una colección de etiquetas que identifican el MetaGraphDef dentro SavedModel para cuantizar. No es necesario especificar si solo tienes un MetaGraphDef.
• signature_keys: Secuencia de claves que identifican SignatureDef que contiene entradas y salidas. Si no se especifica, se usa ["serving_default"].
• quantization_method: Es el método de cuantización que se aplicará. Si no es así especificado, se aplicará una cuantización de STATIC_RANGE.
• op_set: Debe mantenerse como XLA. Actualmente es la opción predeterminada, no es necesario especificarla.
• representa_conjuntos de datos: Especifica el conjunto de datos que se usa para calibrar los parámetros de cuantización.

Compila el conjunto de datos representativo

Un conjunto de datos representativo es, en esencia, un iterable de muestras. Una muestra es un mapa de {input_key: input_value}. Por ejemplo:

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

Los conjuntos de datos representativos deben guardarse como TFRecord con la clase TfRecordRepresentativeDatasetSaver en el paquete tf-nightly de pip. Por ejemplo:

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

Ejemplos

En el siguiente ejemplo, se cuantiza el modelo con la clave de firma de serving_default y el alias de función de tpu_func:

Nota: En el siguiente código, configura CONVERTER_IMAGE como us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:nightly para usar la compilación nocturna.

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

Agregar agrupación en lotes

El convertidor se puede usar para agregar lotes a un modelo. Para ver una descripción del opciones de lote que se pueden ajustar, consulta Definición de las opciones de lotes.

De forma predeterminada, el convertidor agrupará por lotes cualquier función de TPU en el modelo. También puede lote proporcionado por el usuario firmas y funciones lo que puede mejorar aún más el rendimiento. Cualquier función de TPU, función proporcionada por el usuario o firma que se agrupa en lotes, debe cumplir con las requisitos de forma estrictos.

El objeto Converter también puede actualización opciones existentes de lotes. El siguiente es un ejemplo de cómo agregar agrupación en lotes a un modelo. Para ver más información sobre el procesamiento por lotes, consulta Análisis detallado por lotes.

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
}

Inhabilitar las optimizaciones de forma de E/S y bfloat16

Las optimizaciones de forma de IO y BFloat16 están habilitadas de forma predeterminada. Si no funcionan con tu modelo, pueden inhabilitarse.

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

Informe de conversiones

Puedes encontrar este informe de conversiones en el registro después de ejecutar la inferencia Convertidor. Aquí tienen un ejemplo.

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

Este informe estima el costo de procesamiento del modelo de salida en CPU y TPU y desglosa aún más el costo de TPU para cada función, lo que debería reflejar la selección de las funciones de TPU en las opciones del conversor.

Si quieres usar mejor la TPU, es posible que quieras experimentar con el modelo. estructurar y ajustar las opciones del conversor.

Nota: En el informe de conversiones actual, todas las operaciones de las funciones de TPU se se cuenta como costo de TPU.

Preguntas frecuentes

¿Qué funciones debo colocar en la TPU?

Es mejor colocar la mayor cantidad posible de tu modelo en la TPU, ya que el y la gran mayoría de las operaciones se ejecutan más rápido en la TPU.

Si tu modelo no contiene op, cadenas o dispersas incompatibles con TPU tensores, colocar todo el modelo en la TPU suele ser la mejor estrategia. Y puede buscar o crear una función que encapsule todo el modelo crear un alias de función y pasarlo al Converter.

Si el modelo contiene partes que no funcionan en la TPU (por ejemplo, incompatibles con TPU) ops, cadenas o tensores dispersos), la elección de funciones de TPU depende de dónde es la parte incompatible.

• Si es al principio o al final del modelo, puedes refactorizar para mantenerlo en la CPU. Algunos ejemplos son el procesamiento previo y posterior de cadenas etapas. Para obtener más información sobre cómo mover código a la CPU, consulta lo siguiente: “¿Cómo transfiero una parte del modelo a la CPU?” Muestra una forma típica de refactorizar el modelo.
• Si está en el medio del modelo, es mejor dividir el modelo en tres partes y contener todas las operaciones incompatibles con TPU en la parte central y ejecutarlo en la CPU.
• Si se trata de un tensor disperso, considera llamar a tf.sparse.to_dense en el CPU y pasa el tensor denso resultante a la parte de la TPU del modelo.

Otro factor a tener en cuenta es el uso de HBM. Las tablas de incorporación pueden usar muchos o la HBM. Si sobrepasan las limitaciones del hardware de la TPU, se deben en la CPU, junto con las operaciones de búsqueda.

Siempre que sea posible, solo debe existir una función de TPU en una firma. Si la estructura de tu modelo requiere llamar a múltiples funciones de TPU por solicitud de inferencia entrante, debes tener en cuenta la latencia adicional de enviar de tensores entre CPU y TPU.

Una buena manera de evaluar la selección de funciones de TPU es comprobar el Informe de conversiones: Muestra el porcentaje de procesamiento que se aplicó a la TPU y una desglose del costo de cada función de TPU.

¿Cómo paso una parte del modelo a la CPU?

Si tu modelo contiene partes que no se pueden entregar en la TPU, debes refactorizar el modelo para moverlos a la CPU. Este es un ejemplo de juguete. El modelo es en un modelo de lenguaje con una etapa de procesamiento previo. El código de las definiciones de capas y se omiten para mayor simplicidad.

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)

Este modelo no se puede entregar directamente en la TPU por dos motivos. En primer lugar, parámetro es una cadena. En segundo lugar, la función preprocess puede contener muchas cadenas ops. Ninguno de ellos es compatible con TPU.

Para refactorizar este modelo, puedes crear otra función llamada tpu_func para alojar la bert_layer de procesamiento intensivo. Luego, crea un alias de función para tpu_func y pásalo al convertidor. De esta manera, todo lo que está tpu_func se ejecutará en la TPU y todo lo que quede en model_func se ejecutará en la CPU.

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

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

¿Qué debo hacer si el modelo tiene operaciones, cadenas o tensores dispersos no compatibles con TPU?

La mayoría de las ops de TensorFlow estándar son compatibles con la TPU, pero algunas incluidos los tensores dispersos y las cadenas. El usuario que generó una conversión no comprueba si hay operaciones incompatibles con TPU. Entonces, un modelo que contiene esas operaciones puede pasar el conversión. Sin embargo, cuando se ejecute para inferencia, se producirán errores como los siguientes.

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

Si tu modelo tiene operaciones incompatibles con la TPU, deben colocarse fuera de la TPU. . Además, la cadena es un formato de datos no admitido en la TPU. De esta manera, las variables de tipo cadena no se deben colocar en la función de TPU. Además, parámetros y los valores de retorno de la función de TPU no deben ser de tipo cadena como en la nube. Del mismo modo, evita colocar tensores dispersos en la función de TPU, incluso en sus parámetros y valores de retorno.

Por lo general, no es difícil refactorizar la parte incompatible del modelo moverla a la CPU. Este es un ejemplo.

¿Cómo admitir operaciones personalizadas en el modelo?

Si se usan operaciones personalizadas en tu modelo, es posible que el convertidor no las reconozca y no logra convertir el modelo. Esto se debe a que la biblioteca de op de la op personalizada que contiene la definición completa de la op, no está vinculada al convertidor.

Como el código del conversor aún no es de código abierto, que se creó con op personalizada.

¿Qué debo hacer si tengo un modelo de TensorFlow 1?

El convertidor no es compatible con los modelos de TensorFlow 1. Los modelos de TensorFlow 1 deben migrar a TensorFlow 2.

¿Debo habilitar el puente MLIR cuando ejecuto mi modelo?

La mayoría de los modelos convertidos pueden ejecutarse con el puente MLIR TF2XLA más reciente el puente TF2XLA original.

¿Cómo convierto un modelo que ya se exportó sin un alias de función?

Si se exportó un modelo sin un alias de función, la forma más fácil es exportar una y otra vez y Crea un alias de función. Si la reexportación no es una opción, es posible convertir el modelo proporcionando un concrete_function_name Sin embargo, identificar concrete_function_name sí requiere trabajo de detección.

Los alias de función son una asignación de una cadena definida por el usuario a una función concreta de la fuente de datos. Facilitan la referencia a una función específica en el modelo. El El convertidor acepta alias de funciones y nombres de funciones concretos sin procesar.

Los nombres de funciones concretas se pueden encontrar examinando saved_model.pb.

En el siguiente ejemplo, se muestra cómo poner una función concreta llamada __inference_serve_24 en la 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"
    }'

¿Cómo resuelvo un error de restricción de constante de tiempo de compilación?

Tanto para el entrenamiento como para la inferencia, XLA requiere que las entradas a ciertas operaciones tengan un una forma conocida en el momento de la compilación de la TPU. Esto significa que cuando XLA compila la TPU del programa, las entradas a estas operaciones deben tener un forma.

Existen dos maneras de resolver este problema.

• La mejor opción es actualizar las entradas de la op para que tengan un modelo para cuando XLA compila el programa de TPU. Esta compilación se lleva a cabo justo antes de que se ejecute la porción de la TPU del modelo. Esto significa que la forma deberían conocerse estáticamente cuando TpuFunction esté a punto de ejecutarse.
• Otra opción es modificar TpuFunction para que ya no incluya el op problemática.

¿Por qué recibo un error de forma por lotes?

La agrupación en lotes tiene requisitos de forma estrictos que permiten que las solicitudes entrantes se agrupen en lotes junto con su dimensión 0 (es decir, el dimensión de lotes). Estos requisitos de forma provienen del procesamiento por lotes de TensorFlow y no puede ser relajado.

Si no se cumplen estos requisitos, se producirán errores como los siguientes:

1. Los tensores de entrada de lotes deben tener al menos una dimensión.
2. Las dimensiones de las entradas deben coincidir.
3. Los tensores de entrada de lotes proporcionados en una invocación de op dada deben tener la misma tamaño de la 0a dimensión.
4. La 0.a dimensión del tensor de salida por lotes no es igual a la suma de la 0. de los tensores de entrada.

Para cumplir con estos requisitos, considera proporcionar una función o firma para agrupar en lotes. También puede ser necesario modificar las funciones existentes para cumplir con y los requisitos de cumplimiento.

Si una función se procesa en lotes, asegúrate de que las formas de input_signature de @tf.function no tienen ninguna en la dimensión 0, Si se agrega una firma se está agrupando, asegúrate de que todas sus entradas tengan -1 en la dimensión 0.

Para obtener una explicación completa sobre por qué ocurren estos errores y cómo resolverlos puedes consultarlos Batching Deep Dive:

Problemas conocidos

La función de TPU no puede llamar indirectamente a otra función de TPU.

Si bien el convertidor puede controlar la mayoría de las situaciones de llamadas a funciones en la límite de CPU-TPU, hay un caso límite poco frecuente en el que fallaría. Es cuando una TPU llama indirectamente a otra función de TPU.

Esto se debe a que el convertidor modifica el llamador directo de una función de TPU desde que llama a la función de TPU para llamar a un stub de llamadas de TPU. El stub de llamada contiene ops que solo pueden funcionar en la CPU. Cuando una función TPU llama a que finalmente llama al llamador directo, esas operaciones de CPU se podrían llevar en la TPU que se ejecutará, lo que generará errores de kernel faltantes. Ten en cuenta este caso es diferente de una función de TPU que llama directamente a otra función de TPU. En este caso, el convertidor no modifica ninguna de las funciones para llamar al código auxiliar de llamadas, por lo que puede funcionar.

En el convertidor, implementamos la detección de esta situación. Si ves Si se muestra el siguiente error, significa que tu modelo encontró este caso límite:

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 solución general es refactorizar el modelo para evitar que una función llame situación. Si te resulta difícil hacerlo, comunícate con el equipo de Atención al cliente de Google equipo para discutirlo más.

Reference

Opciones del convertidor en 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;
}

Análisis detallado de la agrupación en lotes

La agrupación en lotes se usa para mejorar la capacidad de procesamiento y el uso de TPU. Permite varias solicitudes para que se procesen al mismo tiempo. Durante el entrenamiento, el agrupamiento en lotes se puede hacer con tf.data. Durante la inferencia, por lo general, se hace agregando un op en el gráfico que agrupa en lotes las solicitudes entrantes. La op espera hasta que tenga suficiente de solicitudes o se alcanza un tiempo de espera antes de que genere un gran lote a partir del solicitudes individuales. Consulta Definición de las opciones de lotes para obtener más información sobre las diferentes opciones de lotes que se pueden ajustar incluidos los tamaños de los lotes y los tiempos de espera.

De forma predeterminada, el convertidor inserta la operación de lotes directamente antes de la TPU. el procesamiento. Une las funciones de TPU proporcionadas por el usuario y cualquier TPU preexistente de procesamiento en el modelo con operaciones por lotes. Esto se puede anular comportamiento predeterminado indicando al convertidor qué funciones y/o firmas se pueden agrupar en lotes.

En el siguiente ejemplo, se muestra cómo agregar el procesamiento por lotes predeterminado.

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
}

Agrupación en lotes de firmas

Nota: Esta función aún se encuentra en la etapa experimental. Después de las pruebas, el campo se quitará del wrapper experimental.

El procesamiento por lotes de firmas por lotes a todo el modelo a partir de las entradas de la firma y vaya a los resultados de la firma. A diferencia del agrupador en lotes predeterminado del convertidor comportamiento, el procesamiento por lotes de firmas en lotes tanto el cálculo de TPU como el de la CPU el procesamiento. Esto da entre un 10% y 20% de mejora de rendimiento durante la inferencia en algunos e implementar modelos automáticamente.

Como sucede con todos los lotes, requisitos de forma estrictos. Para garantizar que se cumplan estos requisitos de forma, las entradas de firma deben tener formas que tienen al menos dos dimensiones. La primera dimensión es el tamaño del lote y debe tener un tamaño de -1. Por ejemplo, (-1, 4), (-1) o (-1, 128, 4, 10) son todas formas de entrada válidas. Si no es posible, considera usar el comportamiento por lotes predeterminado o agrupación en lotes de funciones.

Para usar el agrupamiento en lotes de firmas, proporciona los nombres de las firmas como signature_name con 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"
  }
}

Agrupación en lotes de funciones

Nota: Esta función aún se encuentra en la etapa experimental. Después de las pruebas, el campo se quitará del wrapper experimental.

El procesamiento por lotes de funciones se puede usar para indicarle al convertidor qué funciones se deben por lotes. De forma predeterminada, el convertidor agrupará todas las funciones de TPU. Función el procesamiento por lotes anula este comportamiento predeterminado.

El procesamiento por lotes de funciones se puede usar para el procesamiento por lotes de la CPU. Muchos modelos ven un de mejora del rendimiento cuando el procesamiento de la CPU se procesa en lotes. La mejor manera de el procesamiento por lotes de CPU usa lotes de firmas; sin embargo, es posible que no funcione para algunos e implementar modelos automáticamente. En esos casos, la función puede usarse para agrupar en lotes parte de la CPU además del cálculo de TPU. Ten en cuenta que en la op ejecutarse en la TPU, por lo que cualquier función de lote que se proporcione debe llamarse en el CPU.

El procesamiento por lotes de funciones también se puede usar para cumplir con requisitos de forma estrictos que impone la operación de lotes. Si las funciones de la TPU no cumplen con los requisitos de forma de las op en lotes, se puede usar el procesamiento por lotes Convertidor en lotes de diferentes funciones.

Para usar esto, genera un function_alias para la función que debería por lotes. Puedes hacerlo buscando o creando una función en tu modelo que incluye todo lo que quieres agrupar en lotes. Asegúrate de que esta función cumpla con los requisitos de forma estrictos que impone la operación de lotes. Agrega @tf.function si aún no tiene uno. Es importante proporcionar input_signature a @tf.function. El 0 dimensión debe ser None porque es la dimensión del lote, por lo que no puede ser un tamaño fijo. Por ejemplo, [None, 4], [None] o [None, 128, 4, 10] son todos formas de entrada válidas. Cuando guardes el modelo, proporciona SaveOptions como los que se muestran. a continuación para asignarle a model.batch_func el alias “batch_func”. Luego, puedes pasar esto alias de la función al conversor.

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

  ...

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

A continuación, pasa las function_alias con las 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"
  }
}

Definición de las opciones de agrupación en lotes

• num_batch_threads: (número entero) Cantidad de subprocesos de programación para lotes de trabajo en proceso. Determina la cantidad de lotes procesados en en paralelo. Esto debería estar más o menos en línea con la cantidad de núcleos de TPU disponibles.
• max_batch_size: (número entero) Tamaño del lote máximo permitido. Puede ser más grande que allowed_batch_sizes para usar la división por lotes grande.
• batch_timeout_micros: (número entero) Cantidad máxima de microsegundos que se esperará antes de generar un lote incompleto.
• allowed_batch_sizes: (lista de números enteros) Si la lista no está vacía, rellenará los lotes hasta el tamaño más cercano de la lista. La lista debe tener monótonamente creciente y el elemento final debe ser menor o igual que max_batch_size
• max_enqueued_batches: (número entero) Cantidad máxima de lotes en cola para antes de que las solicitudes fallen rápido.

Actualiza las opciones de lotes existentes

Puedes agregar o actualizar las opciones de lote si ejecutas la imagen de Docker y especifica lote_options y configura disable_default_optimizations como verdadero mediante el --converter_options_string. Las opciones de lote se aplicarán a cada una función de TPU o una op de procesamiento por lotes preexistente.

Nota: No puedes actualizar las opciones de lotes si configuras tpu_functions juntos. con --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

Requisitos de forma de agrupación en lotes

Los lotes se crean mediante la concatenación de tensores de entrada en las solicitudes a lo largo de su lote (0a). Los tensores de salida se dividen según su dimensión 0. En orden para realizar estas operaciones, la operación de lotes tiene requisitos estrictos de forma para sus entradas y salidas.

Explicación

Para entender estos requisitos, es útil primero entender cómo por lotes. En el siguiente ejemplo, agruparemos en lotes una tabla tf.matmul operación

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

La primera solicitud de inferencia produce las entradas A y B con las formas (1, 3, 2) y (1, 2, 4), respectivamente. La segunda solicitud de inferencia produce la las entradas A y B con las formas (2, 3, 2) y (2, 2, 4).

Se alcanzó el tiempo de espera de lotes. El modelo admite un tamaño de lote de 3, por lo que las solicitudes de inferencia n.o 1 y n.o 2 se agrupan en lotes sin ningún relleno. El Los tensores en lotes se forman mediante la concatenación de las solicitudes #1 y #2 a lo largo del lote (0) dimensión. Como la A del #1 tiene una forma de (1, 3, 2), y la A del #2 tiene una forma de (2, 3, 2), cuando se concatenan a lo largo de la dimensión del lote (0), el valor la forma resultante es (3, 3, 2).

Se ejecuta tf.matmul y produce un resultado con la forma (3, 3, 4).

El resultado de tf.matmul se agrupa en lotes, por lo que debe volver a dividirse en solicitudes separadas. La op de procesamiento por lotes lo hace dividiendo a lo largo del lote (0) de cada tensor de salida. Decide cómo dividir la dimensión 0 en función según la forma de las entradas originales. Dado que las formas de la solicitud n° 1 tienen un 0 de salida tiene una dimensión 0 de 1 para una forma de (1, 3, 4). Dado que las formas de la solicitud n° 2 tienen una 0a dimensión de 2, su salida tiene un 0. dimensión de 2 para una forma de (2, 3, 4).

Requisitos de forma

Para realizar la concatenación de entrada y la división de salida que se describen arriba, la op de lotes tiene los siguientes requisitos de forma:

1. Las entradas para la agrupación en lotes no pueden ser escalares. Para concatenar el 0.a dimensión, los tensores deben tener al menos dos dimensiones.

   En la explicación anterior, Ni A ni B son escalares.

   Si no se cumple con este requisito, se producirá un error como el siguiente: Batching input tensors must have at least one dimension. Una solución sencilla para este error es convertir el escalar en un vector.

2. En diferentes solicitudes de inferencia (por ejemplo, diferentes sesiones de ejecución invocaciones), los tensores de entrada con el mismo nombre tienen el mismo tamaño para cada más, excepto la 0. Esto permite que las entradas sean limpias se concatenan a lo largo de su dimensión 0.

   En la explicación anterior, la solicitud A de la solicitud núm. 1 tiene una forma de (1, 3, 2). Esto significa que cualquier solicitud futura debe producir una forma con el patrón (X, 3, 2) La solicitud núm. 2 cumple con este requisito con (2, 3, 2). Del mismo modo, la solicitud B de la núm. 1 tiene una forma de (1, 2, 4), por lo que todas las solicitudes futuras deben producirá una forma con el patrón (X, 2, 4).

   Si no se cumple con este requisito, se producirá un error como el siguiente: Dimensions of inputs should match.

3. Para una solicitud de inferencia determinada, todas las entradas deben tener el mismo valor 0. tamaño de la dimensión. Si diferentes tensores de entrada a la op dimensiones 0 diferentes, la operación de lotes no sabe cómo dividir la tensores de salida.

   En la explicación anterior, todos los tensores de la solicitud núm. 1 tienen una dimensión 0. tamaño de 1. Esto le permite a la op de procesamiento por lotes saber que su salida debe tener un valor 0. el tamaño de dimensión de 1. De manera similar, los tensores de la solicitud núm. 2 tienen una dimensión 0. de tamaño 2, por lo que su salida tendrá un tamaño de dimensión 0 de 2. Cuando La operación por lotes divide la forma final de (3, 3, 4) y produce (1, 3, 4) para la solicitud n.o 1 y (2, 3, 4) para la solicitud n.o 2.

   Si no se cumple con este requisito, se producirán errores como los siguientes: Batching input tensors supplied in a given op invocation must have equal 0th-dimension size.

4. El tamaño de dimensión 0 de la forma de cada tensor de salida debe ser la suma de todos los tensores de entrada tamaño de la 0a dimensión (más cualquier relleno ingresado por la operación de lotes para cumplir con la siguiente allowed_batch_size más grande). Esto permite la operación de procesamiento por lotes para dividir los tensores de salida en su dimensión 0. en la dimensión 0 de los tensores de entrada.

   En la explicación anterior, los tensores de entrada tienen una dimensión 0 de 1. de las solicitudes n° 1 y 2 de la solicitud n° 2. Por lo tanto, cada tensor de salida debe tienen una 0a dimensión de 3 porque 1+2=3. El tensor de salida (3, 3, 4) cumpla con este requisito. Si 3 no hubiera sido un tamaño de lote válido, pero 4 lo era, de procesamiento por lotes habría tenido que rellenar la 0a dimensión de las entradas de 3 a 4. En este caso, cada tensor de salida debería tener un tamaño de 0a dimensión de 4.

   Si no se cumple con este requisito, se mostrará un error como el siguiente: Batched output tensor's 0th dimension does not equal the sum of the 0th dimension sizes of the input tensors.

Resuelve errores de requisitos de forma

Para cumplir con estos requisitos, considera proporcionar una función o firma para agrupar en lotes. También puede ser necesario modificar las funciones existentes para cumplir con y los requisitos de cumplimiento.

Si un función se procesa en lotes, asegúrate de que las formas de input_signature de @tf.function Deben tener None en la dimensión 0 (es decir, la dimensión del lote). Si un firma se está agrupando, asegúrate de que todas sus entradas tengan -1 en la dimensión 0.

La op BatchFunction no admite SparseTensors como entrada o salida. Internamente, cada tensor disperso se representa como tres tensores independientes que pueden tienen diferentes tamaños de 0.a dimensión.