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

Introdução

O conversor de inferência da 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 da TPU é recomendado porque vem pré-instalado com o comando as ferramentas de linha 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 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. Agrupamento: adiciona operações de agrupamento ao modelo para permitir o agrupamento no gráfico e melhorar a taxa de transferência.
  3. Conversão de BFloat16: converte o formato de dados do modelo de float32 para bfloat16 para melhorar o desempenho computacional e reduzir o uso de memória de alta largura de banda (HBM, na sigla em inglês) 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 acima: ~20 a 30 minutos

Pré-requisitos

  1. O modelo precisa ser do TF2 e exportado no SavedModel .
  2. O modelo precisa ter um alias de função para a função 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 tem suporte a AVX.
    1. Execute lscpu | grep avx para verificar se o AVX é suportado.

Antes de começar

Antes de começar a configuração, faça o seguinte:

  • Criar um projeto: no console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do 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 disponibilizado nessa VM do 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. Use 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 no shell da VM do TPU ou no 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 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:

  • 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 os auxiliares de credenciais do Docker:

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

  • Extraia a imagem do Docker do Inference Converter:

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

Converter imagem

A imagem é para fazer conversões de modelo únicas. Defina os caminhos do modelo e ajuste as opções do conversor de acordo com suas necessidades. A seção Exemplos de uso 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 para a função da 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 as SaveOptions 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 oferece ganhos na taxa de inferência, além de menor uso de memória e tamanho de armazenamento, à custa 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 que já existe 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

Esta seção define conceitos relacionados especificamente à quantização com o Inference Converter.

Conceitos relacionados a outras configurações de TPU (por exemplo, fatias, 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): a PTQ é uma técnica que reduz o tamanho e a complexidade computacional de um modelo de rede neural sem afetar significativamente a precisão dele. 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, o que gera uma pequena perda de acurácia.

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

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

    • Ele deve representar adequadamente as entradas reais para o modelo durante 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 processar todas as entradas possíveis para o modelo, mesmo que elas não sejam representadas explicitamente no conjunto de dados representativo.
    • Ele deve ser grande o suficiente para coletar estatísticas suficientes e reduzir erro. Como regra geral, é recomendável 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): a 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 todos entradas para o 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 requer menos computação no momento 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 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 dos 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 alcance dinâmico
    Intervalo de valores Determinado uma vez, durante a calibração Determinado para cada entrada
    Precisão Pode ser menos preciso, especialmente para modelos com uma grande variedade 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 à precisão, já que ajuda a preservar a precisão do modelo.

Como usar a quantização

A quantização pode ser aplicada configurando e definindo QuantizationOptions nas opções do conversor. As opções mais importantes são:

  • tags: coleção de tags que identificam o MetaGraphDef dentro do SavedModel para quantização. 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. Se não for especificado, a quantização STATIC_RANGE 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 arquivos TFRecord usando a classe TfRecordRepresentativeDatasetSaver 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})

Exemplos

O exemplo a seguir quantiza o modelo com a chave de assinatura de serving_default e 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 das opções de lote que podem ser ajustadas, consulte Definição das opções de lote.

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 Conversor também pode atualizar as opções de lote atuais. Veja a seguir um exemplo de como adicionar lotes a um modelo. Para mais informações sobre o agrupamento, consulte Aprofundamento sobre o agrupamento.

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 as otimizações de forma bfloat16 e IO

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 o TPU, experimente a estrutura do modelo e ajuste as opções do conversor.

Perguntas frequentes

Quais funções devo colocar na TPU?

É melhor colocar o máximo possível do modelo na TPU, porque a grande 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 isso pode ser feito encontrando ou criando uma função que envolva todo o modelo, criando um alias de função para ele e transmitindo-o ao 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 etapas. 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 da HBM. Incorporar tabelas embutidas pode usar muitos da HBM. Se eles crescerem além da limitação de hardware da TPU, precisarão ser colocados na CPU, junto com as operações de pesquisa.

Sempre que possível, apenas uma função da TPU deve existir em uma assinatura. Se a estrutura do seu modelo exigir a chamada de várias funções de TPU por uma solicitação de inferência recebida, você deve estar ciente da latência adicional de 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. Confira um exemplo de brinquedo. O modelo é um modelo de linguagem com uma fase 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 veiculado diretamente no TPU por dois motivos. Primeiro, o parâmetro é uma string. Segundo, a função preprocess pode conter muitas strings ops. Nenhum dos dois é compatível com TPU.

Para refatorar esse modelo, crie outra função chamada tpu_func para hospedar o bert_layer, que é de computação intensa. Em seguida, crie um alias de função para tpu_func e transmita-o para o 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 tensors esparsas incompatíveis com o TPU?

A maioria das operações padrão do TensorFlow é compatível 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 que contém essas operações pode transmitir a conversão. 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. Portanto, as variáveis do tipo string não podem ser colocadas na função 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 esparsos na função da TPU, inclusive em parâmetros e valores de retorno.

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

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

Se as 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 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 é de 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 é compatível com 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.

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 nomes de função concretos e nomes de função concretos brutos.

Os nomes de função 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?

Para treinamento e inferência, a XLA exige que as entradas de determinadas operações tenham uma forma conhecida no momento da compilação da TPU. Isso significa que, quando o XLA compila a TPU do programa, as entradas dessas operações devem ter uma camada conhecida estaticamente forma

Há duas maneiras de resolver esse problema.

  • A melhor opção é atualizar as entradas da operação para ter uma forma conhecida estática até que o XLA compile o programa TPU. Essa compilação acontece logo antes que a parte da TPU do modelo seja executada. Isso significa que a forma precisa ser conhecida de forma estática 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 forma em lote?

O lote tem requisitos rígidos de formato que permitem que as solicitações recebidas sejam agrupadas ao longo da dimensão 0 (também conhecida como de lote). 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 iguais.
  3. O agrupamento em lote de tensores de entrada fornecidos em uma determinada invocação de operação precisa ter tamanho da dimensão 0.
  4. A dimensão 0 do tensor de saída em lote não é igual à soma dos tamanhos da dimensão 0 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 estiver sendo agrupada, verifique se todas as formas de input_signature da @tf.function têm 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 resolvê-los, consulte Aprofundamento sobre lotes.

Problemas conhecidos

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

Embora o Conversor possa processar a maioria dos cenários de chamada de função em todo o limite da CPU-TPU, há um caso raro de falha. É 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ó podem funcionar 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 que 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.

No conversor, implementamos a detecção desse cenário. Se você encontrar o erro a seguir, significa que seu modelo atingiu este 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 a equipe de suporte do Google para discutir o assunto.

Referência

Opções do 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

O lote é usado para melhorar a capacidade e a utilização da TPU. Isso permite que várias solicitações sejam processadas ao mesmo tempo. Durante o treinamento, lotes pode ser feito usando tf.data. Durante a inferência, isso geralmente é feito adicionando uma operação no gráfico que agrupa as solicitações recebidas. A operação vai aguardar até receber solicitações suficientes ou até que um tempo limite seja atingido antes de gerar um lote grande das solicitações individuais. Consulte Definição de opções de lote para mais informações sobre as diferentes opções de lote que podem ser ajustadas, incluindo tamanhos e tempos limite de lote.

lote no gráfico

Por padrão, o conversor insere a operação de lote diretamente antes da TPU de computação. Ele envolve 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 essa padrão informando ao usuário que fez uma conversão funções e/ou assinaturas devem ser reunidos em lotes.

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
}

Agrupamento de assinaturas

Os lotes de assinaturas agrupam todo o modelo, começando pelas entradas da assinatura e acessar as saídas da assinatura. Ao contrário do comportamento padrão de agrupamento do Converter, o agrupamento de assinaturas agrupa a computação da TPU e da CPU. Isso gera um ganho de 10% a 20% no desempenho durante a inferência em alguns modelos.

Como todos os lotes, o lote de assinaturas tem requisitos de formato rígidos. Para garantir que esses requisitos sejam atendidos, as entradas de assinatura precisam ter formas com pelo menos duas dimensões. A primeira é 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 de lote padrão ou lote de funções em lote.

Para usar o agrupamento de assinaturas, informe 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"
  }
}

Função em lote

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 de TPU. O agrupamento de funções substitui esse comportamento padrão.

O lote de funções pode ser usado para agrupar a computação da CPU. Muitos modelos têm uma de melhoria de desempenho quando a computação da CPU é agrupada em lote. A melhor maneira de processar a CPU em lote é usando a assinatura em lote, mas ela pode não funcionar para alguns modelos. 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 executada na TPU. Portanto, qualquer função de lote fornecida deve ser chamada no CPU.

O agrupamento de funções também pode ser usado para atender aos requisitos de formato estrito imposicionados pela operação de agrupamento. Nos casos em que as funções de TPU não atendem aos requisitos de formato da operação de agrupamento, o agrupamento de funções pode ser usado para informar ao conversor para agrupar diferentes funções.

Para usar isso, gere um function_alias para a função que será em lotes. 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 das opções de lote

  • 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 paralelo. Esse valor precisa ser aproximadamente igual ao número de núcleos de TPU disponíveis.
  • max_batch_size: (número inteiro) tamanho máximo do lote permitido. Pode ser maior do que allowed_batch_sizes para usar a divisão de lotes grandes.
  • 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 precisa ser crescente 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 opções de lotes atuais

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

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

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

solicitação de inferência 1

O tempo limite de lotes foi atingido. O modelo oferece suporte a um tamanho de lote de 3, para que as solicitações de inferência 1 e 2 sejam 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 uma forma de (1, 3, 2) e o A de #2 tem uma 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 de agrupamento faz isso dividindo a dimensão do lote (0) 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 uma dimensão 0 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 entrada 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 da dimensão 0, os tensores precisam ter pelo menos duas dimensões.

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

    Se esse requisito não for atendido, vai ocorrer um erro como: Batching input tensors must have at least one dimension. Uma correção simples para esse erro é transformar 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 1 tem uma forma de (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 uma forma de (1, 2, 4), portanto, 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 lote tiverem uma dimensão 0 diferente, a operação de lote não saberá como dividir de entrada e 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 precisa ter um tamanho de dimensão 0 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 o a operação de lote divide a forma final de (3, 3, 4) e produz (1, 3, 4). para a solicitação 1 e (2, 3, 4) para a solicitação 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 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 encontrar o 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 tensores de entrada têm uma dimensão 0 de 1. nas solicitações 1 e 2 da solicitação 2. Portanto, cada tensor de saída precisa ter 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, o a operação de lote teria que preencher a dimensão 0 das entradas de 3 a 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 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 aceita 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.