Esta página foi traduzida pela API Cloud Translation.

Introdução à inferência da Cloud TPU

Informações gerais e benefícios

A veiculação se refere ao processo de implantação de um modelo de aprendizado de máquina treinado em um ambiente de produção, onde ele pode ser usado para inferência. A inferência é compatível com a TPU v5e e versões mais recentes. Os SLOs de latência são uma prioridade para a veiculação.

Este documento discute a disponibilização de um modelo em uma TPU de host único. As frações de TPU com 8 chips ou menos têm uma VM ou um host de TPU e são chamadas de TPUs de host único.

Primeiros passos

Você vai precisar de uma conta e um projeto do Google Cloud para usar o Cloud TPU. Para mais informações, consulte Configurar um ambiente do Cloud TPU.

É necessário solicitar a cota a seguir para veicular em TPUs:

Recursos do v5e sob demanda: TPUv5 lite pod cores for serving per project per zone
Recursos preemptivos v5e: Preemptible TPU v5 lite pod cores for serving per project per zone
Recursos v6e sob demanda: TPUv6 cores per project per zone
Recursos v6e preemptivos: Preemptible TPUv6 cores per project per zone

Para mais informações sobre a cota de TPU, consulte Cota de TPU.

Inferência e exibição de modelos do Cloud TPU

A maneira de veicular um modelo para inferência depende do framework de ML usado para criar o modelo. O Cloud TPU oferece suporte à veiculação de modelos criados em JAX, TensorFlow e PyTorch.

Inferência e veiculação de modelos JAX

Para disponibilizar um modelo em uma VM da TPU, você precisa:

Serializar seu modelo no formato SavedModel do TensorFlow
Use o Converter de inferência para preparar o modelo salvo para veiculação
Usar o TensorFlow Serving para mostrar o modelo

Formato SavedModel

Um SavedModel contém um programa completo do TensorFlow, incluindo parâmetros treinados e computação. Ele não exige que o código de criação do modelo original seja executado.

Se o modelo foi escrito no JAX, use jax2tf para serializar o modelo no formato SavedModel.

Conversor de inferência

O conversor de inferência da Cloud TPU prepara e otimiza um modelo exportado no formato SavedModel para inferência de TPU. É possível executar o conversor de inferência em um shell local ou na VM de TPU. Recomendamos o uso do shell da VM TPU, porque ele tem todas as ferramentas de linha de comando necessárias para executar o conversor. Para mais informações sobre o Inference Converter, consulte o Guia do usuário do Inference Converter.

Requisitos do Inference Converter

O modelo precisa ser exportado do TensorFlow ou do JAX no formato SavedModel.
É necessário definir um alias de função para a função TPU. Para mais informações, consulte o Guia do usuário do Inference Converter. Os exemplos neste guia usam tpu_func como o alias da função TPU.
Verifique se a CPU da máquina oferece suporte às instruções Advanced Vector eXtensions (AVX), já que a biblioteca do TensorFlow (a dependência do Conversor de inferência do Cloud TPU) é compilada para usar instruções AVX. A maioria das CPUs tem suporte para AVX.

Observação: execute lscpu | grep avx para verificar se o conjunto de instruções AVX tem suporte.

Inferência e veiculação de modelos JAX

Esta seção descreve como disponibilizar modelos do JAX usando jax2tf e TensorFlow Serving.

Use jax2tf para serializar seu modelo no formato SavedModel.
Use o Converter de inferência para preparar o modelo salvo para exibição
Usar o TensorFlow Serving para mostrar o modelo

Use `jax2tf` para serializar um modelo JAX no formato SavedModel

A função Python a seguir mostra como usar jax2tf no código do modelo:

# Inference function
def model_jax(params, inputs):
  return params[0] + params[1] * inputs

# Wrap the parameter constants as tf.Variables; this will signal to the model
# saving code to save those constants as variables, separate from the
# computation graph.
params_vars = tf.nest.map_structure(tf.Variable, params)

# Build the prediction function by closing over the `params_vars`. If you
# instead were to close over `params` your SavedModel would have no variables
# and the parameters will be included in the function graph.
prediction_tf = lambda inputs: jax2tf.convert(model_jax)(params_vars, inputs)

my_model = tf.Module()
# Tell the model saver what the variables are.
my_model._variables = tf.nest.flatten(params_vars)
my_model.f = tf.function(prediction_tf, jit_compile=True, autograph=False)
tf.saved_model.save(my_model)

Para mais informações sobre jax2tf, consulte Interoperação do JAX e do Cloud TPU.

Use o Converter de inferência para preparar o modelo salvo para veiculação

As instruções para usar o Inference Converter estão descritas no guia do Inference Converter.

Usar o TensorFlow Serving

As instruções para usar o TensorFlow Serving estão descritas em Serviço do TensorFlow.

Exemplos de disponibilização de modelos JAX

Pré-requisitos

Configure suas credenciais do Docker e extraia a imagem do Docker do Converter de inferência e do serviço do Cloud TPU:

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
   us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Conecte-se à VM do TPU com SSH e instale o código de demonstração de inferência:

gcloud storage cp \
"gs://cloud-tpu-inference-public/demo" \
. \
--recursive

Instale as dependências da demonstração do JAX:
```
pip install -r ./demo/jax/requirements.txt
```

Exibir o modelo BERT JAX para inferência

Você pode fazer o download do modelo BERT pré-treinado no Hugging Face.

Exporte um modelo salvo do TensorFlow compatível com TPU de um modelo BERT do Flax:
```
cd demo/jax/bert
python3 export_bert_model.py
```

Inicie o contêiner do servidor de modelos do Cloud TPU:

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/jax/bert_tpu,target=/models/bert \
  -e MODEL_NAME=bert \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Cerca de 30 segundos após o início do contêiner, verifique o registro do contêiner do servidor de modelo e verifique se os servidores gRPC e HTTP estão ativos:

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

Se você encontrar uma entrada de registro que termina com as informações a seguir, o servidor estará pronto para atender às solicitações.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

Envie uma solicitação de inferência para o servidor do modelo.

python3 bert_request.py

A saída será semelhante a esta:

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."

Fazer a limpeza.

Limpe o contêiner do Docker antes de executar outras demonstrações.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

Limpe os artefatos do modelo:

sudo rm -rf /tmp/jax/

Servir a difusão estável do JAX para inferência

Faça o download do modelo pré-treinado da Stable Diffusion no Hugging Face.

Faça o download do modelo de difusão estável em um formato de modelo salvo do TF2 compatível com TPU:

Observação: o download leva cerca de três minutos.
```
cd demo/jax/stable_diffusion
python3 export_stable_diffusion_model.py
```

Inicie o contêiner do servidor de modelo do Cloud TPU para o modelo:

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/jax/stable_diffusion_tpu,target=/models/stable_diffusion \
  -e MODEL_NAME=stable_diffusion \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Após cerca de dois minutos, verifique o registro do contêiner do servidor de modelo para garantir que os servidores gRPC e HTTP estejam em execução:

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

Se o registro terminar com as informações a seguir, significa que os servidores estão prontos para atender às solicitações.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

Envie uma solicitação para o servidor do modelo.
```
python3 stable_diffusion_request.py
```
Este script envia "Pintura de um esquilo patinando em Nova York" como instrução. A imagem de saída será salva como stable_diffusion_images.jpg no seu diretório atual.

Fazer a limpeza.

Limpe o contêiner do Docker antes de executar outras demonstrações.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

Limpar os artefatos do modelo

sudo rm -rf /tmp/jax/

TensorFlow Serving

As instruções a seguir demonstram como você pode oferecer seu modelo do TensorFlow em VMs TPU.

Fluxo de trabalho do TensorFlow Serving

Faça o download da imagem do Docker do TensorFlow Serving para sua VM de TPU.

Definir variáveis de ambiente de exemplo

export YOUR_LOCAL_MODEL_PATH=model-path
export MODEL_NAME=model-name
# Note: this image name may change later.
export IMAGE_NAME=us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Fazer o download da imagem do Docker

docker pull ${IMAGE_NAME}

Configure as credenciais do Docker e extraia o Inference Converter e a imagem do Docker do TensorFlow Serving.

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
   us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Faça o download do código de demonstração:

gcloud storage cp \
"gs://cloud-tpu-inference-public/demo" \
. \
--recursive

Instale as dependências da demonstração do TensorFlow:
```
pip install -r ./demo/tf/requirements.txt
```

Exiba o modelo do TensorFlow usando a imagem do Docker do TensorFlow Serving na VM da TPU.

# PORT 8500 is for gRPC model server and 8501 is for HTTP/REST model server.
docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=${YOUR_LOCAL_MODEL_PATH},target=/models/${MODEL_NAME} \
  -e MODEL_NAME=${MODEL_NAME} \
  ${IMAGE_NAME}

Use a API Serving Client para consultar seu modelo.
- API REST Client
- API gRPC do cliente

Executar a demonstração do TensorFlow ResNet-50 Serving

Exporte um modelo salvo do TF2 compatível com TPU do modelo Keras ResNet-50.
```
cd demo/tf/resnet-50
python3 export_resnet_model.py
```

Inicie o contêiner do servidor de modelos do TensorFlow para o modelo.

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/tf/resnet_tpu,target=/models/resnet \
  -e MODEL_NAME=resnet \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Verifique o registro do contêiner do servidor de modelo e verifique se o gRPC e o servidor HTTP estão ativos:

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

Se o registro terminar com as informações a seguir, significa que o servidor está pronto para atender solicitações. Isso leva cerca de 30 segundos.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

Envie a solicitação ao servidor do modelo.

A imagem solicitada é uma banana de https://i.imgur.com/j9xCCzn.jpeg .

python3 resnet_request.py

A saída será semelhante a esta:

Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]

Fazer a limpeza.

Limpe o contêiner do Docker antes de executar outras demonstrações.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

Limpe os artefatos do modelo:

sudo rm -rf /tmp/tf/

Inferência e disponibilização de modelos do PyTorch

Para modelos criados com o PyTorch, o fluxo de trabalho é o seguinte:

Programar um gerenciador de modelo Python para carregar e fazer inferência usando TorchDynamo e PyTorch/XLA
Use TorchModelArchiver para criar um arquivo de modelo
Use TorchServe para disponibilizar o modelo

TorchDynamo e PyTorch/XLA

O TorchDynamo (Dynamo) é um compilador JIT no nível do Python projetado para tornar os programas PyTorch mais rápidos. Ele oferece uma API limpa para que os back-ends do compilador se conectem. Ele modifica dinamicamente o bytecode do Python antes da execução. Na versão PyTorch/XLA 2.0, há um back-end experimental para inferência e treinamento usando o Dynamo.

O Dynamo fornece um gráfico Torch FX (FX) quando reconhece um padrão de modelo, e o PyTorch/XLA usa uma abordagem de tensor preguiçoso para compilar o gráfico FX e retornar a função compilada. Para mais informações sobre o Dynamo, consulte:

Postagem do Pytorch Dev Discussions (em inglês)
Documentação do TorchDynamo
PyTorch 2.0 e XLA para mais detalhes

Confira um exemplo de código para executar a inferência de densenet161 com torch.compile.

import torch
import torchvision
import torch_xla.core.xla_model as xm

def eval_model(loader):
  device = xm.xla_device()
  xla_densenet161 = torchvision.models.densenet161().to(device)
  xla_densenet161.eval()
  dynamo_densenet161 = torch.compile(
      xla_densenet161, backend='torchxla_trace_once')
  for data, _ in loader:
    output = dynamo_densenet161(data)

TorchServe

É possível usar a imagem do Docker torchserve-tpu fornecida para disponibilizar o modelo pytorch arquivado em uma VM TPU.

Configure a autenticação para o Docker:

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
    us-docker.pkg.dev

Extraia a imagem do Docker do TorchServe do Cloud TPU para a VM da TPU:

CLOUD_TPU_TORCHSERVE_IMAGE_URL=us-docker.pkg.dev/cloud-tpu-images/inference/torchserve-tpu:v0.9.0-2.1
docker pull ${CLOUD_TPU_TORCHSERVE_IMAGE_URL}

Coletar artefatos de modelo

Para começar, é necessário fornecer um gerenciador de modelo, que instrui o worker do servidor de modelo do TorchServe a carregar seu modelo, processar os dados de entrada e executar a inferência. É possível usar os gerenciadores de inferência padrão do TorchServe (fonte) ou desenvolver seu próprio gerenciador de modelo personalizado seguindo o base_handler.py. Talvez também seja necessário fornecer o modelo treinado e o arquivo de definição do modelo.

No exemplo de Densenet 161 abaixo, usamos artefatos de modelo e o gerenciador de classificador de imagem padrão fornecido pelo TorchServe:

Configure algumas variáveis de ambiente:

CWD="$(pwd)"

WORKDIR="${CWD}/densenet_161"

mkdir -p ${WORKDIR}/model-store
mkdir -p ${WORKDIR}/logs

Faça o download e copie os artefatos do modelo do exemplo de classificador de imagens do TorchServe:

git clone https://github.com/pytorch/serve.git

cp ${CWD}/serve/examples/image_classifier/densenet_161/model.py ${WORKDIR}
cp ${CWD}/serve/examples/image_classifier/index_to_name.json ${WORKDIR}

Faça o download dos pesos do modelo:

wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth

mv densenet161-8d451a50.pth ${WORKDIR}

Crie um arquivo de configuração do modelo TorchServe para usar o back-end do Dynamo:

echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml

Você vai encontrar os seguintes arquivos e diretórios:

>> ls ${WORKDIR}
model_config.yaml
index_to_name.json
logs
model.py
densenet161-8d451a50.pth
model-store

Gerar um arquivo de modelo

Para exibir seu modelo PyTorch com o TorchServe do Cloud TPU, é necessário empacotar o gerenciador de modelo e todos os artefatos em um arquivo de modelo (*.mar) usando o arquivador de modelo do Torch.

Gere um arquivo de modelo com torch-model-archiver:

MODEL_NAME=Densenet161

docker run \
    --privileged  \
    --shm-size 16G \
    --name torch-model-archiver \
    -it \
    -d \
    --rm \
    --mount type=bind,source=${WORKDIR},target=/home/model-server/ \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torch-model-archiver \
        --model-name ${MODEL_NAME} \
        --version 1.0 \
        --model-file model.py \
        --serialized-file densenet161-8d451a50.pth \
        --handler image_classifier \
        --export-path model-store \
        --extra-files index_to_name.json \
        --config-file model_config.yaml

Você vai encontrar o arquivo de modelo gerado no diretório model-store:

>> ls ${WORKDIR}/model-store
Densenet161.mar

Atender solicitações de inferência

Agora que você tem o arquivo de modelo, pode iniciar o servidor de modelo do TorchServe e atender às solicitações de inferência.

Inicie o servidor de modelo do TorchServe:

docker run \
    --privileged  \
    --shm-size 16G \
    --name torchserve-tpu \
    -it \
    -d \
    --rm \
    -p 7070:7070 \
    -p 7071:7071 \
    -p 8080:8080 \
    -p 8081:8081 \
    -p 8082:8082 \
    -p 9001:9001 \
    -p 9012:9012 \
    --mount type=bind,source=${WORKDIR}/model-store,target=/home/model-server/model-store \
    --mount type=bind,source=${WORKDIR}/logs,target=/home/model-server/logs \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torchserve \
        --start \
        --ncs \
        --models ${MODEL_NAME}.mar \
        --ts-config /home/model-server/config.properties

Consultar a integridade do servidor de modelo:

curl http://localhost:8080/ping

Se o servidor de modelo estiver em execução, você verá:

{
  "status": "Healthy"
}

Para consultar as versões padrão do modelo registrado atual, use:

curl http://localhost:8081/models

O modelo registrado vai aparecer:

{
  "models": [
    {
      "modelName": "Densenet161",
      "modelUrl": "Densenet161.mar"
    }
  ]
}

Para fazer o download de uma imagem para uso de inferência:

curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg

mv kitten_small.jpg ${WORKDIR}

Para enviar uma solicitação de inferência ao servidor do modelo, use:

curl http://localhost:8080/predictions/${MODEL_NAME} -T ${WORKDIR}/kitten_small.jpg

Uma resposta semelhante a esta vai aparecer:

{
  "tabby": 0.47878125309944153,
  "lynx": 0.20393909513950348,
  "tiger_cat": 0.16572578251361847,
  "tiger": 0.061157409101724625,
  "Egyptian_cat": 0.04997897148132324
}

Registros do servidor do modelo

Use os seguintes comandos para acessar os registros:
```
ls ${WORKDIR}/logs/
cat ${WORKDIR}/logs/model_log.log
```
Você vai encontrar a seguinte mensagem no seu registro:
```
"Compiled model with backend torchxla\_trace\_once"
```

Limpar

Pare o contêiner do Docker:

rm -rf serve
rm -rf ${WORKDIR}

docker stop torch-model-archiver
docker stop torchserve-tpu

Criação de perfil

Depois de configurar a inferência, use os profilers para analisar o desempenho e a utilização da TPU. Para mais informações sobre a criação de perfis, consulte:

Introdução à inferência da Cloud TPU

Informações gerais e benefícios

Primeiros passos

Inferência e exibição de modelos do Cloud TPU

Inferência e veiculação de modelos JAX

Formato SavedModel

Conversor de inferência

Requisitos do Inference Converter

Inferência e veiculação de modelos JAX

Use jax2tf para serializar um modelo JAX no formato SavedModel

Use o Converter de inferência para preparar o modelo salvo para veiculação

Usar o TensorFlow Serving

Exemplos de disponibilização de modelos JAX

Pré-requisitos

Exibir o modelo BERT JAX para inferência

Servir a difusão estável do JAX para inferência

TensorFlow Serving

Fluxo de trabalho do TensorFlow Serving

Executar a demonstração do TensorFlow ResNet-50 Serving

Inferência e disponibilização de modelos do PyTorch

TorchDynamo e PyTorch/XLA

TorchServe

Coletar artefatos de modelo

Gerar um arquivo de modelo

Atender solicitações de inferência

Limpar

Criação de perfil

Use `jax2tf` para serializar um modelo JAX no formato SavedModel