Introdução à inferência do Cloud TPU v5e
Informações gerais e benefícios
O Cloud TPU v5e é um acelerador de IA desenvolvido pelo Google e otimizado para treinamento, ajuste e disponibilização baseados em transformador, texto para imagem e CNN (inferência). As frações da TPU v5e podem conter até 256 chips.
A disponibilização se refere ao processo de implantação de um modelo de machine learning treinado em um de produção, onde pode ser usado para inferência. SLOs de latência são a prioridade de veiculação.
Este documento discute a exibição de um modelo em uma TPU de host único. Frações de TPU com Até 8 chips têm uma VM ou host de TPU e são chamados de TPUs de host único.
Primeiros passos
Você vai precisar de cota para TPUs v5e. TPUs sob demanda exigem tpu-v5s-litepod-serving
a cota de ingestão de registros. TPUs reservadas exigem cota de tpu-v5s-litepod-serving-reserved
. Para mais
informações, entre em contato com a equipe de vendas do Cloud.
Você precisará de uma conta e um projeto do Google Cloud para usar a Cloud TPU. Para mais informações, consulte Configurar um ambiente do Cloud TPU
Provisione TPUs v5e usando recursos na fila. Para mais informações sobre as configurações disponíveis da v5e para exibição, consulte Tipos do Cloud TPU v5e para exibição.
Inferência e disponibilização de modelos do Cloud TPU
A forma como um modelo é fornecido para inferência depende do framework de ML em que o modelo foi usado escrito. A TPU v5e oferece suporte a modelos de veiculação escritos em JAX, TensorFlow e o PyTorch.
Inferência e disponibilização de modelos JAX
Para disponibilizar um modelo em uma VM da TPU, é necessário:
- Serializar o modelo no TensorFlow SavedModel formato
- Usar o conversor de inferência para preparar o modelo salvo para veiculação
- Usar o TensorFlow Serving para disponibilizar o modelo
Formato SavedModel
Um SavedModel contém um programa completo do TensorFlow, incluindo parâmetros e computação. Ele não exige o código de criação do modelo original para ser executado.
Se o modelo foi escrito em JAX, será necessário usar jax2tf
para serializar
o modelo no formato SavedModel.
Conversor de inferência
O Conversor de inferência do Cloud TPU prepara e otimiza um modelo exportado em Formato SavedModel para TPU a inferência. É possível executar o conversor de inferência em um shell local ou na VM da TPU. Recomendamos usar o shell da VM da TPU porque ele tem todas as ferramentas de linha de comando necessários para executar o conversor. Para mais informações sobre a inferência Converter, consulte o Guia do usuário do conversor de inferência.
Requisitos do conversor de inferência
Seu modelo deve ser exportado do TensorFlow ou do JAX no SavedModel.
Defina um alias para a função da TPU. Para mais informações, consulte o Guia do usuário do conversor de inferências. Os exemplos neste guia usam
tpu_func
como o alias da função da TPU.Verifique se a CPU da máquina é compatível com Advanced Vector eXtensions (AVX) da biblioteca do TensorFlow, que é a dependência Conversor de inferência do Cloud TPU) é compilado para usar as instruções do AVX. A maioria das CPUs seja compatível com AVX.
Inferência e disponibilização de modelos JAX
Esta seção descreve como disponibilizar modelos JAX usando o jax2tf
e o TensorFlow
Disponibilização.
- Usar
jax2tf
para serializar o modelo para o formato SavedModel - Usar o conversor de inferência para preparar seu modelo salvo para veiculação
- Usar o TensorFlow Serving para disponibilizar o modelo
Usar jax2tf
para serializar um modelo JAX para o formato SavedModel
A seguinte função do Python 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.
Usar o conversor de inferência para preparar o modelo salvo para veiculação
As instruções de uso do Conversor de inferências Guia do conversor de inferências.
Usar o TensorFlow Serving
As instruções para usar o TensorFlow Serving são descritas em Exibição do TensorFlow.
Exemplos de disponibilização de modelos JAX
Pré-requisitos
Configure suas credenciais do Docker e extraia o Inference Converter e o Cloud TPU Disponibilização da imagem Docker:
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 da 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
Disponibilizar o modelo JAX BERT para inferência
É possível fazer o download do modelo BERT pré-treinado do 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 modelo 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 contêiner ser iniciado, verifique o servidor de modelo registro do contêiner 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 uma entrada de registro termina com as informações a seguir, o servidor está 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 de inferência ao 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/
Disponibilizar o JAX Stable Diffusion para inferência
É possível fazer o download do modelo de difusão estável pré-treinado do Hugging Face.
Fazer o download do modelo de difusão estável em um modelo salvo do TF2 compatível com TPU formato:
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 confirmar os servidores gRPC e HTTP estão 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 o e 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 ao servidor do modelo.
python3 stable_diffusion_request.py
Este script envia "Pintura de um esquilo patinando em Nova York" como comando. 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 mostram como disponibilizar modelos do TensorFlow em VMs de TPU.
Fluxo de trabalho do TensorFlow Serving
Faça o download da imagem Docker do TensorFlow Serving para sua VM da TPU.
Definir variáveis de ambiente de amostra
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 Docker
docker pull ${IMAGE_NAME}
Configure as credenciais do Docker, extraia o Inference Converter e 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 de demonstração do TensorFlow:
pip install -r ./demo/tf/requirements.txt
Disponibilizar seu modelo do TensorFlow usando o TensorFlow Serving imagem Docker 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 o modelo.
Executar a demonstração do TensorFlow ResNet-50 Serving
Exportar um modelo TF2 salvo 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 modelo 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 modelos e confirme 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, o servidor está pronto para atender às solicitações. O processo 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 da solicitação é 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 PyTorch
Para modelos escritos com PyTorch, o fluxo de trabalho é:
- Escrever um gerenciador de modelos em Python para carregar e inferir usando
TorchDynamo
e PyTorch/XLA - Use
TorchModelArchiver
para criar um arquivo de modelo - Use
TorchServe
para exibir o modelo
TorchDynamo e PyTorch/XLA
O TorchDynamo (Dynamo) é um aplicativo Compilador JIT projetado para tornar os programas PyTorch mais rápidos. Ele fornece uma API limpa para que os back-ends do compilador se conectem. Ele modifica dinamicamente o bytecode Python antes da execução. Na versão 2.0 do PyTorch/XLA, 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 PyTorch/XLA usa uma abordagem de tensor lento para compilar o gráfico FX e retornar a função compilada. Para mais informações sobre a Dynamo, consulte:
- Postagem de discussões sobre desenvolvedores sobre o Pytorch (em inglês)
- Documentação do TorchDynamo
- PyTorch 2.0 e XLA para mais detalhes
Confira um pequeno exemplo de código da execução da inferência densnet161 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 exibir seu arquivo
pytorch em uma VM da TPU.
Configure a autenticação do Docker:
sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
us-docker.pkg.dev
Extraia a imagem do Cloud TPU TorchServe Docker para sua 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, você precisa fornecer um gerenciador de modelos, que instrui o Worker do servidor de modelo do TorchServe para carregar o modelo, processar os dados de entrada e e executar inferência. É possível usar os gerenciadores de inferência padrão do TorchServe. (fonte) ou desenvolver seu próprio gerenciador de modelos personalizado seguindo base_handler.py. Talvez também seja necessário fornecer o modelo treinado e o arquivo de definição do modelo.
No exemplo do Densenet 161 a seguir, usamos artefatos de modelo e o padrão gerenciador de classificador de imagem 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 imagem 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 de modelo TorchServe para usar o back-end do Dynamo:
echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml
Você verá 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 o modelo PyTorch com o Cloud TPU TorchServe, você precisa empacotar
o gerenciador e todos os artefatos do modelo em um arquivo de modelo
(*.mar)
usando o Toch Model Archiver.
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
O arquivo do modelo gerado no diretório model-store vai aparecer:
>> ls ${WORKDIR}/model-store
Densenet161.mar
Disponibilizar solicitações de inferência
Agora que você tem o arquivo do modelo, é possível iniciar o servidor do modelo TorchServe e exibir solicitações de inferência.
Inicie o servidor do modelo 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
Integridade do servidor do modelo de consulta:
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
Você verá o modelo registrado:
{ "models": [ { "modelName": "Densenet161", "modelUrl": "Densenet161.mar" } ] }
Para fazer o download de uma imagem para inferência, use:
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 de 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
A seguinte mensagem será exibida no seu registro:
"Compiled model with backend torchxla\_trace\_once"
Limpar
Interrompa 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, é possível usar os criadores de perfil para analisar o desempenho e o uso da TPU. Para mais informações sobre a criação de perfil, consulte:
- Como criar perfis no Cloud TPU
- Criação de perfis do TensorFlow
- Criação de perfis do PyTorch
- Criação de perfis do JAX