此页面由 Cloud Translation API 翻译。

Cloud TPU v5e 推断简介

概览和优势

Cloud TPU v5e 是 Google 开发的 AI 加速器，针对基于 Transformer、文本到图像和基于 CNN 的训练、微调和传送，（推理）。TPU v5e slice 最多可包含 256 个芯片。

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

本文档介绍如何在单主机 TPU 上提供模型。TPU 切片 8 个或更少的芯片具有一个 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 格式

DDA 包含完整的 TensorFlow 程序，包括经过训练的参数和计算。它不需要原始模型构建代码运行。

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

推理转换器

Cloud TPU 推断转换器准备并优化适用于 TPU 的 SavedModel 格式推理。您可以在本地 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 提供 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 保存模型格式的 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 已保存模型。
```
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 Server 已启动：

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 编写的模型，工作流为：

使用 TorchDynamo 编写 Python 模型处理程序，以进行加载和推理和 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 默认的推理处理程序（来源），或按照 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 格式