此页面由 Cloud Translation API 翻译。

Cloud TPU v5e 推断简介

概览和优势

Cloud TPU v5e 是 Google 开发的 AI 加速器，针对基于 Transformer、文本到图像和基于 CNN 的训练、微调和传送，（推理）。TPU v5e 切片最多可包含 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 虚拟机中运行推理转换器。我们建议您使用 TPU 虚拟机 shell，因为它包含所有命令行工具所需的资源。如需详细了解推断转换器，请参阅推断转换器用户指南。

推断转换器要求

您的模型必须从 TensorFlow 或 JAX 导出 SavedModel 格式。
您必须为 TPU 函数定义函数别名。如需更多信息请参阅推断转换器用户指南。本指南中的示例使用 tpu_func 作为 TPU 函数别名。
确保您的机器 CPU 支持高级矢量扩展 (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 模型以进行推理

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

从 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 以进行推理

您可以下载预训练的 Stable Diffusion 模型来自 Hugging Face。

在与 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

以下说明演示了如何提供 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

使用 TensorFlow Serving 应用 TensorFlow 模型您的 TPU 虚拟机上的 Docker 映像。

# 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 客户端 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 级别的旨在提高 PyTorch 程序速度的 JIT 编译器。它提供了一个简洁的 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 映像来传送您的 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 格式