Introdução ao conversor de inferência do Cloud TPU v5e

Introdução

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

  1. Conversão de TPU: adiciona TPUPartitionedCall e outras operações de TPU ao para torná-lo veiculável na TPU. Por padrão, um modelo exportado para a inferência não tem essas operações e não pode ser exibida na TPU, mesmo foi treinado na TPU.
  2. Lotes: adiciona operações em lote ao modelo para ativar o agrupamento em lote no gráfico. para melhorar a capacidade de processamento.
  3. Conversão BFloat16: converte o formato de dados do modelo dos float32 a bfloat16 para melhorar desempenho computacional e menos memória com alta largura de banda (HBM) na TPU.
  4. Otimização de forma de E/S: otimiza as formas de 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 todas as funções que gostariam de executar na TPU. Eles passam essas funções para o conversor e o O conversor coloca as chaves na TPU e as otimiza.

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 a 30 minutos

Pré-requisitos

  1. O modelo precisa ser do TF2 e exportado no formato SavedModel.
  2. O modelo precisa ter um alias de função para a função da TPU. Consulte a exemplo de código para saber como fazer isso. Os exemplos a seguir usam tpu_func como o alias da função TPU.
  3. Certifique-se de que a CPU do computador seja compatível com Advanced Vector eXtensions (AVX) da biblioteca do TensorFlow, que é a dependência do Cloud TPU Conversor de inferência) é compilado para usar as instruções do AVX. A maioria das CPUs seja compatível com AVX.
    1. Execute lscpu | grep avx para verificar se o AVX é suportado.

Antes de começar

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

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

  • Configure uma VM da TPU: Crie uma nova VM de TPU usando o console do Google Cloud ou gcloud, ou use uma VM de TPU existente para executar inferência com a variável na VM da TPU.

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

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

    Conecte-se à instância com SSH 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 comandos a seguir para permitir o uso não raiz do Docker:

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

  • Inicialize seus 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:

  • Permita o uso não raiz do Docker:

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

  • Faça login no seu ambiente:

    gcloud auth login

  • Inicialize seus 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 de conversão

A imagem é para fazer conversões de modelo únicas. Definir os caminhos do modelo e ajustar opções do conversor para atender às suas necessidades. O Exemplos de uso seção apresenta 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 da 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 TPU

  1. Encontre ou crie uma função no modelo que agrupe 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 dar model.tpu_func um alias func_on_tpu.
  3. Você pode passar 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 de TPU

É possível colocar várias funções na TPU. Basta criar várias funções e transmita-os por converter_options_string ao 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 usada para representar os parâmetros de um modelo. Isso resulta em um modelo menor e computação mais rápida. Um modelo quantizado fornece ganhos capacidade de inferência, bem como menor uso de memória e tamanho de armazenamento, em detrimento de pequenas quedas na precisão.

O novo recurso de quantização pós-treinamento do TensorFlow destinada a a TPU, é desenvolvida a partir do recurso similar existente no TensorFlow Lite. usado para segmentar dispositivos móveis e de borda. Para saber mais sobre quantização em geral, confira Documento do TensorFlow Lite.

Conceitos de quantização

Nesta seção, definimos conceitos especificamente relacionados à 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 no Arquitetura do sistema de TPUs.

  • 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 o que afeta significativamente sua precisão. A PTQ funciona convertendo pesos de ponto flutuante e ativações de um modelo treinado para números inteiros de menor precisão, 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, com apenas uma pequena perda de precisão.

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

  • Conjunto de dados representativo: um conjunto de dados representativo para quantização é um pequeno conjunto de dados que representa os dados de entrada reais do modelo. É usado durante a etapa de calibração da quantização para coletar estatísticas sobre o intervalo de valores dos quais os pesos e as ativações o desempenho do modelo. O conjunto de dados representativo precisa atender propriedades a seguir:

    • Ele precisa representar corretamente as entradas reais para o modelo a inferência. Isso significa que ele deve abranger o intervalo de valores que que o modelo provavelmente verá no mundo real.
    • Elas precisam passar coletivamente por cada ramificação de condicionais (como tf.cond), se houver algum. Isso é importante porque o processo de quantização precisa ser capaz de lidar com todas as entradas possíveis ao modelo, mesmo que não sejam representados explicitamente conjunto de dados representativo.
    • Ele deve ser grande o suficiente para coletar estatísticas suficientes e reduzir erro. Como regra geral, recomendamos usar mais de 200 amostras representativas.

    O conjunto de dados representativo pode ser um subconjunto do conjunto de dados de treinamento ou um conjunto de dados separado projetado especificamente para representar as entradas do mundo real no modelo. A escolha do 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. os pesos e as ativações de um modelo de rede neural uma vez, durante o etapa de calibragem. Isso significa que o mesmo intervalo de valores é usado para todas as entradas do modelo. Isso pode ser menos preciso do que o intervalo dinâmico quantização específica, especialmente para modelos com uma ampla gama de valores de entrada. No entanto, a quantização de intervalo estático exige menos computação na 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 valores para os pesos e ativações de um modelo de rede neural para cada entrada. Isso permite que o modelo se adapte ao intervalo de valores de os dados de entrada, o que pode melhorar a precisão. No entanto, a quantização de intervalo dinâmico exige mais computação no momento da 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 calibragem Determinado para cada entrada
    Precisão Pode ser menos preciso, especialmente para modelos com uma ampla gama de valores de entrada Pode ser mais preciso, especialmente para modelos com um amplo intervalo de valores de entrada.
    Complexidade Mais simples Mais complexo
    Computação no ambiente de execução Menos computação Mais computação

  • Quantização somente de peso: a quantização somente de peso é um tipo de que só quantiza os pesos de um modelo de rede neural, enquanto as ativações ficam em ponto flutuante. Essa pode ser uma boa opção para modelos sensíveis à acurácia, já que pode ajudar preservar a acurácia do modelo.

Como usar a quantização

A quantização pode ser aplicada configurando QuantizationOptions às opções do conversor. As principais opções são:

  • tags: coleção de tags que identificam o MetaGraphDef em a SavedModel para quantizar. Não é necessário especificar se você tem apenas um MetaGraphDef.
  • assinatura_chaves: sequência de chaves que identificam 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. Caso contrário especificado, a quantização STATIC_RANGE vai ser aplicada.
  • op_set: precisa ser mantido como XLA. Atualmente, essa é a opção padrão, então não é preciso especificar.
  • representante_datasets: especifica o conjunto de dados usado para calibrar os parâmetros de quantização.

Como criar o conjunto de dados representativo

Um conjunto de dados representativo é essencialmente um iterável de amostras. Em que 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 TFRecord. arquivos usando a classe TfRecordRepresentativeDatasetSaver disponível no pacote tf-nightly do 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})

Exemplos

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

Adicionar lotes

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

Por padrão, o conversor agrupa todas as funções de TPU em lote no modelo. Ela também pode lote fornecido pelo usuário assinaturas e Funções o que pode melhorar ainda mais o desempenho. Qualquer função da TPU, função fornecida pelo usuário ou uma assinatura em lote, precisa atender aos requisitos requisitos rígidos de formato.

O usuário que fez a conversão também pode atualizar opções de lote atuais. Veja a seguir um exemplo de como adicionar lotes a um modelo. Para mais informações sobre lotes, consulte Análise detalhada em lote.

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 formato bfloat16 e E/S

O BFloat16 e as otimizações de forma de E/S são ativados por padrão. Se não funcionarem funcionam bem com seu modelo, elas podem ser desativadas.

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

Relatório de conversão

Esse relatório de conversão pode ser encontrado no registro depois de executar a inferência Conversor. 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
--------------------------------

Esse relatório estima o custo computacional do modelo de saída na CPU e no TPU e detalha o custo do TPU para cada função, o que deve refletir sua seleção das funções do TPU nas opções do conversor.

Se você quiser utilizar melhor a TPU, convém testar o modelo e ajustar as opções do conversor.

Perguntas frequentes

Quais funções devo colocar na TPU?

É melhor colocar o máximo possível do seu modelo na TPU, porque a a maioria das operações é executada mais rapidamente na TPU.

Se o modelo não tiver operações, strings ou operações esparsas incompatíveis com TPU Portanto, colocar todo o modelo na TPU geralmente é a melhor estratégia. E faça isso encontrando ou criando uma função que envolva todo o modelo, criar um alias de função para ela e passá-la para o conversor.

Se o modelo tiver peças que não funcionam na TPU (por exemplo, peças incompatíveis com TPU ops, strings ou tensores esparsos), a escolha das funções da TPU depende de onde é a parte incompatível.

  • Se for no início ou no fim do modelo, é possível refatorar para mantê-lo na CPU. Exemplos são pré e pós-processamento de strings fases. Para mais informações sobre como mover códigos para a CPU, consulte "Como faço para mover uma parte do modelo para CPU?" Ele mostra uma maneira típica de refatorar o modelo.
  • Se estiver no meio do modelo, é melhor dividir o modelo em três partes e conter todas as operações incompatíveis com TPU na parte do meio e executá-lo na CPU.
  • Se for um tensor esparso, chame tf.sparse.to_dense nos CPU e passando o tensor denso resultante para a parte da TPU do modelo.

Outro fator a considerar é o uso do HBM. A incorporação de tabelas pode usar muita 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, deve existir apenas uma função de TPU em uma assinatura. Se a estrutura do modelo exigir a chamada de várias funções de TPU por pedido de inferência, você precisará considerar a latência adicional do envio de tensores entre a CPU e a TPU.

Uma boa maneira de avaliar a seleção de funções de TPU é verificar Relatório de conversão. Ele mostra a porcentagem de computação que foi colocada na TPU e um detalhe do custo de cada função da TPU.

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

Se o modelo tiver peças que não podem ser veiculadas na TPU, será preciso refatorar o modelo para movê-los para a CPU. Aqui está um exemplo de brinquedo. O modelo é um modelo de linguagem com um estágio de pré-processamento. O código das definições de camadas e são omitidas para simplificar.

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 disponibilizado diretamente na TPU por dois motivos. Primeiro, os é uma string. Segundo, a função preprocess pode conter muitas strings ops. Os dois 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 tpu_func será executado na TPU, e tudo o restante em model_func será executado em e a 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 são compatíveis com a TPU, mas algumas incluindo tensores e strings esparsos não são aceitos. O conversor não verifica operações incompatíveis com a TPU. Assim, um modelo com essas operações pode passar o e conversão em massa. No entanto, ao executá-lo para inferência, ocorrerão erros como o mostrado abaixo.

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

Se o modelo tiver operações incompatíveis com TPU, elas deverão ser colocadas fora da TPU função. Além disso, a string é um formato de dados incompatível com a TPU. Então, variáveis do tipo string não devem ser colocadas na função da TPU. E o parâmetros e valores de retorno da função da TPU não devem ser do tipo string, como muito bem. Da mesma forma, evite colocar tensores esparsas na função TPU, incluindo os parâmetros e valores de retorno.

Geralmente não é difícil refatorar a parte incompatível do modelo e movê-lo para a CPU. Confira um exemplo.

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

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

Como o código do conversor ainda não tem código aberto, ele não pode ser criados com operações personalizadas.

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

O conversor não oferece suporte aos modelos do TensorFlow 1. Os modelos do TensorFlow 1 precisam serão 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 TF2XLA ou a ponte TF2XLA original.

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. Se a reexportação não for uma opção, ainda será possível converter o modelo fornecendo um concrete_function_name. No entanto, identificar O concrete_function_name exige trabalho de detetive.

Aliases de função são um mapeamento de uma string definida pelo usuário para uma função concreta nome. Eles facilitam a consulta a uma função específica no modelo. O O conversor aceita aliases de função e nomes de funções concretos brutos.

Nomes de funções concretos 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 constante no tempo de compilação?

Tanto para treinamento quanto para inferência, o XLA exige que as entradas de determinadas operações tenham conhecida no momento da compilação da TPU. Isso significa que, quando o XLA compila a TPU do programa, as entradas dessas operações precisam ter uma conexão forma

Há duas maneiras de resolver esse problema.

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

Por que estou recebendo um erro de formas em lote?

O agrupamento tem requisitos de forma rígidos que permitem que as solicitações recebidas sejam agrupadas na dimensão 0 (também conhecida como dimensão de agrupamento). Esses requisitos de formato vêm do fluxo de trabalho em lote do TensorFlow e não pode ser relaxado.

O não cumprimento desses requisitos resultará em erros como:

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

Para atender a esses requisitos, considere fornecer uma função ou assinatura aos lotes. Também pode ser necessário modificar as funções existentes para atender a essas e cumprimento de requisitos regulatórios.

Se uma função está sendo agrupado, verifique se as formas input_signature da assinatura de @tf.function ter "None" na dimensão 0. Se uma assinatura está sendo agrupado, verifique se todas as entradas têm -1 na dimensão 0.

Para uma explicação completa sobre por que esses erros estão acontecendo e como resolver eles, consulte Análise detalhada de lotes.

Problemas conhecidos

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

Embora o conversor possa lidar com a maioria dos cenários de chamada de função em todo o limite CPU-TPU, há um caso de borda raro em que ele falharia. É quando uma função TPU chama indiretamente outra função TPU.

Isso ocorre porque o conversor modifica o autor da chamada direta de uma função da TPU do 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 chama o autor da chamada direta, essas operações da CPU podem ser trazidas para a TPU para execução, o que gera erros de kernel ausentes. Observe este 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.

No conversor, implementamos a detecção desse cenário. Se você vir o erro a seguir significa que o 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 que essa função chame diferente. Se isso for difícil, entre em contato com o Suporte do Google equipe para discutir mais.

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 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álise detalhada sobre lotes

A criação de lotes é usada para melhorar a capacidade de processamento e a utilização da TPU. Permite várias solicitações sejam processadas ao mesmo tempo. Durante o treinamento, o agrupamento pode ser feito usando tf.data. Durante a inferência, isso normalmente é feito adicionando um no gráfico que agrupa as solicitações recebidas em lote. A operação espera até ter o suficiente solicitações ou quando o tempo limite for atingido antes de gerar um grande lote da solicitações individuais. Consulte Definição de opções de lotes para mais informações sobre as diferentes opções de lotes que podem ser ajustadas, incluindo tamanhos de lote e tempos limite.

lote no gráfico

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

O exemplo a seguir mostra como adicionar o lote 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
}

Lote de assinaturas

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

Assim como em todos os lotes, os lotes de assinaturas têm requisitos rígidos de formato. Para ajudar a garantir que esses requisitos de formato sejam atendidos, as entradas de assinatura devem ter formas que têm pelo menos duas dimensões. A primeira dimensão é o tamanho do lote, e devem ter o 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, considere usar o comportamento padrão de agrupamento ou agrupamento de funções.

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

Agrupamento em lote de funções

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

É possível usar o agrupamento de funções para fazer a computação da CPU em lote. Muitos modelos apresentam uma melhoria no desempenho quando a computação da CPU é agrupada. A melhor maneira de A computação de CPU em lote está usando lotes de assinaturas, mas pode não funcionar para alguns de modelos de machine learning. Nesses casos, o agrupamento de funções pode ser usado para agrupar parte da CPU computação em nuvem, além da TPU. A operação de lote não pode ser executados na TPU. Portanto, qualquer função de lote fornecida precisa ser chamada no CPU.

O agrupamento de funções também pode ser usado para satisfazer requisitos rígidos de formato imposto pela operação de lote. Nos casos em que as funções da TPU não atendem aos requisitos de agrupar os requisitos de formato das operações, o agrupamento de funções pode ser usado para informar Conversor para agrupar funções diferentes.

Para usar isso, gere uma function_alias para a função que precisa ser processada em lote. Para isso, encontre ou crie uma função no modelo que une tudo o que você quer agrupar. Certifique-se de que essa função atenda requisitos rígidos de formato imposto pela operação de lote. Adicione @tf.function caso ainda não tenha um. É importante fornecer o input_signature ao @tf.function. O dia 0 a dimensão deve ser None porque é a dimensão de lote. Portanto, não pode ser um tamanho fixo. Por exemplo, [None, 4], [None] ou [None, 128, 4, 10] são todos 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, transmita 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 as function_alias usando o 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 lotes

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

Como atualizar opções de lotes atuais

É possível adicionar ou atualizar opções de lote executando a imagem Docker especificando batch_options e definindo disable_default_optimizations como verdadeiro usando sinalização --converter_options_string. As opções de lote serão aplicadas a cada Função da TPU ou operação de lote pré-existente.

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

Lotes são criados pela concatenação de tensores de entrada entre solicitações junto ao lote (0a). Os tensores de saída são divididos ao longo da dimensão 0. Para para realizar essas operações, a operação de lotes tem requisitos rígidos de formato para as entradas e saídas.

Tutorial

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

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 o entradas A e B com as formas (2, 3, 2) e (2, 2, 4).

solicitação de inferência 1

O tempo limite do lote foi atingido. O modelo aceita um tamanho de lote de 3, as solicitações de inferência 1 e 2 são agrupadas sem preenchimento. O os tensores em lote são formados pela concatenação das solicitações 1 e 2 no lote (0th) dimensão. Como o A de #1 tem a forma (1, 3, 2) e o A de #2 tem a forma de (2, 3, 2), quando eles são concatenados ao longo da dimensão do lote (0), a forma resultante é (3, 3, 2).

solicitação em lote

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

solicitação de matriz em lote

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

resultados da solicitação de inferência

Requisitos de forma

Para realizar a concatenação de entradas e a divisão de saída descritas acima, a operação de lote tem os seguintes requisitos de forma:

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

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

    O não cumprimento desse requisito causará um erro como: Batching input tensors must have at least one dimension. Uma correção simples para esse erro é transforme o escalar em um vetor.

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

    No tutorial acima, a solicitação A da solicitação no 1 tem o formato (1, 3, 2). Isso significa que qualquer solicitação futura deve 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 o formato (1, 2, 4), então todas as solicitações futuras precisam produzir uma forma com o padrão (X, 2, 4).

    O não cumprimento desse requisito causará um erro como: Dimensions of inputs should match.

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

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

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

  4. O tamanho da dimensão 0 do formato de cada tensor de saída precisa ser a soma: de todos os tensores de entrada, Tamanho da dimensão 0 (mais qualquer padding introduzido pelo a operação de lote para atender ao próximo maior allowed_batch_size). Isso permite a operação de lote para dividir os tensores de saída ao longo da dimensão 0 com base na dimensão 0 dos tensores de entrada.

    No tutorial acima, os tensors de entrada têm uma dimensão 0 de 1 da solicitação 1 e 2 da solicitação 2. Portanto, cada tensor de saída têm uma dimensão 0 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 teria que preencher a dimensão 0 das entradas de 3 para 4. Nesse caso, cada tensor de saída precisa 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, considere fornecer uma função ou assinatura aos lotes. Também pode ser necessário modificar as funções existentes para atender a essas e cumprimento de requisitos regulatórios.

Se um função está sendo agrupado, verifique se as formas input_signature de @tf.function estão têm None na dimensão 0 (também conhecida como dimensão de lote). Se um assinatura está sendo agrupado, verifique se todas as entradas têm -1 na dimensão 0.

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