Questa pagina è stata tradotta dall'API Cloud Translation.

Introduzione all'inferenza Cloud TPU v5e

Panoramica e vantaggi

Cloud TPU v5e è un acceleratore AI 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 in un ambiente di produzione, dove 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. I segmenti TPU con meno di 8 chip hanno una VM o un host TPU e sono chiamati TPU a singolo host.

Inizia

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

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

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

Inferenza e pubblicazione di modelli Cloud TPU

La modalità di pubblicazione di un modello per l'inferenza dipende dal framework ML con cui è stato scritto. TPU v5e supporta la pubblicazione 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
Utilizza TensorFlow Serving per pubblicare il modello

Formato SavedModel

Un SavedModel contiene un programma TensorFlow completo, inclusi parametri addestrati e calcoli. Non richiede l'esecuzione del codice di creazione del modello originale.

Se il modello è stato scritto in JAX, dovrai utilizzare jax2tf per serializzarlo 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 su Inference Converter, consulta la Guida dell'utente di Inference Converter.

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 di Inference Converter. Gli esempi in questa guida utilizzano tpu_func come alias della funzione TPU.
Assicurati che la CPU della tua macchina supporti le istruzioni AVX (Advanced Vector eXtensions), poiché la libreria TensorFlow (la dipendenza di Cloud TPU Inference Converter) viene compilata per utilizzare le istruzioni AVX. La maggior parte delle CPU supporta 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 pubblicare modelli JAX utilizzando jax2tf e TensorFlow Serving.

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

Utilizza `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 saperne di più su jax2tf, consulta la sezione Interoperabilità di JAX e Cloud TPU.

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

Le istruzioni per l'utilizzo di Inference Converter sono descritte nella guida di Inference Converter.

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

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}

Elimina gli artefatti del modello:

sudo rm -rf /tmp/jax/

Esegui Stable Diffusion di JAX per l'inferenza

Puoi scaricare il modello di diffusione stabile preaddestrato da Hugging Face.

Scarica il modello di diffusione stabile in un formato di modello salvato TF2 compatibile con TPU:

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 contenitore del server di modelli per assicurarti che i server gRPC e HTTP siano 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 nella directory corrente.

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}

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 sulle VM TPU.

Workflow di pubblicazione di TensorFlow

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 della 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 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 contenitore 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 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 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}

Elimina gli artefatti del modello:

sudo rm -rf /tmp/tf/

Inferenza e pubblicazione del modello PyTorch

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

Scrivi un gestore del modello 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) è un compilatore JIT a livello di Python progettato per velocizzare i programmi PyTorch. Fornisce un'API pulita per il collegamento ai backend del compilatore. Modifica dinamicamente il bytecode di Python poco 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 per l'esecuzione dell'inferenza di 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 pubblicare il modello PyTorch archiviato 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 gli handler di inferenza predefiniti di TorchServe (source) o sviluppare il tuo handler 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 elementi del modello dall'esempio di classificatore di immagini 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}

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 seguenti file e directory:

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

Generare 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 del modello 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 dell'archivio del modello generato nella directory model-store:

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

Gestire le richieste di inferenza

Ora che hai il file dell'archivio del modello, puoi avviare il server di modelli TorchServe e gestire le 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 registrato corrente, utilizza:

curl http://localhost:8081/models

Dovresti vedere il modello registrato:

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

Per scaricare un'immagine da utilizzare 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 di modelli, 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 del server del modello

Utilizza i seguenti comandi per accedere ai log:
```
ls ${WORKDIR}/logs/
cat ${WORKDIR}/logs/model_log.log
```
Nel log dovresti vedere 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 configurato l'inferenza, puoi utilizzare i profiler per analizzare il rendimento e l'utilizzo delle TPU. Per ulteriori informazioni sulla profilazione, vedi:

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

Utilizza 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

Esegui Stable Diffusion di JAX per l'inferenza

TensorFlow Serving

Workflow di pubblicazione di TensorFlow

Esegui la demo di TensorFlow ResNet-50 Serving

Inferenza e pubblicazione del modello PyTorch

TorchDynamo e PyTorch/XLA

TorchServe

Raccogliere gli artefatti del modello

Generare un file di archivio del modello

Gestire le richieste di inferenza

Esegui la pulizia

Profilazione

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