Questa pagina è stata tradotta dall'API Cloud Translation.

Introduzione all'inferenza Cloud TPU v5e

Panoramica e vantaggi

Cloud TPU v5e è un acceleratore IA sviluppato da Google ottimizzato per l'addestramento, l'ottimizzazione e il servizio (inferenza) basati su transformer, conversione di testo in immagine e CNN. Le sezioni TPU v5e possono contenere fino a 256 chip.

Per pubblicazione si intende il processo di deployment di un modello di machine learning addestrato a un ambiente di produzione in cui può essere utilizzato per l'inferenza. Gli SLO relativi alla latenza sono una priorità per il servizio.

Questo documento illustra la pubblicazione di un modello su una TPU a un solo host. Sezione TPU con Fino a 8 chip hanno una VM o un host TPU e sono chiamati TPU a host singolo.

Inizia

Avrai bisogno di una quota per le TPU v5e. Le TPU on demand richiedono una quota di tpu-v5s-litepod-serving. Le TPU riservate richiedono una quota pari a tpu-v5s-litepod-serving-reserved. Per ulteriori informazioni, contatta il team di vendita Cloud.

Per utilizzare Cloud TPU, devi avere un account e un progetto Google Cloud. Per maggiori informazioni per informazioni, consulta Configurare un ambiente Cloud TPU

Esegui il provisioning delle TPU v5e utilizzando le risorse in coda. Per ulteriori informazioni sulle configurazioni v5e disponibili per la pubblicazione, consulta Tipi di Cloud TPU v5e per la gestione

Inferenza e pubblicazione di modelli Cloud TPU

Il modo in cui fornisci un modello per l'inferenza dipende dal framework ML in cui il modello è stato con cui scrivere. TPU v5e supporta il servizio di modelli scritti in JAX, TensorFlow e PyTorch.

Inferenza e pubblicazione del modello JAX

Per pubblicare un modello su una VM TPU, devi:

Serializza il modello nel formato SavedModel di TensorFlow
Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione
Utilizzo di TensorFlow Serving per fornire il modello

Formato SavedModel

Un SavedModel contiene un programma TensorFlow completo, che include parametri e calcolo. Non richiede il codice edilizio originale eseguire.

Se il modello è scritto in JAX, devi usare jax2tf per serializzare del modello nel formato SavedModel.

Convertitore di inferenza

Il convertitore di inferenza Cloud TPU prepara e ottimizza un modello esportato in formato SavedModel per l'inferenza TPU. Puoi eseguire il convertitore di inferenza in una shell locale o nella VM TPU. Ti consigliamo di utilizzare la shell della VM TPU perché contiene tutti gli strumenti a riga di comando necessari per l'esecuzione del convertitore. Per ulteriori informazioni sull'Inferenza Convertitore, consulta la Guida dell'utente per la conversione di inferenza.

Requisiti di Inference Converter

Il modello deve essere esportato da TensorFlow o JAX nel formato SavedModel.
Devi definire un alias di funzione per la funzione TPU. Per ulteriori informazioni, consulta la Guida dell'utente del convertitore di inferenza. Negli esempi in questa guida viene utilizzato tpu_func come alias della funzione TPU.
Assicurati che la CPU della macchina supporti Advanced Vector eXtensions (AVX) di base, la libreria TensorFlow (la dipendenza il convertitore di inferenza di Cloud TPU) è compilato in modo da utilizzare le istruzioni AVX. La maggior parte delle CPU il supporto AVX.

Nota: puoi eseguire lscpu | grep avx per verificare se l'insieme di istruzioni AVX è supportato.

Inferenza e pubblicazione del modello JAX

Questa sezione descrive come gestire modelli JAX utilizzando jax2tf e TensorFlow Pubblicazione.

Utilizza jax2tf per serializzare il modello nel formato SavedModel
Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione
Utilizzo di TensorFlow Serving per fornire il modello

Usa `jax2tf` per serializzare un modello JAX nel formato SavedModel

La seguente funzione Python mostra come utilizzare jax2tf all'interno del codice del modello:

# 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)

Per ulteriori informazioni su jax2tf, consulta la sezione sull'interazione tra JAX e Cloud TPU.

Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione

Le istruzioni per l'utilizzo del convertitore di inferenza sono descritte nella sezione Guida per la conversione di inferenza.

Utilizzare TensorFlow Serving

Le istruzioni per l'utilizzo di TensorFlow Serving sono descritte in TensorFlow Serving.

Esempi di erogazione del modello JAX

Prerequisiti

Configura le tue credenziali Docker ed esegui il pull del convertitore di inferenza e di Cloud TPU Immagine Docker in uso:

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

Connettiti alla VM TPU con SSH e installa il codice di dimostrazione dell'inferenza:

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

Installa le dipendenze della demo JAX:

pip install -r ./demo/jax/requirements.txt

Pubblica il modello JAX BERT per l'inferenza

Puoi scaricare il modello BERT preaddestrato da Hugging Face.

Esporta un modello salvato TensorFlow compatibile con TPU da un modello Flax BERT:
```
cd demo/jax/bert
python3 export_bert_model.py
```

Avvia il container del server di modelli 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

Circa 30 secondi dopo l'avvio del contenitore, controlla il log del contenitore del server di modelli e assicurati che i server gRPC e HTTP siano attivi:

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

Se viene visualizzata una voce di log che termina con le seguenti informazioni, il server è pronto per gestire le richieste.

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 ...

Invia una richiesta di inferenza al server del modello.

python3 bert_request.py

L'output sarà simile al seguente:

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?."

Eseguire la pulizia.

Assicurati di ripulire il contenitore Docker prima di eseguire altre demo.
```
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}
```
Esegui la pulizia degli artefatti del modello:
```
sudo rm -rf /tmp/jax/
```

Fornisci lo Stable Diffusion della JAX per l'inferenza

Puoi scaricare il modello Stable Diffusion preaddestrato di Hugging Face.

Scarica il modello Stable Diffusion in un modello salvato TF2 compatibile con TPU formato:

Nota: il download richiede circa 3 minuti.
```
cd demo/jax/stable_diffusion
python3 export_stable_diffusion_model.py
```

Avvia il contenitore del server del modello Cloud TPU per il modello:

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

Dopo circa due minuti, controlla il log del container del server del modello per verificare sui server gRPC e HTTP sono in esecuzione:

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

Se il log termina con le seguenti informazioni, significa che i server sono pronti a gestire le richieste.

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 ...

Invia una richiesta al server del modello.
```
python3 stable_diffusion_request.py
```
Questo script invia "Dipinto di uno scoiattolo che pattina a New York" come prompt. L'immagine di output verrà salvata come stable_diffusion_images.jpg nel della directory corrente.

Eseguire la pulizia.

Assicurati di ripulire il container Docker prima di eseguire altre demo.

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

Elimina gli artefatti del modello

sudo rm -rf /tmp/jax/

TensorFlow Serving

Le istruzioni riportate di seguito mostrano come puoi pubblicare il tuo modello TensorFlow su VM TPU.

Flusso di lavoro di TensorFlow Serving

Scarica l'immagine Docker di TensorFlow Serving per la tua VM TPU.

Impostare le variabili di ambiente di esempio

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

Scarica l'immagine Docker

docker pull ${IMAGE_NAME}

Configura le credenziali Docker e recupera l'immagine Docker di Inference Converter e 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

Scarica il codice demo:

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

Installa le dipendenze demo di TensorFlow:

pip install -r ./demo/tf/requirements.txt

Pubblica il tuo modello TensorFlow utilizzando l'immagine Docker di TensorFlow Serving sulla VM 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}

Utilizza l'API Serving Client per eseguire query sul tuo modello.
- API client REST
- API client gRPC

Esegui la demo di TensorFlow ResNet-50 Serving

Esporta un modello salvato TF2 compatibile con TPU dal modello Keras ResNet-50.
```
cd demo/tf/resnet-50
python3 export_resnet_model.py
```

Avvia il container del server del modello TensorFlow per il modello.

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

Controlla il log del container del server del modello e assicurati che gRPC e server HTTP siano attivi:

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

Se il log termina con le seguenti informazioni, significa che il server è pronto a gestire le richieste. L'operazione richiede circa 30 secondi.

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 ...

Invia la richiesta al server del modello.

L'immagine richiesta è una banana da https://i.imgur.com/j9xCCzn.jpeg.

python3 resnet_request.py

L'output sarà simile al seguente:

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

Eseguire la pulizia.

Assicurati di ripulire il contenitore Docker prima di eseguire altre demo.
```
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}
```
Esegui la pulizia degli artefatti del modello:
```
sudo rm -rf /tmp/tf/
```

Inferenza e pubblicazione del modello PyTorch

Per i modelli scritti con PyTorch, il flusso di lavoro è:

Scrivi un gestore di modelli Python per il caricamento e l'inferenza utilizzando TorchDynamo e PyTorch/XLA
Usa TorchModelArchiver per creare un archivio di modelli
Utilizza TorchServe per pubblicare il modello

TorchDynamo e PyTorch/XLA

TorchDynamo (Dynamo) è una funzionalità di livello Python Compilatore JIT progettato per rendere più veloci i programmi PyTorch. Fornisce un'API pulita per il collegamento ai backend del compilatore. Modifica dinamicamente il bytecode Python prima dell'esecuzione. Nella release PyTorch/XLA 2.0 è presente un backend sperimentale per l'inferenza e l'addestramento utilizzando Dynamo.

Dynamo fornisce un grafo Torch FX (FX) quando riconosce un pattern del modello e PyTorch/XLA utilizza un approccio tensoriale lazy per compilare il grafo FX e restituire la funzione compilata. Per ulteriori informazioni su Dynamo, consulta:

Post di Pytorch Dev Discussions
Documentazione di TorchDynamo
PyTorch 2.0 e XLA per ulteriori dettagli

Ecco un piccolo esempio di codice dell'esecuzione dell'inferenza densenet161 con 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

Puoi utilizzare l'immagine Docker torchserve-tpu fornita per gestire i dati archiviati un modello pytorch su una VM TPU.

Configura l'autenticazione per Docker:

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

Estrai l'immagine Docker TorchServe Cloud TPU nella VM 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}

Raccogliere gli artefatti del modello

Per iniziare, devi fornire un gestore del modello, che indica al worker del server di modelli TorchServe di caricare il modello, elaborare i dati di input ed eseguire l'inferenza. Puoi utilizzare i gestori di inferenza predefiniti di TorchServe (fonte) oppure sviluppare il tuo gestore del modello personalizzato seguendo base_handler.py. Potresti anche dover fornire il modello addestrato e il file di definizione del modello.

Nel seguente esempio di Densenet 161, utilizziamo gli elementi del modello e l'handle predefinito per il classificatore delle immagini fornito da TorchServe:

Configura alcune variabili di ambiente:

CWD="$(pwd)"

WORKDIR="${CWD}/densenet_161"

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

Scarica e copia gli artefatti del modello dal classificatore di immagini TorchServe (esempio):

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}

Scarica i pesi del modello:

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

mv densenet161-8d451a50.pth ${WORKDIR}

Crea un file di configurazione del modello TorchServe per utilizzare il backend Dynamo:

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

Dovresti vedere i file e le directory seguenti:

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

Genera un file di archivio del modello

Per pubblicare il modello PyTorch con Cloud TPU TorchServe, devi pacchettizzare il gestore del modello e tutti gli elementi del modello in un file di archivio del modello (*.mar) utilizzando Torch Model Archiver.

Genera un file di archivio dei modelli con 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

Dovresti vedere il file di archivio dei modelli generato nella directory model-store:

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

Gestisci le richieste di inferenza

Ora che hai il file di archivio dei modelli, puoi avviare il server del modello TorchServe e gestire richieste di inferenza.

Avvia il server di modelli 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

Controlla l'integrità del server del modello di query:

curl http://localhost:8080/ping

Se il server del modello è attivo e funzionante, vedrai:

{
  "status": "Healthy"
}

Per eseguire query sulle versioni predefinite del modello attualmente registrato, usa:

curl http://localhost:8081/models

Dovresti vedere il modello registrato:

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

Per scaricare un'immagine da usare per l'inferenza:

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

mv kitten_small.jpg ${WORKDIR}

Per inviare una richiesta di inferenza al server del modello, utilizza:

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

Dovresti visualizzare una risposta simile alla seguente:

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

Log dei server del modello

Utilizza i comandi seguenti per accedere ai log:
```
ls ${WORKDIR}/logs/
cat ${WORKDIR}/logs/model_log.log
```
Nel log dovrebbe essere visualizzato il seguente messaggio:
```
"Compiled model with backend torchxla\_trace\_once"
```

Esegui la pulizia

Interrompi il container Docker:

rm -rf serve
rm -rf ${WORKDIR}

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

Profilazione

Dopo aver impostato l'inferenza, puoi usare i profiler per analizzare le prestazioni e l'utilizzo della TPU. Per ulteriori informazioni sulla profilazione, consulta:

Introduzione all'inferenza Cloud TPU v5e

Panoramica e vantaggi

Inizia

Inferenza e pubblicazione di modelli Cloud TPU

Inferenza e pubblicazione del modello JAX

Formato SavedModel

Convertitore di inferenza

Requisiti di Inference Converter

Inferenza e pubblicazione del modello JAX

Usa jax2tf per serializzare un modello JAX nel formato SavedModel

Utilizza il convertitore di inferenza per preparare il modello salvato per la pubblicazione

Utilizzare TensorFlow Serving

Esempi di erogazione del modello JAX

Prerequisiti

Pubblica il modello JAX BERT per l'inferenza

Fornisci lo Stable Diffusion della JAX per l'inferenza

TensorFlow Serving

Flusso di lavoro di TensorFlow Serving

Esegui la demo di TensorFlow ResNet-50 Serving

Inferenza e pubblicazione del modello PyTorch

TorchDynamo e PyTorch/XLA

TorchServe

Raccogliere gli artefatti del modello

Genera un file di archivio del modello

Gestisci le richieste di inferenza

Esegui la pulizia

Profilazione

Usa `jax2tf` per serializzare un modello JAX nel formato SavedModel