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 de TPU. El conversor se ejecuta en una shell local o de VM de TPU. Se recomienda la shell de VM de TPU porque viene preinstalada con las herramientas de línea de comandos necesarias para el convertidor. Toma un SavedModel exportado y realiza los siguientes pasos:

1. Conversión de TPU: agrega TPUPartitionedCall y otras operaciones de TPU al modelo para que se pueda entregar en la TPU. De forma predeterminada, un modelo exportado para inferencia no tiene esas operaciones y no se puede entregar en la TPU, incluso si se entrenó en ella.
2. Agrupación en lotes: agrega operaciones de procesamiento por lotes al modelo para habilitar el procesamiento por lotes en el grafo y obtener una mejor capacidad de procesamiento.
3. Conversión BFloat16: Convierte el formato de datos del modelo de float32 a bfloat16 para mejorar el rendimiento computacional y reducir el uso de la memoria de alto ancho de banda (HBM) en la TPU.
4. Optimización de la forma de E/S: optimiza las formas del tensor para los datos transferidos entre la CPU y la TPU a fin de mejorar la utilización del ancho de banda.

Cuando se exporta un modelo, los usuarios crean alias de funciones para cualquier función que les gustaría ejecutar en la TPU. Pasan estas funciones al convertidor, y este las ubica en la TPU y las optimiza.

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

Tiempo estimado para completar los pasos anteriores: aproximadamente 20 min - 30 min

Requisitos previos

1. El modelo debe ser de TF2 y se debe exportar en formato de SavedModel.
2. El modelo debe tener un alias de función para la función de TPU. Consulta el ejemplo de código para ver cómo hacerlo. En los siguientes ejemplos, se usa tpu_func como el alias de la función de TPU.
3. Asegúrate de que la CPU de tu máquina admita las instrucciones de EXtensiones vectoriales avanzadas (AVX), ya que la biblioteca de TensorFlow (la dependencia del convertidor de inferencia de Cloud TPU) se compila para usar instrucciones de AVX. La mayoría de las CPUs son compatibles con AVX.
   1. Puedes ejecutar lscpu | grep avx para verificar si se admite el conjunto de instrucciones de AVX.

Antes de comenzar

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

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

• Configura una VM de TPU: Crea una VM de TPU nueva con la consola de Google Cloud o gcloud, o usa una VM de TPU existente para ejecutar inferencias con el modelo convertido 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 el convertidor de inferencia de Cloud TPU. 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. Utilizas estas herramientas para interactuar con la imagen del convertidor.

  Conéctate a la instancia con SSH con el 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 VM de TPU o desde la 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 que no sea raíz:

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

• Inicializa los auxiliares de credenciales de Docker:

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

Shell local

En tu shell local, configura el entorno con 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 que no sea raíz:

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

• Accede a tu entorno:

  gcloud auth login
  

• Inicializa los 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 es para realizar conversiones de modelos únicas. Establece las rutas del modelo y ajusta las opciones del convertidor según tus necesidades. En la sección Ejemplos de uso, se describen 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. Busca o crea una función en tu modelo que una todo lo que deseas ejecutar en la TPU. Si @tf.function no existe, agrégalo.
2. Cuando guardes el modelo, proporciona SaveOptions, como se muestra a continuación, para darle a model.tpu_func un alias func_on_tpu.
3. Puedes pasar el alias de esta 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 incluir varias funciones en la TPU. Simplemente crea varios alias de funciones y pásalos en converter_options_string al convertidor.

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 del modelo cuantizado antes de usarlo en producción. TensorFlow proporciona compatibilidad con TPU v5e para la cuantización.

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

La nueva función de cuantización posterior al entrenamiento en TensorFlow que se orienta a la TPU se desarrolló a partir de la función existente similar en TensorFlow Lite que se usa para apuntar a dispositivos móviles y perimetrales. Para obtener más información sobre la cuantización en general, consulta el documento de TensorFlow Lite.

Conceptos de cuantización

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

Los conceptos relacionados con otras opciones de configuración de TPU (por ejemplo, porciones, hosts, chips y TensorCores) se describen en la página Arquitectura del sistema de TPU.

• Cuantización posterior al entrenamiento (PTQ): La PTQ es una técnica que reduce el tamaño y la complejidad de procesamiento de un modelo de red neuronal sin afectar significativamente su exactitud. PTQ funciona mediante la conversión de los pesos de punto flotante y las activaciones de un modelo entrenado en números enteros de menor precisión, como números enteros de 8 bits o 16 bits. Esto puede causar una reducción significativa en el tamaño del modelo y la latencia de inferencia, mientras solo genera 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 toman las ponderaciones y las activaciones de un modelo de red neuronal. Esta información se usa para determinar los parámetros de cuantización del modelo, que son los valores que se usarán para convertir los pesos de punto flotante y las activaciones en números enteros.

• Conjunto de datos representativo: 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 tomarán los pesos y las activaciones del modelo. El conjunto de datos representativo debe cumplir las siguientes propiedades:

  • Debe representar de forma correcta las entradas reales al modelo durante la inferencia. Esto significa que debe abarcar el rango de valores que el modelo puede ver en el mundo real.
  • Debería fluir de forma colectiva a través de cada rama de condicionales (como tf.cond), si hay alguno. Esto es importante porque el proceso de cuantización debe poder controlar todas las entradas posibles al modelo, incluso si no están representadas de forma explícita en el conjunto de datos representativo.
  • Debe ser lo suficientemente grande como para recopilar suficientes estadísticas y reducir los errores. 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 diseñado de forma específica para ser representativo de las entradas del mundo real al modelo. La elección del conjunto de datos que se usará depende de la aplicación específica.

• Cuantización del rango estático (SRQ): SRQ determina el rango de valores para las ponderaciones y activaciones de un modelo de red neuronal una sola vez, durante el paso de calibración. Esto significa que se usa el mismo rango de valores para todas las entradas del modelo. Esto puede ser menos preciso que la cuantización del 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 durante el tiempo de ejecución que la cuantización del rango dinámico.

• Cuantización del rango dinámico (DRQ): DRQ determina el rango de valores de las ponderaciones y 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 la exactitud. Sin embargo, la cuantización del rango dinámico requiere más procesamiento en el tiempo de ejecución que la cuantización del rango estático.

  Característica	Cuantización del rango estático	Cuantización del rango dinámico
  Rango de valores	Se determina una vez durante la calibración	Determinado para cada entrada
  Exactitud	Puede ser menos precisa, en especial para modelos con un amplio rango de valores de entrada.	Puede ser más exacta, especialmente para modelos con un amplio rango de valores de entrada.
  Complejidad	Es más simple	Es más complejo
  Cálculo en el tiempo de ejecución	Menos procesamiento	Más procesamiento

• Cuantización de solo peso: La cuantización de solo peso es un tipo de cuantización que solo cuantiza los pesos de un modelo de red neuronal, a la vez que deja 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 a preservar la precisión del modelo.

Cómo usar la cuantización

La cuantización se puede aplicar configurando y configurando QuantizationOptions en las opciones del conversor. Las opciones más destacadas son las siguientes:

• etiquetas: colección de etiquetas que identifican el MetaGraphDef dentro de la SavedModel para cuantizar. No es necesario especificar si solo tienes un MetaGraphDef.
• firma_keys: Secuencia de claves que identifican SignatureDef que contienen entradas y salidas. Si no se especifica, se usa ["serving_default"].
• quantization_method: Método de cuantización que se aplica. Si no se especifica, se aplicará la cuantización STATIC_RANGE.
• op_set: Debe mantenerse como XLA. Actualmente, es la opción predeterminada, no es necesario especificarla.
• Representativo_datasets: especifica el conjunto de datos que se usa para calibrar los parámetros de cuantización.

Crea el conjunto de datos representativo

Un conjunto de datos representativo es, básicamente, un iterable de muestras. Donde 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 archivos TFRecord con la clase TfRecordRepresentativeDatasetSaver disponible actualmente en el paquete tf-nightly 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 conversor se puede usar para agregar lotes a un modelo. Para obtener una descripción de las opciones de lotes que pueden ajustarse, consulta Definición de las opciones de lotes.

De forma predeterminada, el conversor agrupará en lotes todas las funciones de TPU del modelo. También puede agrupar firmas y funciones proporcionadas por el usuario, lo que puede mejorar aún más el rendimiento. Cualquier función de TPU, función proporcionada por el usuario o firma por lotes, debe cumplir con los requisitos de forma estrictos de la operación de lotes.

El conversor también puede actualizar las opciones de lotes existentes. El siguiente es un ejemplo de cómo agregar el procesamiento por lotes a un modelo. Para obtener más información sobre el procesamiento por lotes, consulta el análisis detallado de la agrupación en 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 bien con tu modelo, puedes inhabilitarlas.

# 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 el convertidor de inferencia. Aquí encontrarás 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
--------------------------------

En este informe, se estima el costo de procesamiento del modelo de salida en CPU y TPU, y se 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 de convertidor.

Si deseas usar mejor la TPU, tal vez debas experimentar con la estructura del modelo y ajustar las opciones del conversor.

Nota: En el informe de conversiones actual, todas las operaciones de las funciones de TPU se registran como costos 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 la gran mayoría de las operaciones se ejecutan más rápido en la TPU.

Si tu modelo no contiene ninguna op, strings o tensores dispersos incompatibles con TPU, colocar todo el modelo en la TPU suele ser la mejor estrategia. Puedes hacerlo buscando o creando una función que una todo el modelo, creando un alias de función para él y pasándolo al convertidor.

Si tu modelo contiene partes que no pueden funcionar en la TPU (p.ej.,operaciones incompatibles con TPU, strings o tensores dispersos), la elección de las funciones de TPU depende de dónde se encuentre la parte incompatible.

• Si está al principio o al final del modelo, puedes refactorizarlo para mantenerlo en la CPU. Algunos ejemplos son las etapas de procesamiento previo y posterior de la cadena. Para obtener más información sobre cómo mover código a la CPU, consulta “¿Cómo muevo 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 dividirlo en tres partes, contener todas las ops incompatibles con TPU en la parte intermedia y hacer que se ejecute en la CPU.
• Si se trata de un tensor disperso, considera llamar a tf.sparse.to_dense en la CPU y pasar el tensor denso resultante a la parte de la TPU del modelo.

Otro factor que se debe tener en cuenta es el uso de HBM. Las tablas de incorporación pueden usar mucha HBM. Si superan la limitación de hardware de la TPU, se deben colocar en la CPU, junto con las operaciones de búsqueda.

Siempre que sea posible, solo debe existir una función de TPU bajo una firma. Si la estructura de tu modelo requiere llamar a varias funciones de TPU por solicitud de inferencia entrante, debes tener en cuenta la latencia adicional del envío de tensores entre la CPU y la TPU.

Una buena manera de evaluar la selección de funciones de TPU es consultar el Informe de conversiones. Muestra el porcentaje de procesamiento que se colocó en la TPU y un desglose del costo de cada función de TPU.

¿Cómo muevo 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 moverlas a la CPU. Este es un ejemplo de un juguete. El modelo es un modelo de lenguaje con una etapa de procesamiento previo. El código para las funciones y definiciones de capa 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. Primero, el parámetro es una cadena. En segundo lugar, la función preprocess puede contener muchas operaciones de string. Ninguno de los dos es compatible con TPU.

Para refactorizar este modelo, puedes crear otra función llamada tpu_func a fin de alojar el 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é dentro de 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, strings o tensores dispersos incompatibles con TPU?

La mayoría de las ops estándar de TensorFlow son compatibles con la TPU, pero algunas, incluidas las strings y los tensores dispersos, no lo son. El conversor no verifica si hay operaciones incompatibles con TPU. Por lo tanto, un modelo que contenga esas operaciones puede pasar la conversión. Pero cuando se ejecuta para inferencias, se producirán errores como los que se muestran a continuación.

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

Si tu modelo tiene operaciones incompatibles con TPU, deben colocarse fuera de la función de TPU. Además, la string es un formato de datos no compatible en la TPU. Por lo tanto, las variables de tipo string no deben colocarse en la función TPU. Además, los parámetros y los valores que se muestran de la función de TPU no deben ser de tipo string. De manera similar, evita colocar tensores dispersos en la función de TPU, incluidos sus parámetros y valores de retorno.

Por lo general, no es difícil refactorizar la parte incompatible del modelo y moverla a la CPU. Aquí se muestra un ejemplo:

¿Cómo se admiten operaciones personalizadas en el modelo?

Si se usan operaciones personalizadas en tu modelo, es posible que el conversor no las reconozca y no pueda 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 actualmente el código del conversor aún no es de código abierto, no se puede compilar el conversor con una operación personalizada.

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

El conversor 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 se pueden ejecutar con el puente MLIR TF2XLA más reciente o con el puente TF2XLA original.

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

Si un modelo se exportó sin un alias de función, la forma más fácil es volver a exportarlo y crear un alias de función. Si la reexportación no es una opción, aún es posible convertir el modelo proporcionando un concrete_function_name. Sin embargo, identificar la concrete_function_name correcta requiere un poco de investigación.

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

Para encontrar los nombres de las funciones concretas, debes examinar saved_model.pb.

En el siguiente ejemplo, se muestra cómo colocar 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 tiempo de compilación constante?

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

Existen dos maneras de resolver este problema.

• La mejor opción es actualizar las entradas de las operaciones para que tengan una forma conocida estáticamente antes de que XLA compile el programa de TPU. Esta compilación se realiza justo antes de que se ejecute la parte de TPU del modelo. Eso significa que la forma debe conocerse de forma estática para el momento en que TpuFunction esté a punto de ejecutarse.
• Otra opción es modificar TpuFunction para dejar de incluir la 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 en su dimensión 0 (también conocida como la dimensión de lotes). Estos requisitos de forma provienen de la operación por lotes de TensorFlow y no se pueden flexibilizar.

El incumplimiento de estos requisitos dará lugar a los siguientes errores:

1. Los tensores de entrada en 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 operación determinada deben tener el mismo tamaño de dimensión 0.
4. La dimensión 0 del tensor de salida por lotes no es igual a la suma de los tamaños de la 0a dimensión de los tensores de entrada.

Para cumplir con estos requisitos, considera proporcionar una función o una firma diferente para por lotes. También puede ser necesario modificar las funciones existentes para cumplir con estos requisitos.

Si una función se agrupa en lotes, asegúrate de que las formas input_signature de @tf.function tengan el valor None en la dimensión 0. Si una firma se agrupa en lotes, asegúrate de que todas sus entradas tengan -1 en la dimensión 0.

Para obtener una explicación completa de por qué ocurren estos errores y cómo resolverlos, consulta Cómo agrupar en lotes.

Problemas conocidos

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

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

Esto se debe a que el convertidor modifica al emisor directo de una función de TPU para que pase de llamar a la función de TPU a llamar a un stub de llamada de TPU. El stub de llamada contiene operaciones que solo pueden funcionar en la CPU. Cuando una función de TPU llama a cualquier función que finalmente llama al llamador directo, esas operaciones de CPU podrían activarse en la TPU para su ejecución, lo que generará errores de kernel faltantes. Ten en cuenta que 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 stub de la llamada, por lo que puede funcionar.

En el convertidor, implementamos la detección de esta situación. Si ves el siguiente error, significa que tu modelo alcanzó 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 esta situación de llamada a función. Si te resulta difícil hacerlo, comunícate con el equipo de Atención al cliente de Google para obtener más información.

Reference

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

El procesamiento en lotes se usa para mejorar la capacidad de procesamiento y el uso de TPU. Permite que se procesen varias solicitudes al mismo tiempo. Durante el entrenamiento, el procesamiento por lotes se puede realizar con tf.data. Durante la inferencia, se suele agregar una operación en el grafo que agrupa las solicitudes entrantes. La op espera hasta que tenga suficientes solicitudes o hasta que se alcance un tiempo de espera antes de generar un lote grande a partir de las solicitudes individuales. Consulta Definición de opciones de lotes para obtener más información sobre las diferentes opciones de procesamiento por lotes que se pueden ajustar, incluidos los tamaños de lote y los tiempos de espera.

De forma predeterminada, el convertidor inserta la op directamente antes del cálculo de la TPU. Une las funciones de TPU que proporciona el usuario y cualquier cálculo de TPU preexistente en el modelo con operaciones de agrupación en lotes. Es posible anular este comportamiento predeterminado si le indicas al convertidor qué funciones o firmas se deben 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 las firmas

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

El procesamiento por lotes de firmas agrupa todo el modelo en lotes, desde las entradas de la firma hasta las salidas de esta. A diferencia del comportamiento de lotes predeterminado del convertidor, el procesamiento por lotes de firmas agrupa el procesamiento de la TPU y el de la CPU. Esto proporciona un aumento de rendimiento de un 10% a un 20% durante la inferencia en algunos modelos.

Como sucede con todos los lotes, este tipo de lotes tiene requisitos de forma estrictos. Para garantizar que se cumplan estos requisitos de forma, las entradas de firma deben tener formas que tengan 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 formas de entrada válidas. Si no es posible, considera usar el comportamiento de lotes predeterminado o la agrupación en lotes de funciones.

Para usar el procesamiento por lotes de firmas, proporciona los nombres de firma 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 las funciones

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

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

El procesamiento por lotes de funciones se puede usar para el procesamiento de la CPU por lotes. Muchos modelos ven una mejora en el rendimiento cuando el procesamiento de la CPU se procesa por lotes. La mejor manera de realizar el procesamiento de CPU por lotes es usar el procesamiento por lotes de firmas, pero es posible que no funcione en algunos modelos. En esos casos, el procesamiento en lotes de funciones se puede usar para agrupar por lotes parte del procesamiento de la CPU además del cálculo de la TPU. Ten en cuenta que la operación de lotes no puede ejecutarse en la TPU, por lo que cualquier función de lotes que se proporcione debe llamarse en la CPU.

El procesamiento en lotes de las funciones también se puede usar para satisfacer los requisitos de forma estrictos que impone la operación por lotes. En los casos en que las funciones de TPU no cumplen con los requisitos de forma de la operación, se puede usar el procesamiento en lotes de funciones para indicarle al convertidor que agrupe diferentes funciones.

Con el fin de usar esto, genera un function_alias para la función que se debe agrupar en lotes. Puedes hacerlo buscando o creando una función en tu modelo que una 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 por lotes. Agrega @tf.function si aún no tiene uno. Es importante proporcionar el input_signature a @tf.function. La 0a 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 formas de entrada válidas. Cuando guardes el modelo, proporciona SaveOptions como los que se muestran a continuación para asignar a model.batch_func un alias "batch_func". Luego, puedes pasar el alias de esta función al convertidor.

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

• num_batch_threads (número entero): Cantidad de subprocesos de programación para procesar lotes de trabajo. Determina la cantidad de lotes procesados en paralelo. Debería ser similar a la cantidad de núcleos de TPU disponibles.
• max_batch_size: (número entero) Tamaño del lote máximo permitido. Puede ser mayor que allowed_batch_sizes para usar una gran división por lotes.
• batch_timeout_micros (número entero): Cantidad máxima de microsegundos que se esperarán antes de generar un lote incompleto.
• allowed_batch_sizes (lista de números enteros): Si la lista no está vacía, rellenará lotes hasta el tamaño más cercano en la lista. La lista debe aumentar monótonamente y el elemento final debe ser menor o igual que max_batch_size.
• max_enqueued_batches (número entero): Es la cantidad máxima de lotes en cola para procesar antes de que las solicitudes fallen rápidamente.

Actualiza las opciones de lotes existentes

Puedes agregar o actualizar las opciones de lotes si ejecutas la imagen de Docker que especifica lotes_options y configuras disable_default_optimizations como verdadero con la marca --converter_options_string. Las opciones de lotes se aplicarán a cada función de TPU o a cada op preexistente.

Nota: No puedes actualizar las opciones de lotes si configuras tpu_functions juntas 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 la 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 dimensión de lotes (0). Los tensores de salida se dividen en su dimensión 0. A fin de realizar estas operaciones, la operación de lotes tiene requisitos de forma estrictos para sus entradas y salidas.

Explicación

Para entender estos requisitos, es útil primero comprender cómo se realiza el procesamiento por lotes. En el siguiente ejemplo, se agrupa en lotes una op tf.matmul simple.

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 las entradas A y B con las formas (2, 3, 2) y (2, 2, 4).

Se alcanzó el tiempo de espera del lote. 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. Los tensores en lotes se forman mediante la concatenación de las solicitudes n.o 1 y n.o 2 a lo largo de la dimensión del lote (0). Como la A del n.o 1 tiene una forma de (1, 3, 2) y la A del n.o 2 tiene una forma de (2, 3, 2), cuando se concatenan a lo largo de la dimensión del lote (0), 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. Para ello, la op lo hace dividiendo por la dimensión de lote (0) de cada tensor de salida. Decide cómo dividir la dimensión 0 en función de la forma de las entradas originales. Dado que las formas de la solicitud n.o 1 tienen una dimensión 0 de 1, su resultado tiene una dimensión 0 de 1 para una forma de (1, 3, 4). Dado que las formas de la solicitud n.o 2 tienen una dimensión 0 de 2, su resultado tiene una dimensión 0 de 2 para una forma de (2, 3, 4).

Requisitos de la forma

Para realizar la concatenación de entradas y la división de resultados descrita anteriormente, la operación de lotes tiene los siguientes requisitos de forma:

1. Las entradas al procesamiento por lotes no pueden ser escalares. Para concatenar a lo largo de la dimensión 0, los tensores deben tener al menos dos dimensiones.

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

   Si no cumples con este requisito, se mostrará un error como el siguiente: Batching input tensors must have at least one dimension. Una solución simple para este error es hacer que el escalar sea un vector.

2. En diferentes solicitudes de inferencia (por ejemplo, distintas invocaciones de ejecución de sesión), los tensores de entrada con el mismo nombre tienen el mismo tamaño para cada dimensión, excepto la dimensión 0. Esto permite que las entradas se concatenen de forma limpia a lo largo de su dimensión 0.

   En la explicación anterior, la A de la solicitud n.o 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 2 cumple con este requisito con (2, 3, 2). De manera similar, la B de la solicitud n.o 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 cumples con este requisito, se mostrará un error como el siguiente: Dimensions of inputs should match.

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

   En la explicación anterior, todos los tensores de la solicitud #1 tienen un tamaño de dimensión 0 de 1. Esto permite que la operación de lotes sepa que su resultado debe tener un tamaño de dimensión 0 de 1. De manera similar, los tensores de solicitud n.o 2 tienen un tamaño de dimensión 0 de 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), produce (1, 3, 4) para la solicitud n° 1 y (2, 3, 4) para la solicitud n° 2.

   Si no lo cumples, 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 la 0a dimensión de la forma de cada tensor de salida debe ser la suma del tamaño 0 de todos los tensores de entrada (más cualquier relleno ingresado por la op para cumplir con el siguiente allowed_batch_size más grande). Esto permite que la op más grande divida los tensores de salida en su 0a dimensión en función de la dimensión 0 de la entrada.

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

   Si no cumples 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.

Cómo resolver errores de requisitos de forma

Para cumplir con estos requisitos, considera proporcionar una función o una firma diferente para por lotes. También puede ser necesario modificar las funciones existentes para cumplir con estos requisitos.

Si se está agrupando una función en lotes, asegúrate de que todas las formas de input_signature de @tf.function tengan None en la dimensión 0 (también conocida como la dimensión del lote). Si se está agrupando una firma, asegúrate de que todas sus entradas tengan -1 en la dimensión 0.

La op BatchFunction no admite SparseTensors como entradas o salidas. Internamente, cada tensor disperso se representa como tres tensores separados que pueden tener diferentes tamaños de 0.a dimensión.