Introdução ao conversor de inferência do Cloud TPU v5e [pré-lançamento público]

Introdução

O conversor de inferência do Cloud TPU prepara e otimiza um modelo do TensorFlow 2 (TF2) para inferência da TPU. O conversor é executado em um shell de VM local ou de TPU. O shell da VM da TPU é recomendado porque vem pré-instalado com as ferramentas de linha de comando necessárias para o conversor. Ele usa um SavedModel exportado e executa as seguintes etapas:

  1. Conversão de TPU: adiciona TPUPartitionedCall e outras operações de TPU ao modelo para torná-lo veiculável na TPU. Por padrão, um modelo exportado para inferência não tem essas operações e não pode ser exibido na TPU, mesmo que tenha sido treinado na TPU.
  2. Lotes: adiciona operações de lote ao modelo para ativar os lotes no gráfico e melhorar a capacidade.
  3. Conversão BFloat16: ele converte o formato de dados do modelo de float32 para bfloat16, melhorando o desempenho computacional e reduzindo o uso de memória de alta largura de banda (HBM, na sigla em inglês) na TPU.
  4. Otimização das formas de E/S: otimiza as formas do tensor para dados transferidos entre a CPU e a TPU para melhorar a utilização da largura de banda.

Ao exportar um modelo, os usuários criam aliases de função para qualquer função que queiram executar na TPU. Essas funções são transmitidas ao conversor, e ele as coloca na TPU e as otimiza. Se essas funções contiverem operações incompatíveis com TPU, será possível ativar o posicionamento do dispositivo flexível para que as operações incompatíveis sejam executadas na CPU.

O conversor de inferência do Cloud TPU está disponível como uma imagem do Docker que pode ser executada em qualquer ambiente com o Docker instalado.

Tempo estimado para concluir as etapas mostradas acima: cerca de 20 minutos a 30 minutos

Pré-requisitos

  1. O modelo precisa ser um modelo TF2 e exportado no formato TASK (link em inglês).
  2. O modelo precisa ter um alias de função para a função da TPU. Consulte o exemplo de código para saber como fazer isso. Os exemplos a seguir usam tpu_func como o alias da função da TPU.
  3. Verifique se a CPU da sua máquina é compatível com as instruções Advanced Vector eXtensions (AVX), porque a biblioteca TensorFlow (a dependência do Cloud TPU Inference Converter) é compilada para usar instruções do AVX. A maioria das CPUs tem suporte a AVX.
    1. Você pode executar lscpu | grep avx para verificar se o conjunto de instruções do AVX tem suporte.

Antes de começar

Antes de iniciar a configuração, faça o seguinte:

  • Criar um novo projeto: na página do seletor de projetos do console do Google Cloud, selecione ou crie um projeto do Cloud.

  • Configure uma VM do Cloud TPU: crie uma nova VM do Cloud TPU usando o console do Google Cloud ou o gcloud ou use uma VM atual do Cloud TPU para executar a inferência com o modelo convertido na VM do Cloud TPU.

    • Verifique se a imagem da VM do Cloud TPU é baseada no TensorFlow. Por exemplo, --version=tpu-vm-tf-2.11.0.
    • O modelo convertido será carregado e exibido nesta VM do Cloud TPU.

  • Verifique se você tem as ferramentas de linha de comando necessárias para usar o conversor de inferência do Cloud TPU. É possível instalar o SDK Google Cloud e o Docker localmente ou usar uma VM do Cloud TPU que tenha esse software instalado por padrão. Use essas ferramentas para interagir com a imagem do conversor.

    É possível executar o SSH na VM do Cloud TPU usando o seguinte comando: 

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

Configuração do ambiente

Configure o ambiente a partir do shell da VM da TPU ou do shell local.

Shell da VM da TPU

  • No shell da VM da TPU, execute os seguintes comandos para permitir o uso do Docker não raiz:

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

  • Inicialize os auxiliares de credenciais do Docker:

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

Shell local

No shell local, configure o ambiente seguindo estas etapas:

  • Instale o SDK Cloud, que inclui a ferramenta de linha de comando gcloud.

  • Instale o Docker:

  • Permitir o uso do Docker não raiz:

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

  • Faça login no seu ambiente:

    gcloud auth login

  • Inicialize os auxiliares de credenciais do Docker:

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

  • Extraia a imagem do Docker do conversor de inferência:

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

Imagem do conversor

A imagem serve para fazer conversões únicas de modelo. Defina os caminhos do modelo e ajuste as opções do conversor para atender às suas necessidades. A seção Exemplos de uso fornece vários casos de uso comuns.

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

Inferência com o modelo convertido na VM do Cloud 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))

Exemplos de uso

Adicionar um alias de função para a função da TPU

  1. Encontre ou crie uma função no modelo que envolva tudo o que você quer executar na TPU. Se @tf.function não existir, adicione-o.
  2. Ao salvar o modelo, forneça SaveOptions como abaixo para fornecer a model.tpu_func um alias func_on_tpu.
  3. Você pode transmitir esse alias de função para o 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)

Converter um modelo com várias funções da TPU

É possível colocar várias funções na TPU. Basta criar vários aliases de função e transmiti-los em converter_options_string para o conversor.

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

Quantização

A quantização é uma técnica que reduz a precisão dos números usados para representar os parâmetros de um modelo. Isso resulta em um tamanho de modelo menor e computação mais rápida. Um modelo quantizado fornece ganhos na capacidade de inferência, bem como menor uso de memória e tamanho de armazenamento ao custo de pequenas quedas de precisão.

O novo recurso de quantização pós-treinamento do TensorFlow, que é voltado para a TPU, foi desenvolvido a partir de um recurso semelhante no TensorFlow Lite, usado para segmentar dispositivos móveis e de borda. Para saber mais sobre a quantização em geral, consulte o documento do TensorFlow Lite.

Conceitos de quantização

Esta seção define conceitos relacionados especificamente à quantização com o conversor de inferência.

Conceitos relacionados a outras configurações de TPU (por exemplo, frações, hosts, chips e TensorCores) são descritos na página Arquitetura do sistema de TPU.

  • Quantização pós-treinamento (PTQ, na sigla em inglês): é uma técnica que reduz o tamanho e a complexidade computacional de um modelo de rede neural sem afetar significativamente a acurácia. A PTQ converte os pesos e ativações de ponto flutuante de um modelo treinado em números inteiros de precisão mais baixa, como números inteiros de 8 ou 16 bits. Isso pode causar uma redução significativa no tamanho do modelo e na latência de inferência, mas causar uma pequena perda na acurácia.

  • Calibração: a etapa de calibração para a quantização é o processo de coletar estatísticas sobre o intervalo de valores que os pesos e ativações de um modelo de rede neural usam. Essas informações são usadas para determinar os parâmetros de quantização do modelo, que são os valores usados para converter os pesos e ativações de ponto flutuante em números inteiros.

  • Conjunto de dados representativo: um conjunto representativo para quantização é um pequeno conjunto de dados que representa os dados reais de entrada para o modelo. Ela é usada durante a etapa de calibração da quantização para coletar estatísticas sobre o intervalo de valores que os pesos e ativações do modelo vão assumir. O conjunto de dados representativo precisa satisfazer as seguintes propriedades:

    • Ela precisa representar corretamente as entradas reais no modelo durante a inferência. Isso significa que ele precisa cobrir o intervalo de valores que o modelo provavelmente verá no mundo real.
    • Eles fluem coletivamente por cada ramificação de condicionais (como tf.cond), se houver. Isso é importante porque o processo de quantização precisa ser capaz de processar todas as entradas possíveis para o modelo, mesmo que elas não estejam representadas explicitamente no conjunto de dados representativo.
    • Ela precisa ser grande o suficiente para coletar estatísticas suficientes e reduzir erros. Como regra geral, é recomendável usar mais de 200 amostras representativas.

    O conjunto de dados representativo pode ser um subconjunto do conjunto de treinamento ou um conjunto separado, projetado especificamente para representar as entradas reais do modelo. A escolha de qual conjunto de dados usar depende do aplicativo específico.

  • Quantização de intervalo estático (SRQ, na sigla em inglês): o SRQ determina o intervalo de valores para os pesos e as ativações de um modelo de rede neural uma vez durante a etapa de calibração. Isso significa que o mesmo intervalo de valores é usado para todas as entradas do modelo. Isso pode ser menos preciso do que a quantização de intervalo dinâmico, especialmente para modelos com uma ampla gama de valores de entrada. No entanto, a quantização de intervalo estático requer menos computação no tempo de execução do que a quantização de intervalo dinâmico.

  • Quantização de intervalo dinâmico (DRQ, na sigla em inglês): a DRQ determina o intervalo de valores para os pesos e as ativações de um modelo de rede neural para cada entrada. Isso permite que o modelo se adapte ao intervalo de valores dos dados de entrada, o que pode melhorar a acurácia. No entanto, a quantização de intervalo dinâmico requer mais computação no ambiente de execução do que a quantização de intervalo estático.

    Recurso Quantização de intervalo estático Quantização de intervalo dinâmico
    Intervalo de valores Determinado uma vez, durante a calibração Determinado para cada entrada
    Acurácia Pode ser menos preciso, especialmente para modelos com uma ampla gama de valores de entrada. Pode ser mais preciso, especialmente para modelos com uma grande variedade de valores de entrada.
    Complexidade Mais simples Mais complexo
    Computação no tempo de execução Menos computação Mais computação

  • Quantização somente de peso: esse tipo de quantização apenas quantiza os pesos de um modelo de rede neural, deixando as ativações em ponto flutuante. Essa pode ser uma boa opção para modelos sensíveis à acurácia, porque ajuda a preservar a acurácia.

Como usar a quantização

A quantização pode ser aplicada configurando QuantizationOptions para as opções do conversor. As opções mais importantes são as seguintes:

  • Tags: coleção de tags que identificam o MetaGraphDef na SavedModel para quantificar. Não é necessário especificar se você tem apenas um MetaGraphDef.
  • signature_keys: sequência de chaves que identifica SignatureDef que contêm entradas e saídas. Se não for especificado, ["serving_default"] será usado.
  • quantization_method: método de quantização a ser aplicado. Se não for especificada, a quantização STATIC_RANGE vai ser aplicada.
  • op_set: precisa ser mantido como XLA. Essa é a opção padrão e não precisa ser especificada.
  • Representação_de_dados: especifica o conjunto de dados usado para calibrar os parâmetros de quantização.

Criar o conjunto de dados representativo

Um conjunto de dados representativo é essencialmente um iterável de amostras. Onde um exemplo é um mapa de: {input_key: input_value} Exemplo:

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

Os conjuntos de dados representativos precisam ser salvos como arquivos TFRecord usando a classe TfRecordRepresentativeDatasetSaver atualmente disponível no pacote tf-nightly pip. Exemplo:

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

Examples

O exemplo a seguir quantifica o modelo com a chave de assinatura de serving_default e o alias de função de tpu_func:

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

Ativar a colocação de dispositivo de software

O posicionamento de dispositivo flexível pode mover operações incompatíveis com TPU para serem executadas na CPU em tempo de execução. Caso o modelo tenha uma operação incompatível com TPU dentro da função da TPU que não possa ser separada facilmente, considere ativar o posicionamento do dispositivo flexível configurando enable_soft_device_placement=True na função da TPU. Essa sinalização sinaliza o ambiente de execução para verificar se há operações incompatíveis com TPU e executá-las na CPU.

A posição do dispositivo de software também é útil para depuração. Com ela, é possível usar tf.print, que é incompatível com TPU, na função da TPU para imprimir os valores dos tensores.

O posicionamento de dispositivo soft funciona melhor com a ponte MLIR. Uma sinalização extra precisa ser definida.

tpu_functions {
  function_alias: "tpu_func"
  enable_soft_device_placement: true
}

Adicionar lotes

O conversor pode ser usado para adicionar lotes a um modelo. Para ver uma descrição das opções de lote que podem ser ajustadas, consulte Definição das opções de lote.

Por padrão, o conversor agrupa em lote todas as funções de TPU no modelo. Ele também pode agrupar assinaturas e funções fornecidas pelo usuário para melhorar ainda mais o desempenho. Qualquer função da TPU, função ou assinatura fornecida pelo usuário que é agrupada em lote precisa atender aos requisitos rígidos de forma da operação em lote.

O conversor também pode atualizar as opções de lote atuais. Confira abaixo um exemplo de como adicionar lotes a um modelo. Para mais informações sobre lotes, consulte Análise detalhada em 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
}

Desativar otimizações de forma bfloat16 e E/S

BFloat16 e otimizações de forma de E/S são ativadas por padrão. Se não funcionarem bem com o modelo, elas poderão ser desativadas.

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

Relatório de conversões

Você pode encontrar esse relatório de conversão no registro depois de executar o conversor de inferência. Veja um exemplo abaixo.

-------- 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 relatório estima o custo computacional do modelo de saída na CPU e na TPU e detalha ainda mais o custo da TPU para cada função, o que reflete a seleção das funções de TPU nas opções de conversor.

Para usar melhor a TPU, teste a estrutura do modelo e ajuste as opções de conversor.

Perguntas frequentes

Quais funções devo colocar na TPU?

É melhor colocar o máximo possível do modelo na TPU, já que a grande maioria das operações é executada com mais rapidez na TPU.

Se o modelo não tiver operações, strings ou tensores esparsos incompatíveis com TPU, a melhor estratégia geralmente é colocar todo o modelo na TPU. E você pode fazer isso encontrando ou criando uma função que une todo o modelo, criando um alias de função para ela e transmitindo isso para o conversor.

Caso o modelo contenha partes que não funcionam na TPU (por exemplo, operações, strings ou tensores esparsos incompatíveis com TPU), a escolha das funções da TPU vai depender do local em que a parte incompatível está.

  • Se estiver no início ou no fim do modelo, será possível refatorar o modelo para mantê-lo na CPU. Exemplos são os estágios de pré e pós-processamento de strings. Para saber mais sobre como mover o código para a CPU, consulte "Como mover uma parte do modelo para a CPU?". Ele mostra uma maneira típica de refatorar o modelo.
  • Se ele estiver no meio do modelo e tiver apenas algumas operações, o posicionamento do dispositivo flexível poderá ser considerado. No entanto, ela pode gerar uma latência de inferência maior devido às comunicações extras entre a CPU e a TPU.
  • Se ele estiver no meio do modelo e tiver um grande bloco de operações, o posicionamento flexível do dispositivo tornaria a inferência muito cara. Portanto, é melhor dividir o modelo em três partes e conter todas as operações incompatíveis com TPU na parte do meio. Em seguida, faça com que ele seja executado na CPU.
  • Se for um tensor esparso, chame tf.sparse.to_dense na CPU e transmita o tensor denso resultante para a parte da TPU do modelo.

Outro fator a considerar é o uso de HBM. A incorporação de tabelas pode usar muito HBM. Se elas crescerem além da limitação de hardware da TPU, elas precisarão ser colocadas na CPU, junto com as operações de pesquisa.

Sempre que possível, só é permitido ter uma função de TPU em uma assinatura. Se a estrutura do modelo exigir a chamada de várias funções de TPU por solicitação de inferência recebida, considere a latência extra do envio de tensores entre a CPU e a TPU.

Uma boa maneira de avaliar a seleção de funções da TPU é conferir o Relatório de conversões. Ele mostra a porcentagem de computação realizada na TPU e os detalhes do custo de cada função da TPU.

Como faço para mover uma parte do modelo para a CPU?

Se o modelo contiver peças que não possam ser exibidas na TPU, será necessário refatorar o modelo para movê-las para a CPU. Aqui está um exemplo de brinquedo. O modelo é um modelo de linguagem com um estágio de pré-processamento. Para simplificar, o código das definições e funções de camada é omitido.

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)

Esse modelo não pode ser exibido diretamente na TPU por dois motivos. Primeiro, o parâmetro é uma string. Segundo, a função preprocess pode conter muitas operações de string. As duas não são compatíveis com TPU.

Para refatorar esse modelo, crie outra função chamada tpu_func para hospedar o bert_layer de uso intensivo de computação. Em seguida, crie um alias de função para tpu_func e transmita-o ao conversor. Dessa forma, tudo dentro de tpu_func será executado na TPU e tudo o que restar em model_func será executado na 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)

O que devo fazer se o modelo tiver operações, strings ou tensores esparsos incompatíveis com TPU?

A maioria das operações padrão do TensorFlow é compatível com a TPU, mas algumas, incluindo tensores e strings esparsas, não são. O conversor não verifica se há operações incompatíveis com TPU. Portanto, um modelo com essas operações pode passar a conversão. No entanto, ao executá-lo para inferência, ocorrerão erros como abaixo.

'tf.StringToNumber' op isn't compilable for TPU device. enable soft_device_placement option to run on CPU

Se o modelo tem operações incompatíveis com TPU, elas precisam ser colocadas fora da função de TPU. Além disso, a string é um formato de dados incompatível na TPU. Portanto, não coloque variáveis do tipo string na função da TPU. Além disso, os parâmetros e valores de retorno da função da TPU também não podem ser digitados como string. Da mesma forma, evite colocar tensores esparsos na função da TPU, incluindo os parâmetros e valores de retorno dela.

Geralmente, não é difícil refatorar a parte incompatível do modelo e movê-la para a CPU. Confira um exemplo. No entanto, há casos em que as operações incompatíveis com TPU ficam profundamente aninhadas no modelo e são difíceis de separar. Para esses modelos, considere ativar o posicionamento de dispositivos não associados para executar as operações na CPU no momento da execução. Mas lembre-se do impacto no desempenho.

Como oferecer suporte a operações personalizadas no modelo?

Se operações personalizadas forem usadas no seu modelo, o conversor poderá não reconhecê-las e não converter o modelo. Isso ocorre porque a biblioteca de operações da operação personalizada, que contém a definição completa dela, não está vinculada ao conversor.

Como, no momento, o conversor ainda não tem código aberto, ele não pode ser criado com uma operação personalizada.

O que devo fazer se eu tiver um modelo do TensorFlow 1?

O conversor não oferece suporte a modelos do TensorFlow 1. Os modelos do TensorFlow 1 precisam ser migrados para o TensorFlow 2.

Preciso ativar a ponte MLIR ao executar meu modelo?

A maioria dos modelos convertidos pode ser executada com a ponte MLIR mais recente do TF2XLA ou a ponte TF2XLA original.

Durante o treinamento e a inferência, o posicionamento do dispositivo soft no TensorFlow 2 funciona melhor com a ponte MLIR. A ponte MLIR pode ser ativada usando estas instruções.

Como converter um modelo que já foi exportado sem um alias de função?

Se um modelo foi exportado sem um alias de função, a maneira mais fácil é exportá-lo novamente e criar um alias de função. Mesmo que a reexportação não seja uma opção, ainda será possível converter o modelo fornecendo um concrete_function_name. No entanto, identificar o concrete_function_name correto requer algum trabalho de detetive.

Os aliases de função são um mapeamento de uma string definida pelo usuário para um nome de função concreto. Eles facilitam a referência a uma função específica no modelo. O conversor aceita aliases de função e nomes de função brutos concretos.

Os nomes concretos de funções podem ser encontrados examinando o saved_model.pb.

O exemplo a seguir mostra como colocar uma função concreta chamada __inference_serve_24 na 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"
    }'

Como resolver um erro de restrição de constante tempo de compilação?

Para treinamento e inferência, o XLA exige que as entradas de determinadas operações tenham um formato conhecido no tempo de compilação da TPU. Isso significa que, quando o XLA compila a parte da TPU do programa, as entradas para essas operações precisam ter uma forma conhecida estaticamente.

Há duas maneiras de resolver esse problema.

  • A melhor opção é atualizar as entradas da operação para que tenham um formato conhecido estaticamente no momento em que o XLA compilar o programa de TPU. Essa compilação acontece logo antes de a parte da TPU do modelo ser executada. Isso significa que a forma precisa ser conhecida estaticamente no momento em que a TpuFunction estiver prestes a ser executada.
  • Outra opção é modificar o TpuFunction para não incluir mais a operação problemática.

Por que estou recebendo um erro de lote de formas?

O lote tem requisitos de forma rigorosos que permitem que as solicitações recebidas sejam agrupadas na 0a dimensão (também conhecida como dimensão de lote). Esses requisitos de forma vêm das operações de lote do TensorFlow e não podem ser atenuados.

O não cumprimento desses requisitos resulta em erros como os seguintes:

  1. Os tensores de entrada em lote precisam ter pelo menos uma dimensão.
  2. As dimensões das entradas precisam ser iguais.
  3. Os tensores de entrada em lote fornecidos em uma determinada invocação de operação precisam ter o mesmo tamanho de dimensão 0.
  4. A dimensão 0 do tensor de saída em lote não é igual à soma dos tamanhos de 0a dimensão dos tensores de entrada.

Para atender a esses requisitos, forneça uma função ou assinatura diferente para lotes. Também pode ser necessário modificar as funções atuais para atender a esses requisitos.

Se uma função estiver sendo loteada, verifique se todas as formas de input_signature da @tf.function têm None na dimensão 0. Se uma assinatura estiver sendo enviada em lote, verifique se todas as entradas têm -1 na 0a dimensão.

Para uma explicação completa sobre o motivo desses erros e como resolvê-los, consulte Análise detalhada dos lotes.

Problemas conhecidos

A função da TPU não pode chamar indiretamente outra função da TPU

O conversor consegue lidar com a maioria dos cenários de chamada de função além do limite da CPU-TPU, mas há um caso extremo raro em que ele falharia. É quando uma função de TPU chama indiretamente outra função de TPU.

Isso ocorre porque o conversor modifica o autor da chamada direta de uma função de TPU, deixando de chamar a própria função da TPU para chamar um stub de chamada da TPU. O stub de chamada contém operações que só funcionam na CPU. Quando uma função da TPU chama qualquer função que chame o autor da chamada direto, essas operações de CPU podem ser executadas na TPU para execução, gerando erros de kernel ausentes. Esse caso é diferente de uma função de TPU que chama diretamente outra função de TPU. Nesse caso, o conversor não modifica nenhuma função para chamar o stub de chamada para que ele funcione.

No conversor, implementamos a detecção desse cenário. Se você vir o erro a seguir, isso significa que seu modelo atingiu esse caso extremo:

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.

A solução geral é refatorar o modelo para evitar esse tipo de cenário de chamada de função. Se você achar isso difícil, entre em contato com a equipe de suporte do Google para mais informações.

Referência

Opções de conversor no 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 do not 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. Or soft
    // device placement can be enabled. But this will incur performance loss.
    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 such as
    // enable_soft_device_placement 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;
  }

  // Set to true to enable outside compilation for this TPU function. If the TPU
  // function has TPU-incompatible ops, outside compilation can automatically
  // move the ops to execute on the CPU at the runtime.
  bool enable_soft_device_placement = 2;
}

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

Informações detalhadas sobre lotes

Lotes são usados para melhorar a capacidade e o uso da TPU. Ele permite que várias solicitações sejam processadas ao mesmo tempo. Durante o treinamento, os lotes podem ser feitos usando tf.data. Durante a inferência, isso geralmente é feito adicionando uma operação ao gráfico que agrupa as solicitações recebidas. A operação espera até ter solicitações suficientes ou atingir o tempo limite antes de gerar um lote grande das solicitações individuais. Consulte Definição das opções de lote para mais informações sobre as diferentes opções de lote que podem ser ajustadas, incluindo tamanhos de lote e tempos limite.

Lotes no gráfico

Por padrão, o conversor insere a operação de lote diretamente antes do cálculo da TPU. Ele une as funções de TPU fornecidas pelo usuário e qualquer computação de TPU preexistente no modelo com operações de lote. É possível substituir esse comportamento padrão informando ao conversor quais funções e/ou assinaturas precisam ser agrupadas em lote.

O exemplo a seguir mostra como adicionar os lotes padrão.

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
}

Enviar assinaturas em lote

O agrupamento de assinaturas agrupa todo o modelo, começando nas entradas da assinatura e indo até as saídas da assinatura. Ao contrário do comportamento de lote padrão do conversor, o lote de assinaturas agrupa o cálculo da TPU e da CPU. Isso dá a um ganho de desempenho de 10% a 20% durante a inferência em alguns modelos.

Como todos os lotes, os lotes de assinatura têm requisitos rigorosos de forma. Para ajudar a garantir que esses requisitos de forma sejam atendidos, as entradas de assinatura precisam ter formas com pelo menos duas dimensões. A primeira dimensão é o tamanho do lote e deve ter um tamanho -1. Por exemplo, (-1, 4), (-1) ou (-1, 128, 4, 10) são formas de entrada válidas. Se isso não for possível, use o comportamento de lote padrão ou o envio de lotes de funções.

Para usar o agrupamento de assinaturas, forneça os nomes das assinaturas como signature_name(s) usando 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"
  }
}

Lotes de funções

O lote de funções pode ser usado para informar ao conversor quais funções devem ser agrupadas. Por padrão, o conversor agrupa em lote todas as funções da TPU. O lote de funções substitui esse comportamento padrão.

O lote de funções pode ser usado para calcular em lote a CPU. Muitos modelos apresentam uma melhoria de desempenho quando a computação da CPU é feita em lote. A melhor maneira de lotar a computação de CPU é usando lotes de assinatura. No entanto, isso pode não funcionar em alguns modelos. Nesses casos, é possível usar o lote de funções para agrupar parte do cálculo da CPU, além do cálculo da TPU. A operação de lote não pode ser executada na TPU. Portanto, qualquer função de lote fornecida precisa ser chamada na CPU.

O agrupamento de funções também pode ser usado para atender aos requisitos de formato rígidos impostos pela operação de lote. Quando as funções do TPU não atendem aos requisitos de forma da operação de lote, o lote da função pode ser usado para instruir o conversor a agrupar diferentes funções.

Para usar isso, gere um function_alias para a função que precisa ser agrupada. Para fazer isso, encontre ou crie uma função no modelo que encapsule tudo o que você quer em lote. Verifique se a função atende aos requisitos rígidos de forma impostos pela operação em lote. Adicione @tf.function se ainda não houver um. É importante fornecer o input_signature ao @tf.function. A dimensão 0 precisa ser None porque é a dimensão do lote, portanto, não pode ter um tamanho fixo. Por exemplo, [None, 4], [None] ou [None, 128, 4, 10] são formas de entrada válidas. Ao salvar o modelo, forneça SaveOptions como os mostrados abaixo para dar a model.batch_func um alias "batch_func". Em seguida, é possível transmitir esse alias de função para o 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)

Em seguida, transmita os function_alias(s) usando as 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"
  }
}

Definição de opções de lote

  • num_batch_threads (número inteiro): número de linhas de execução de programação para processar lotes de trabalho. Determina o número de lotes processados em paralelo. Isso deve estar aproximadamente de acordo com o número de núcleos de TPU disponíveis.
  • max_batch_size: (número inteiro) tamanho máximo do lote permitido. Pode ser maior que allowed_batch_sizes para utilizar a divisão em lote grande.
  • batch_timeout_micros (número inteiro): número máximo de microssegundos a serem aguardados antes de gerar um lote incompleto.
  • allowed_batch_sizes: (lista de números inteiros) se a lista não estiver vazia, ela preencherá os lotes até o tamanho mais próximo na lista. A lista precisa estar aumentada monotonicamente e o elemento final precisa ser menor ou igual a max_batch_size.
  • max_enqueued_batches (número inteiro): número máximo de lotes enfileirados para processamento antes que as solicitações falhem rapidamente.

Como atualizar as opções de lote atuais

É possível adicionar ou atualizar as opções de lote executando a imagem do Docker que especifica batch_options e definindo disable_default_optimizations como verdadeiro usando a sinalização --converter_options_string. As opções de lote serão aplicadas a cada função de TPU ou operação de lote preexistente.

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 lote

Os lotes são criados pela concatenação dos tensores de entrada nas solicitações ao longo da dimensão de lote (0a). Os tensores de saída são divididos ao longo de sua dimensão 0. Para executar essas operações, a operação em lote tem requisitos rigorosos de forma para entradas e saídas.

Tutorial

Para entender esses requisitos, é útil primeiro entender como os lotes são executados. No exemplo abaixo, estamos agrupando uma op tf.matmul simples.

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

A primeira solicitação de inferência produz as entradas A e B com as formas (1, 3, 2) e (1, 2, 4), respectivamente. A segunda solicitação de inferência produz as entradas A e B com as formas (2, 3, 2) e (2, 2, 4).

solicitação de inferência 1

O tempo limite dos lotes foi atingido. O modelo aceita um tamanho de lote de 3, portanto, as solicitações de inferência 1 e 2 são agrupadas sem preenchimento. Os tensores agrupados são formados pela concatenação das solicitações 1 e 2 ao longo da dimensão do lote (0a). Como o A do número 1 tem a forma (1, 3, 2) e o A do 2 tem uma forma de (2, 3, 2), quando eles são concatenados ao longo da dimensão do lote (0a), a forma resultante é (3, 3, 2).

solicitação em lote

A tf.matmul é executada e produz uma saída com a forma (3, 3, 4).

solicitação matmul em lote

A saída de tf.matmul é agrupada em lote, portanto, precisa ser dividida novamente em solicitações separadas. A operação em lote faz isso dividindo ao longo da dimensão do lote (0o) de cada tensor de saída. e decide como dividir a dimensão 0 com base na forma das entradas originais. Como as formas da solicitação no 1 têm uma dimensão 0 de 1, a saída tem uma 0a dimensão de 1 para uma forma de (1, 3, 4). Como as formas da solicitação no 2 têm uma dimensão 0 de 2, a saída tem uma dimensão 0 de 2 para uma forma de (2, 3, 4).

resultados da solicitação de inferência

Requisitos de forma

Para executar a concatenação de entrada e a divisão de saída descritas acima, a op em lote tem os seguintes requisitos de forma:

  1. As entradas para lotes não podem ser escalares. Para concatenar ao longo da dimensão 0, os tensores precisam ter pelo menos duas dimensões.

    No tutorial acima. Nem A nem B são escalares.

    Se você não atender a esse requisito, o seguinte erro será exibido: Batching input tensors must have at least one dimension. Uma correção simples para esse erro é tornar o escalar um vetor.

  2. Em solicitações de inferência distintas (por exemplo, invocações de execução de sessão distintas), os tensores de entrada com o mesmo nome têm o mesmo tamanho para cada dimensão, exceto a 0. Isso permite que as entradas sejam concatenadas corretamente ao longo da 0a dimensão.

    No tutorial acima, a solicitação A da solicitação 1 tem a forma (1, 3, 2). Isso significa que qualquer solicitação futura precisa produzir uma forma com o padrão (X, 3, 2). A solicitação 2 atende a esse requisito com (2, 3, 2). Da mesma forma, a solicitação B da solicitação 1 tem a forma (1, 2, 4). Portanto, todas as solicitações futuras precisam produzir uma forma com o padrão (X, 2, 4).

    Se você não atender a esse requisito, o seguinte erro será exibido: Dimensions of inputs should match.

  3. Para uma solicitação de inferência, todas as entradas precisam ter o mesmo tamanho de dimensão 0. Se diferentes tensores de entrada para a operação de lote tiverem dimensões 0 diferentes, a operação em lote não saberá como dividir os tensores de saída.

    No tutorial acima, todos os tensores da solicitação no 1 têm um tamanho de dimensão 0 de 1. Isso permite que a operação de lote saiba que a saída precisa ter um tamanho de dimensão 0 de 1. Da mesma forma, os tensores da solicitação no 2 têm um tamanho de 0 a dimensão de 2. Portanto, a saída terá um tamanho de dimensão 0 de 2. Quando a operação de lote divide a forma final de (3, 3, 4), ela produz (1, 3, 4) para a solicitação no 1 e (2, 3, 4) para a solicitação no 2.

    O não cumprimento desse requisito resultará em erros como: Batching input tensors supplied in a given op invocation must have equal 0th-dimension size.

  4. O tamanho 0 da forma de cada tensor de saída precisa ser a soma do tamanho de 0 dos tensores de entrada, mais qualquer padding introduzido pela operação em lote para atender à próxima maior allowed_batch_size. Isso permite que a operação em lote divida os tensores de saída ao longo da 0a dimensão com base na 0a dimensão dos tensores de entrada.

    No tutorial acima, os tensores de entrada têm uma 0a dimensão de 1 da solicitação 1 e 2 da solicitação no 2. Portanto, cada tensor de saída precisa ter uma 0a dimensão de 3, porque 1+2=3. O tensor de saída (3, 3, 4) atende a esse requisito. Se 3 não fosse um tamanho de lote válido, mas 4 fosse, a operação de lote precisaria preencher a dimensão 0 das entradas de 3 a 4. Nesse caso, cada tensor de saída teria que ter um tamanho de dimensão 0 de 4.

    O não cumprimento desse requisito resultará em um erro como: Batched output tensor's 0th dimension does not equal the sum of the 0th dimension sizes of the input tensors.

Como resolver erros de requisitos de forma

Para atender a esses requisitos, forneça uma função ou assinatura diferente para lotes. Também pode ser necessário modificar as funções atuais para atender a esses requisitos.

Se uma função estiver sendo loteada, verifique se todas as formas de input_signature da @tf.function têm None na 0a dimensão (também conhecida como dimensão do lote). Se uma assinatura estiver sendo loteada, verifique se todas as entradas dela têm -1 na dimensão 0.

A op BatchFunction não é compatível com SparseTensors como entradas ou saídas. Internamente, cada tensor esparso é representado como três tensores separados que podem ter diferentes tamanhos de 0a dimensão.