此页面由 Cloud Translation API 翻译。

Cloud TPU v5e 推理简介

概览和优势

Cloud TPU v5e 是 Google 开发的 AI 加速器，针对基于转换器的文本到图像模型以及基于 CNN 的训练、微调和服务（推理）进行了优化。TPU v5e slice 最多可包含 256 个芯片。

服务是指将训练好的机器学习模型部署到生产环境中的过程，以便在该环境中进行推理。延迟时间 SLO 是服务的优先事项。

本文档介绍了如何在单主机 TPU 上提供模型。芯片数量不超过 8 个的 TPU 切片只有一个 TPU 虚拟机或主机，称为单主机 TPU。

开始使用

您需要 v5e TPU 配额。按需 TPU 需要 tpu-v5s-litepod-serving 配额。预留的 TPU 需要 tpu-v5s-litepod-serving-reserved 配额。如需了解详情，请与 Cloud 销售团队联系。

您需要拥有 Google Cloud 账号和项目才能使用 Cloud TPU。如需了解详情，请参阅设置 Cloud TPU 环境

您可以使用已排队资源预配 v5e TPU。如需详细了解可用于服务的 v5e 配置，请参阅用于服务的 Cloud TPU v5e 类型。

Cloud TPU 模型推理和服务

如何提供模型以进行推理取决于模型的编写所用机器学习框架。TPU v5e 支持提供使用 JAX、TensorFlow 和 PyTorch 编写的模型。

JAX 模型推理和服务

如需在 TPU 虚拟机上提供模型，您需要：

以 TensorFlow SavedModel 格式序列化模型
使用推理转换器准备已保存的模型以供分发
使用 TensorFlow Serving 提供模型

SavedModel 格式

SavedModel 包含一个完整的 TensorFlow 程序，其中包括经过训练的参数和计算。无需运行原始模型构建代码。

如果您的模型是使用 JAX 编写的，则需要使用 jax2tf 以 SavedModel 格式序列化模型。

推理转换器

Cloud TPU 推理转换器会准备和优化以 SavedModel 格式导出的模型，以便进行 TPU 推理。您可以在本地 Shell 或 TPU VM 中运行推理转换器。我们建议您使用 TPU 虚拟机 Shell，因为它包含运行转换器所需的所有命令行工具。如需详细了解推理转换器，请参阅推理转换器用户指南。

推理转换器要求

您的模型必须从 TensorFlow 或 JAX 导出，且格式为 SavedModel。
您必须为 TPU 函数定义函数别名。如需了解详情，请参阅推理转换器用户指南。本指南中的示例使用 tpu_func 作为 TPU 函数别名。
请确保您的机器 CPU 支持 Advanced Vector eXtensions (AVX) 指令，因为 TensorFlow 库（Cloud TPU 推理转换器的依赖项）会编译为使用 AVX 指令。大多数 CPU 都支持 AVX。

注意：您可以运行 lscpu | grep avx 来检查是否支持 AVX 指令集。

JAX 模型推理和服务

本部分介绍了如何使用 jax2tf 和 TensorFlow Serving 提供 JAX 模型。

使用 jax2tf 将模型序列化为 SavedModel 格式
使用推理转换器准备要部署的已保存模型
使用 TensorFlow Serving 提供模型

使用 `jax2tf` 将 JAX 模型序列化为 SavedModel 格式

以下 Python 函数展示了如何在模型代码中使用 jax2tf：

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

如需详细了解 jax2tf，请参阅 JAX 与 Cloud TPU 的互操作性。

使用推理转换器准备已保存的模型以供分发

如需了解如何使用推理转换器，请参阅推理转换器指南。

使用 TensorFlow Serving

如需了解如何使用 TensorFlow Serving，请参阅 TensorFlow Serving。

JAX 模型部署示例

前提条件

设置 Docker 凭据，然后拉取推理转换器和 Cloud TPU 服务 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

使用 SSH 连接到 TPU 虚拟机，然后安装推理演示代码：

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

安装 JAX 演示依赖项：

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

提供 JAX BERT 模型以进行推理

您可以从 Hugging Face 下载预训练的 BERT 模型。

从 Flax BERT 模型导出与 TPU 兼容的 TensorFlow 保存模型：
```
cd demo/jax/bert
python3 export_bert_model.py
```

启动 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

容器启动大约 30 秒后，请检查模型服务器容器日志，并确保 gRPC 和 HTTP 服务器已启动：

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

如果您看到以以下信息结尾的日志条目，则表示服务器已准备好处理请求。

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

向模型服务器发送推理请求。

python3 bert_request.py

输出将如下所示：

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

清理。

请务必先清理 Docker 容器，然后再运行其他演示。

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

清理模型工件：

sudo rm -rf /tmp/jax/

提供 JAX Stable Diffusion 以进行推理

您可以从 Hugging Face 下载预训练的稳定型 Diffusion 模型。

下载与 TPU 兼容的 TF2 SavedModel 格式的 Stable Diffusion 模型：

注意：下载过程大约需要 3 分钟。
```
cd demo/jax/stable_diffusion
python3 export_stable_diffusion_model.py
```

启动模型的 Cloud TPU 模型服务器容器：

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

大约两分钟后，检查模型服务器容器日志，确保 gRPC 和 HTTP 服务器正在运行：

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

如果您看到日志以以下信息结尾，则表示服务器已准备好处理请求。

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

向模型服务器发送请求。
```
python3 stable_diffusion_request.py
```
此脚本会发送“Painting of a squirrel skating in New York”（一幅画，描绘一只松鼠在纽约滑冰）作为问题。输出图片将保存为 stable_diffusion_images.jpg 并存储在当前目录中。

清理。

请务必先清理 Docker 容器，然后再运行其他演示。

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

清理模型工件

sudo rm -rf /tmp/jax/

TensorFlow Serving

以下说明演示了如何在 TPU 虚拟机上部署 TensorFlow 模型。

TensorFlow Serving 工作流

下载适用于 TPU 虚拟机的 TensorFlow Serving Docker 映像。

设置示例环境变量

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

下载 Docker 映像

docker pull ${IMAGE_NAME}

设置 Docker 凭据，然后拉取推理转换器和 TensorFlow Serving 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

下载演示代码：

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

安装 TensorFlow 演示依赖项：

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

在 TPU 虚拟机上使用 TensorFlow Serving Docker 映像部署 TensorFlow 模型。

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

使用 Serving Client API 查询模型。
- REST 客户端 API
- gRPC Client API

运行 TensorFlow ResNet-50 Serving 演示

从 Keras ResNet-50 模型导出与 TPU 兼容的 TF2 SavedModel。
```
cd demo/tf/resnet-50
python3 export_resnet_model.py
```

为模型启动 TensorFlow 模型服务器容器。

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

检查模型服务器容器日志，并确保 gRPC 和 HTTP 服务器已启动：

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

如果您看到日志以以下信息结尾，则表示服务器已准备好处理请求。此过程大约需要 30 秒。

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

将请求发送到模型服务器。

请求图片是 https://i.imgur.com/j9xCCzn.jpeg 中的香蕉。

python3 resnet_request.py

输出将如下所示：

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

清理。

请务必先清理 Docker 容器，然后再运行其他演示。

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

清理模型工件：

sudo rm -rf /tmp/tf/

PyTorch 模型推理和服务

对于使用 PyTorch 编写的模型，工作流如下：

编写 Python 模型处理程序，以使用 TorchDynamo 和 PyTorch/XLA 加载和推理
使用 TorchModelArchiver 创建模型归档
使用 TorchServe 提供模型

TorchDynamo 和 PyTorch/XLA

TorchDynamo (Dynamo) 是一个 Python 级 JIT 编译器，旨在加快 PyTorch 程序的运行速度。它提供了一个整洁有序的 API，供编译器后端钩入。它会在执行前动态修改 Python 字节码。PyTorch/XLA 2.0 版本中提供了一个实验性后端，用于使用 Dynamo 进行推理和训练。

当 Dynamo 识别出模型模式时，它会提供 Torch FX (FX) 图，而 PyTorch/XLA 会使用延迟张量方法编译 FX 图并返回已编译的函数。如需详细了解 Dynamo，请参阅：

以下是一个使用 torch.compile 运行 densenet161 推理的简短代码示例。

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

您可以使用提供的 torchserve-tpu Docker 映像在 TPU 虚拟机上提供已归档的 PyTorch 模型。

为 Docker 设置身份验证：

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

将 Cloud TPU TorchServe Docker 映像拉取到 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}

收集模型工件

首先，您需要提供模型处理脚本，该脚本会指示 TorchServe 模型服务器工作器加载模型、处理输入数据并运行推理。您可以使用 TorchServe 默认推理处理程序 (source)，也可以按照 base_handler.py 开发自己的自定义模型处理程序。您可能还需要提供训练好的模型和模型定义文件。

在以下 Densenet 161 示例中，我们使用了模型工件和 TorchServe 提供的默认图像分类器处理程序：

配置一些环境变量：

CWD="$(pwd)"

WORKDIR="${CWD}/densenet_161"

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

从 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}

下载模型权重：

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

mv densenet161-8d451a50.pth ${WORKDIR}

创建 TorchServe 模型配置文件以使用 Dynamo 后端：

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

您应该会看到以下文件和目录：

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

生成模型归档文件

如需使用 Cloud TPU TorchServe 提供 PyTorch 模型，您需要使用 Torch Model Archiver 将模型处理程序和所有模型工件打包到模型归档文件 (*.mar) 中。

使用 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

您应该会在 model-store 目录中看到生成的模型归档文件：

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

处理推理请求

现在，您已经有了模型归档文件，可以启动 TorchServe 模型服务器并处理推理请求。

启动 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

查询模型服务器运行状况：

curl http://localhost:8080/ping

如果模型服务器已启动且正在运行，您会看到：

{
  "status": "Healthy"
}

如需查询当前已注册模型的默认版本，请使用以下命令：

curl http://localhost:8081/models

您应该会看到已注册的模型：

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

如需下载图片以进行推理，请使用以下命令：

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

mv kitten_small.jpg ${WORKDIR}

如需向模型服务器发送推理请求，请使用以下命令：

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

您应该会看到如下所示的响应：

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

模型服务器日志

使用以下命令访问日志：

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

您应该会在日志中看到以下消息：

"Compiled model with backend torchxla\_trace\_once"

清理

停止 Docker 容器：

rm -rf serve
rm -rf ${WORKDIR}

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

分析

设置推理后，您可以使用性能分析器来分析性能和 TPU 利用率。如需详细了解性能分析，请参阅：

Cloud TPU v5e 推理简介

概览和优势

开始使用

Cloud TPU 模型推理和服务

JAX 模型推理和服务

SavedModel 格式

推理转换器

推理转换器要求

JAX 模型推理和服务

使用 jax2tf 将 JAX 模型序列化为 SavedModel 格式

使用推理转换器准备已保存的模型以供分发

使用 TensorFlow Serving

JAX 模型部署示例

前提条件

提供 JAX BERT 模型以进行推理

提供 JAX Stable Diffusion 以进行推理

TensorFlow Serving

TensorFlow Serving 工作流

运行 TensorFlow ResNet-50 Serving 演示

PyTorch 模型推理和服务

TorchDynamo 和 PyTorch/XLA

TorchServe

收集模型工件

生成模型归档文件

处理推理请求

清理

分析

使用 `jax2tf` 将 JAX 模型序列化为 SavedModel 格式