Halaman ini diterjemahkan oleh Cloud Translation API.

Pengantar Inferensi Cloud TPU v5e

Ringkasan dan manfaat

Cloud TPU v5e adalah akselerator AI yang dikembangkan Google dan dioptimalkan untuk pelatihan, fine-tuning, dan penayangan (inferensi) berbasis transformer, text-to-image, dan CNN. Slice TPU v5e dapat berisi hingga 256 chip.

Penyertaan mengacu pada proses deployment model machine learning terlatih ke lingkungan produksi, tempat model tersebut dapat digunakan untuk inferensi. SLO latensi adalah prioritas untuk penayangan.

Dokumen ini membahas penayangan model di TPU host tunggal. Slice TPU dengan 8 chip atau kurang memiliki satu VM atau host TPU dan disebut TPU host tunggal.

Mulai

Anda memerlukan kuota untuk TPU v5e. TPU on-demand memerlukan kuota tpu-v5s-litepod-serving. TPU yang dicadangkan memerlukan kuota tpu-v5s-litepod-serving-reserved. Untuk informasi selengkapnya, hubungi Cloud Sales.

Anda memerlukan akun dan project Google Cloud untuk menggunakan Cloud TPU. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud TPU

Anda menyediakan TPU v5e menggunakan Resource dalam antrean. Untuk mengetahui informasi selengkapnya tentang konfigurasi v5e yang tersedia untuk penayangan, lihat Jenis Cloud TPU v5e untuk penayangan.

Inferensi dan penayangan model Cloud TPU

Cara Anda menayangkan model untuk inferensi bergantung pada framework ML yang digunakan untuk menulis model Anda. TPU v5e mendukung model penayangan yang ditulis dalam JAX, TensorFlow, dan PyTorch.

Inferensi dan penayangan model JAX

Untuk menayangkan model di VM TPU, Anda harus:

Serasikan model Anda dalam format SavedModel TensorFlow
Menggunakan Inference Converter untuk menyiapkan model tersimpan untuk ditayangkan
Menggunakan TensorFlow Serving untuk menayangkan model

Format SavedModel

SavedModel berisi program TensorFlow lengkap, termasuk parameter dan komputasi yang dilatih. Kode ini tidak memerlukan kode pembuatan model asli untuk dijalankan.

Jika model ditulis dalam JAX, Anda harus menggunakan jax2tf untuk melakukan serialisasi model dalam format SavedModel.

Konverter Inferensi

Konverter Inferensi Cloud TPU menyiapkan dan mengoptimalkan model yang diekspor dalam format SavedModel untuk inferensi TPU. Anda dapat menjalankan pengonversi inferensi di shell lokal atau VM TPU. Sebaiknya gunakan shell VM TPU karena memiliki semua alat command line yang diperlukan untuk menjalankan pengonversi. Untuk informasi selengkapnya tentang Inference Converter, lihat Panduan Pengguna Inference Converter.

Persyaratan Konverter Inferensi

Model Anda harus diekspor dari TensorFlow atau JAX dalam format SavedModel.
Anda harus menentukan alias fungsi untuk fungsi TPU. Untuk informasi selengkapnya, lihat Panduan Pengguna Inference Converter. Contoh dalam panduan ini menggunakan tpu_func sebagai alias fungsi TPU.
Pastikan CPU komputer Anda mendukung instruksi Advanced Vector eXtensions (AVX), karena library TensorFlow (dependensi Cloud TPU Inference Converter) dikompilasi untuk menggunakan instruksi AVX. Sebagian besar CPU memiliki dukungan AVX.

Catatan: Anda dapat menjalankan lscpu | grep avx untuk memeriksa apakah set petunjuk AVX didukung.

Inferensi dan penayangan model JAX

Bagian ini menjelaskan cara menayangkan model JAX menggunakan jax2tf dan TensorFlow Serving.

Menggunakan jax2tf untuk melakukan serialisasi model ke dalam format SavedModel
Menggunakan Inference Converter untuk menyiapkan model tersimpan Anda untuk ditayangkan
Menggunakan TensorFlow Serving untuk menayangkan model

Menggunakan `jax2tf` untuk melakukan serialisasi model JAX ke format SavedModel

Fungsi Python berikut menunjukkan cara menggunakan jax2tf dalam kode model Anda:

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

Untuk mengetahui informasi selengkapnya tentang jax2tf, lihat Interoperabilitas JAX dan Cloud TPU.

Menggunakan Inference Converter untuk menyiapkan model tersimpan untuk ditayangkan

Petunjuk untuk menggunakan Pengonversi Inferensi dijelaskan dalam panduan Pengonversi inferensi.

Menggunakan TensorFlow Serving

Petunjuk untuk menggunakan TensorFlow Serving dijelaskan dalam TensorFlow serving.

Contoh penayangan model JAX

Prasyarat

Siapkan kredensial Docker Anda dan ambil image Docker Inference Converter dan 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

Hubungkan ke VM TPU Anda dengan SSH dan instal kode demo inferensi:

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

Instal dependensi demo JAX:

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

Menayangkan model BERT JAX untuk inferensi

Anda dapat mendownload model BERT terlatih dari Hugging Face.

Mengekspor model tersimpan TensorFlow yang kompatibel dengan TPU dari model Flax BERT:
```
cd demo/jax/bert
python3 export_bert_model.py
```

Mulai penampung server model 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

Sekitar 30 detik setelah container dimulai, periksa log penampung server model dan pastikan server gRPC dan HTTP sudah aktif:

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

Jika Anda melihat entri log yang diakhiri dengan informasi berikut, server sudah siap menayangkan permintaan.

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

Kirim permintaan inferensi ke server model.

python3 bert_request.py

Outputnya akan mirip dengan berikut ini:

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

Jalankan pembersihan.

Pastikan untuk membersihkan penampung Docker sebelum menjalankan demo lainnya.

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

Bersihkan artefak model:

sudo rm -rf /tmp/jax/

Menayangkan Stable Diffusion JAX untuk inferensi

Anda dapat mendownload model Stable Diffusion terlatih dari Hugging Face.

Download model Stable Diffusion dalam format saved model TF2 yang kompatibel dengan TPU:

Catatan: Download memerlukan waktu sekitar 3 menit.
```
cd demo/jax/stable_diffusion
python3 export_stable_diffusion_model.py
```

Mulai penampung server model Cloud TPU untuk model:

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

Setelah sekitar dua menit, periksa log penampung server model untuk memastikan server gRPC dan HTTP berjalan:

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

Jika Anda melihat log yang diakhiri dengan informasi berikut, berarti server siap melayani permintaan.

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

Kirim permintaan ke server model.
```
python3 stable_diffusion_request.py
```
Skrip ini mengirim "Lukisan tupai yang bermain skate di New York" sebagai perintah. Gambar output akan disimpan sebagai stable_diffusion_images.jpg di direktori saat ini.

Jalankan pembersihan.

Pastikan untuk membersihkan penampung Docker sebelum menjalankan demo lainnya.

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

Membersihkan artefak model

sudo rm -rf /tmp/jax/

TensorFlow Serving

Petunjuk berikut menunjukkan cara menayangkan model TensorFlow di VM TPU.

Alur kerja TensorFlow Serving

Download image Docker TensorFlow Serving untuk VM TPU Anda.

Menetapkan contoh variabel lingkungan

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

Mendownload image Docker

docker pull ${IMAGE_NAME}

Siapkan kredensial Docker dan ambil image Docker Inference Converter dan 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

Download kode demo:

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

Instal dependensi demo TensorFlow:

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

Menayangkan model TensorFlow menggunakan image Docker TensorFlow Serving di 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}

Gunakan Serving Client API untuk membuat kueri model Anda.
- REST Client API
- gRPC Client API

Menjalankan demo TensorFlow ResNet-50 Serving

Mengekspor model tersimpan TF2 yang kompatibel dengan TPU dari model Keras ResNet-50.
```
cd demo/tf/resnet-50
python3 export_resnet_model.py
```

Luncurkan penampung server model TensorFlow untuk model.

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

Periksa log penampung server model dan pastikan Server gRPC dan HTTP sudah aktif:

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

Jika Anda melihat log yang diakhiri dengan informasi berikut, berarti server siap untuk melayani permintaan. Prosesnya memerlukan waktu sekitar 30 detik.

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

Kirim permintaan ke server model.

Gambar permintaan adalah pisang dari https://i.imgur.com/j9xCCzn.jpeg .

python3 resnet_request.py

Outputnya akan mirip dengan berikut ini:

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

Jalankan pembersihan.

Pastikan untuk membersihkan penampung Docker sebelum menjalankan demo lainnya.

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

Bersihkan artefak model:

sudo rm -rf /tmp/tf/

Inferensi dan penayangan model PyTorch

Untuk model yang ditulis dengan PyTorch, alur kerjanya adalah:

Menulis pengendali model Python untuk memuat dan melakukan inferensi menggunakan TorchDynamo dan PyTorch/XLA
Menggunakan TorchModelArchiver untuk membuat arsip model
Menggunakan TorchServe untuk menayangkan model

TorchDynamo dan PyTorch/XLA

TorchDynamo (Dynamo) adalah compiler JIT tingkat Python yang dirancang untuk mempercepat program PyTorch. API ini menyediakan API yang bersih untuk dihubungkan dengan backend compiler. Alat ini mengubah bytecode Python secara dinamis tepat sebelum dieksekusi. Dalam rilis PyTorch/XLA 2.0, ada backend eksperimental untuk inferensi dan pelatihan menggunakan Dynamo.

Dynamo menyediakan grafik Torch FX (FX) saat mengenali pola model dan PyTorch/XLA menggunakan pendekatan tensor lambat untuk mengompilasi grafik FX dan menampilkan fungsi yang dikompilasi. Untuk informasi selengkapnya tentang Dynamo, lihat:

Postingan Diskusi Developer Pytorch
Dokumentasi TorchDynamo
PyTorch 2.0 & XLA untuk mengetahui detail selengkapnya

Berikut adalah contoh kode kecil untuk menjalankan inferensi densenet161 dengan 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

Anda dapat menggunakan image Docker torchserve-tpu yang disediakan untuk menayangkan model pytorch yang diarsipkan di VM TPU.

Siapkan autentikasi untuk Docker:

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

Ambil image Docker TorchServe Cloud TPU ke VM TPU Anda:

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}

Mengumpulkan artefak model

Untuk memulai, Anda perlu menyediakan pengendali model, yang menginstruksikan pekerja server model TorchServe untuk memuat model Anda, memproses data input, dan menjalankan inferensi. Anda dapat menggunakan pengendali inferensi default TorchServe (sumber), atau mengembangkan pengendali model kustom Anda sendiri dengan mengikuti base_handler.py. Anda mungkin juga perlu menyediakan model yang dilatih, dan file definisi model.

Dalam contoh Densenet 161 berikut, kita menggunakan artefak model dan pengendali pengklasifikasi gambar default yang disediakan oleh TorchServe:

Konfigurasikan beberapa variabel lingkungan:

CWD="$(pwd)"

WORKDIR="${CWD}/densenet_161"

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

Download dan salin artefak model dari contoh pengklasifikasi gambar 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}

Download bobot model:

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

mv densenet161-8d451a50.pth ${WORKDIR}

Buat file konfigurasi model TorchServe untuk menggunakan backend Dynamo:

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

Anda akan melihat file dan direktori berikut:

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

Membuat file arsip model

Untuk menayangkan model PyTorch dengan Cloud TPU TorchServe, Anda perlu mengemas pengendali model dan semua artefak model ke dalam file arsip model (*.mar) menggunakan Torch Model Archiver.

Buat file arsip model dengan 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

Anda akan melihat file arsip model yang dihasilkan di direktori model-store:

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

Menayangkan permintaan inferensi

Setelah memiliki file arsip model, Anda dapat memulai server model TorchServe dan menayangkan permintaan inferensi.

Mulai server model 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

Buat kueri kesehatan server model:

curl http://localhost:8080/ping

Jika server model sudah aktif dan berjalan, Anda akan melihat:

{
  "status": "Healthy"
}

Untuk membuat kueri versi default model terdaftar saat ini, gunakan:

curl http://localhost:8081/models

Anda akan melihat model terdaftar:

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

Untuk mendownload gambar yang akan digunakan untuk inferensi:

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

mv kitten_small.jpg ${WORKDIR}

Untuk mengirim permintaan inferensi ke server model, gunakan:

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

Anda akan melihat respons seperti berikut:

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

Log server model

Gunakan perintah berikut untuk mengakses log:

ls ${WORKDIR}/logs/
cat ${WORKDIR}/logs/model_log.log

Anda akan melihat pesan berikut di log:

"Compiled model with backend torchxla\_trace\_once"

Pembersihan

Hentikan container Docker:

rm -rf serve
rm -rf ${WORKDIR}

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

Pembuatan profil

Setelah menyiapkan inferensi, Anda dapat menggunakan profiler untuk menganalisis performa dan penggunaan TPU. Untuk informasi selengkapnya tentang pembuatan profil, lihat:

Pengantar Inferensi Cloud TPU v5e

Ringkasan dan manfaat

Mulai

Inferensi dan penayangan model Cloud TPU

Inferensi dan penayangan model JAX

Format SavedModel

Konverter Inferensi

Persyaratan Konverter Inferensi

Inferensi dan penayangan model JAX

Menggunakan jax2tf untuk melakukan serialisasi model JAX ke format SavedModel

Menggunakan Inference Converter untuk menyiapkan model tersimpan untuk ditayangkan

Menggunakan TensorFlow Serving

Contoh penayangan model JAX

Prasyarat

Menayangkan model BERT JAX untuk inferensi

Menayangkan Stable Diffusion JAX untuk inferensi

TensorFlow Serving

Alur kerja TensorFlow Serving

Menjalankan demo TensorFlow ResNet-50 Serving

Inferensi dan penayangan model PyTorch

TorchDynamo dan PyTorch/XLA

TorchServe

Mengumpulkan artefak model

Membuat file arsip model

Menayangkan permintaan inferensi

Pembersihan

Pembuatan profil

Menggunakan `jax2tf` untuk melakukan serialisasi model JAX ke format SavedModel