优化的 TensorFlow 运行时

与基于开源的预构建 TensorFlow Serving 容器相比,优化的 TensorFlow 运行时优化了模型,以实现更快、费用更低的推理。优化的 TensorFlow 运行时利用 Google 的专有开源技术来实现这一点。

机器学习 (ML) 模型越大,服务费用越高。与使用基于开源的 TensorFlow 运行时相比,使用优化的 TensorFlow 运行时后,传送机器学习模型的费用和推理速度可降低。使用 Vertex AI Prediction 时,为了利用优化的 TensorFlow 运行时,您无需修改代码。而是需要选择使用该产品的服务容器映像。

优化的 TensorFlow 运行时向后兼容预构建 TensorFlow Serving 容器。如果您运行具有预构建容器的 TensorFlow 模型,则可以轻松切换到优化的 TensorFlow 运行时容器。

虽然在使用优化的 TensorFlow 运行时时模型的性能会提高,但该性能影响会因不同类型的模型而异。

优化的 TensorFlow 运行时概览

与基于开源的预构建 TensorFlow Serving 容器相比,优化的 TensorFlow 运行时使用模型优化和新的专有 Google 技术来提高推理速度并降低其费用。

当 Vertex AI 在运行模型之前上传模型时,会进行优化。将模型部署到端点后,系统会将优化日志添加到预测日志中。您可以使用这些日志来排查优化过程中可能出现的问题。

以下主题介绍了优化的 TensorFlow 运行时中的优化改进。

模型优化

优化的 TensorFlow 运行时提供了以下三项模型优化。

模型 XLA 预编译

TensorFlow 模型运行时,所有操作都单独运行。运行单个操作会产生少量开销。优化的 TensorFlow 运行时可以利用 XLA 将所有或部分 TensorFlow 图预编译为更大的内核,从而消除这种开销。

模型 XLA 预编译功能是可选的,默认情况下处于停用状态。如需了解如何在部署期间启用模型 XLA 预编译,请参阅启用模型 XLA 预编译

模型压缩优化

优化的 TensorFlow 运行时可以更快地运行某些模型,并且对模型精确率的影响微乎其微,因而能够实现模型压缩优化。启用模型压缩优化后,优化的 TensorFlow 运行时会利用量化和权重修剪等方法来加快模型的运行速度。

模型压缩优化功能默认处于停用状态。如需了解如何在部署期间启用模型压缩优化,请参阅启用模型压缩优化

改进了 GPU 的表格模型性能

TensorFlow 表格模型通常在 CPU 上应用,因为它们无法有效利用加速器。优化的 TensorFlow 运行时解决了这个问题,它在 GPU 上运行模型中计算开销高的部分。模型的其余部分通过最小化主机和加速器之间的通信在 CPU 上运行。在 GPU 上运行费用高昂的部分并在 CPU 上运行其余部分,从而加快表格模型处理速度,降低费用。

优化的 TensorFlow 运行时优化了对以下表格模型类型的处理。

Cloud TPU 的自动模型优化

支持 Cloud TPU 的预构建优化 TensorFlow 运行时容器可以自动对要在 TPU 上运行的模型进行分区和优化。如需了解详情,请参阅部署到 Cloud TPU

TensorFlow 运行时 (TFRT) 的使用

优化的 TensorFlow 运行时可以使用 TensorFlow 运行时 (TFRT)。TFRT 可高效利用多线程主机 CPU、支持异步编程模型,并针对低效率问题进行了优化。

除了版本 2.8 之外,所有优化的 TensorFlow 运行时 CPU 容器映像中都启用了 TFRT CPU。如需停用 TFRT CPU,请将 use_tfrt 标志设置为 false

夜间版优化的 TensorFlow 运行时 GPU 容器映像和稳定版优化的 TensorFlow 运行时 GPU 容器映像版本 2.13 及更高版本中提供 TFRT GPU。如需启用 TFRT GPU,请将 use_tfrtallow_precompilation 标志设置为 true。GPU 容器映像上的 TFRT 可最大限度地减少主机 CPU 和 GPU 之间的数据传输开销。启用 TFRT 后,它可以与 XLA 编译搭配使用。由于 XLA 预编译已启用,因此您可能会遇到一些副作用,例如第一个请求的延迟时间增加。如需了解详情,请参阅启用模型 XLA 预编译

Google 运行时的使用

由于优化的 TensorFlow 运行时是使用 Google 的内部堆栈构建的,因此它可以在 Google 的专有运行时环境中运行。

优化的 TensorFlow 运行时容器映像

Vertex AI Prediction 提供两种类型的优化 TensorFlow 运行时容器映像:稳定映像和夜间映像。

稳定容器映像

与基于开源的预构建 TensorFlow Serving 容器一样,稳定版优化的 TensorFlow 运行时容器会与特定 TensorFlow 版本绑定。绑定到特定版本的优化 TensorFlow 运行时容器的保留时长与绑定到同一版本的开源构建相同。优化的 TensorFlow 运行时构建具有与开源 TensorFlow 构建相同的属性,但推理速度更快。

构建向后兼容。这意味着您应该可以使用新版容器运行在旧版 TensorFlow 上训练的模型。新版容器的性能应该优于旧版容器。在极少数情况下,在旧版 TensorFlow 上训练的模型可能无法使用新版容器来运行。

夜间容器映像

夜间版优化的 TensorFlow 运行时构建包括最新的改进和优化,但可能不如稳定版构建可靠。它们主要用于实验目的。夜间构建的名称会包含 nightly 标签。与稳定的容器映像不同,夜间容器不在 Vertex AI 服务等级协议 (SLA) 的涵盖范围内。

可用容器映像

以下夜间版和稳定版优化的 TensorFlow 运行时 Docker 容器映像可用。

机器学习框架版本 支持的加速器(以及 CUDA 版本,如适用) 补丁程序和支持日期结束 提供终止日期 支持的图片
Nightly 仅 CPU 不适用 不适用
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest
Nightly GPU (CUDA 12.x) 不适用 不适用
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.nightly:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.nightly:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.nightly:latest
Nightly Cloud TPU 不适用 不适用
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.nightly:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.nightly:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.nightly:latest
2.15 仅 CPU 2024 年 8 月 15 日 2025 年 8 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-15:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-15:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-15:latest
2.15 GPU (CUDA 12.x) 2024 年 8 月 15 日 2025 年 8 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-15:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-15:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-15:latest
2.15 Cloud TPU 2024 年 8 月 15 日 2025 年 8 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest
2.14 仅 CPU 2024 年 8 月 15 日 2025 年 8 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-14:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-14:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-14:latest
2.14 GPU (CUDA 12.x) 2024 年 8 月 15 日 2025 年 8 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-14:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-14:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-14:latest
2.13 仅 CPU 2024 年 8 月 15 日 2025 年 8 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-13:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-13:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-13:latest
2.13 GPU (CUDA 11.x) 2024 年 8 月 15 日 2025 年 8 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-13:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-13:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-13:latest
2.12 仅 CPU 2024 年 5 月 15 日 2025 年 5 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-12:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-12:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-12:latest
2.12 GPU (CUDA 11.x) 2024 年 5 月 15 日 2025 年 5 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-12:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-12:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-12:latest
2.11 仅 CPU 2023 年 11 月 15 日 2024 年 11 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-11:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-11:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-11:latest
2.11 GPU (CUDA 11.x) 2023 年 11 月 15 日 2024 年 11 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-11:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-11:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-11:latest
2.10 仅 CPU 2023 年 11 月 15 日 2024 年 11 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-10:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-10:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-10:latest
2.10 GPU (CUDA 11.x) 2023 年 11 月 15 日 2024 年 11 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-10:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-10:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-10:latest
2.9 仅 CPU 2023 年 11 月 15 日 2024 年 11 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-9:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-9:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-9:latest
2.9 GPU (CUDA 11.x) 2023 年 11 月 15 日 2024 年 11 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-9:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-9:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-9:latest
2.8 仅 CPU 2023 年 11 月 15 日 2024 年 11 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-8:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-8:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-8:latest
2.8 GPU (CUDA 11.x) 2023 年 11 月 15 日 2024 年 11 月 15 日
  • us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-8:latest
  • europe-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-8:latest
  • asia-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-8:latest

将优化的 TensorFlow 运行时与专用端点结合使用

使用专用端点通过 Vertex AI 提供在线预测,可实现与 Vertex AI 在线预测服务之间的低延迟、安全连接,并且速度会比使用公共端点更快。优化的 TensorFlow 运行时可能会需要处理延迟敏感型模型,因此您可以考虑将其与专用端点结合使用。如需了解详情,请参阅使用专用端点进行在线预测

使用优化的 TensorFlow 运行时部署模型

使用优化的 TensorFlow 运行时部署预测模型的过程与使用基于开源的预构建 TensorFlow Serving 容器部署模型的过程基本相同。唯一的区别是,您在创建模型时需要指定使用优化的 TensorFlow 运行时的容器映像,并且可以启用本主题前面介绍的优化标志。例如,如果您已使用 us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest 容器部署模型,则现在您可以使用 us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.2-8:latest 容器通过优化的 TensorFlow 运行时处理相同的模型。

以下代码示例演示了如何使用 us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest 优化的 TensorFlow 运行时容器创建模型。如需部署此模型,请使用与通过其他预构建 TensorFlow Serving 容器部署模型相同的过程。

如需详细了解此示例中使用的 ModelServiceClient,请参阅 ModelServiceClient 类。如需详细了解如何使用 Vertex AI 部署模型,请参阅使用 Vertex AI API 部署模型。如需详细了解 allow_precompilationallow_compression 设置,请参阅本主题前面介绍的模型优化

    from google.cloud.aiplatform import gapic as aip
    PROJECT_ID = PROJECT_ID
    REGION = LOCATION
    API_ENDPOINT = f"{REGION}-aiplatform.googleapis.com"
    PARENT = f"projects/{PROJECT_ID}/locations/{REGION}"

    client_options = {"api_endpoint": API_ENDPOINT}
    model_service_client = aip.ModelServiceClient(client_options=client_options)
    tf_opt_model_dict = {
        "display_name": "DISPLAY_NAME",
        "metadata_schema_uri": "",
        "artifact_uri": "MODEL_URI",
        "container_spec": {
            "image_uri": "us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest",
            "args": [
                # The optimized TensorFlow runtime includes the following
                # options that can be set here.
                # "--allow_precompilation=true" - enable XLA precompilation
                # "--allow_compression=true" - enable
                #    model compression optimization
            ],
        },
    }
    tf_opt_model = model_service_client.upload_model(
        parent=PARENT,
        model=tf_opt_model_dict).result(timeout=180).model

模型优化标志

使用优化的 TensorFlow 运行时部署模型时,您可以启用两项功能来进一步优化 TensorFlow 模型的处理。

  1. 模型 XLA 预编译
  2. 影响模型压缩的优化

您可以同时启用模型 XLA 预编译和模型压缩优化。以下主题介绍了如何在部署期间使用标志启用这些选项。

启用模型 XLA 预编译

如需将优化的 TensorFlow 运行时配置为预编译模型,请将 allow_precompilation 标志设置为 true。模型 XLA 预编译适用于多种类型的模型,并且在大多数情况下都能提高性能。XLA 预编译最适合批量大小较大的请求。

当使用新批次大小的第一个请求到达时,就会进行模型 XLA 预编译。如需确保在第一个请求之前初始化运行时,您可以添加预热请求文件。如需了解详情,请参阅 TensorFlow 网站上的 SavedModel 预热

XLA 预编译需要几秒钟到几分钟的时间才能完成,具体取决于模型复杂程度。如果您使用模型 XLA 预编译,则应考虑以下事项。

  • 如果您使用预热文件,请尝试添加包含批次大小的请求,这些批次大小应代表您希望模型接收的批次大小。在预热文件中提供大量请求会降低模型服务器的启动速度。

  • 如果您希望模型接收批次大小不同的请求,则可能需要启用具有一组固定 allow_batch_sizes 值的服务器端批处理。如需详细了解如何启用服务器端

    批处理,请参阅 TensorFlow 网站上的为 TensorFlow 启用服务器端请求批处理

  • 由于 XLA 预编译会增加内存开销,因此某些大型模型可能会在 GPU 上出现“内存不足”错误,从而导致模型处理失败。

建议您先对模型进行 XLA 预编译测试,然后再在生产环境中启用此功能。

启用模型压缩优化

如需将优化的 TensorFlow 运行时配置为使用模型压缩优化,请将其 allow_compression 标志设置为 true。先测试启用此标志对模型精确率的影响,然后再确定是否要在生产环境中启用此功能。

停用优化

如需将优化的 TensorFlow 运行时配置为在未优化的情况下运行模型,请将其 disable_optimizer 标志设置为 true

优化的 TensorFlow 运行时限制

优化的 TensorFlow 运行时具有以下限制:

  • 优化的 TensorFlow 运行时与旧版 NVIDIA GPU 不兼容,例如 Tesla K80、Tesla P4 和 Tesla P100。
  • 优化的 TensorFlow 运行时目前仅支持采样 Shapely 可解释性。

价格

使用优化的 TensorFlow 运行时部署模型不会产生额外的费用。它与其他预测部署相同,都是根据使用的虚拟机和加速器数量收费。如需了解详情,请参阅 Vertex AI 价格