在Amazon SageMaker上使用Triton托管ML模型：ONNX模型

ONNX（Open Neural Network Exchange）是一个开源标准，用于表示深度学习模型，得到了许多提供商的广泛支持。 ONNX提供了优化和量化模型的工具，以减少运行机器学习（ML）模型所需的内存和计算量。 ONNX最大的好处之一是它提供了一种标准化的格式，用于在不同的框架和工具之间表示和交换ML模型。这使得开发人员可以在一个框架中训练他们的模型，并在另一个框架中部署它们，而无需进行大量的模型转换或重新训练。基于这些原因，ONNX在ML社区中变得越来越重要。

在本文中，我们展示了如何部署基于ONNX的模型，用于使用GPU的多模型端点（MMEs）。这是“在Amazon SageMaker多模型端点上运行多个深度学习模型”的文章的延续，我们展示了如何在Nvidia的Triton推理服务器上部署ResNet50模型的PyTorch和TensorRT版本。在本文中，我们使用相同的ResNet50模型，以ONNX格式以及另一个自然语言处理（NLP）示例模型，以展示如何在Triton上部署它。此外，我们对ResNet50模型进行基准测试，并比较了与PyTorch和TensorRT版本相同模型的性能优势，使用相同的输入。

ONNX运行时

ONNX Runtime是一种用于ML推理的运行时引擎，旨在优化跨多个硬件平台（包括CPU和GPU）的模型的性能。它允许使用像PyTorch和TensorFlow这样的ML框架。它促进了性能调整，以在目标硬件上以成本效益的方式运行模型，并支持量化和硬件加速等功能，使其成为部署高效，高性能ML应用程序的理想选择之一。有关如何将ONNX模型优化为带有TensorRT的Nvidia GPU的示例，请参阅TensorRT优化（ORT-TRT）和使用TensorRT优化的ONNX运行时。

Amazon SageMaker Triton容器流程如下图所示。

用户可以发送带有实时推理输入负载的HTTPS请求，位于SageMaker端点后面。用户可以指定包含请求目标要调用的模型名称的TargetModel标头。在内部，SageMaker Triton容器实现了一个HTTP服务器，具有与如何容器提供请求中提到的相同的合同。它支持动态批处理，并支持Triton提供的所有后端。根据配置，调用ONNX运行时，并根据用户提供的模型配置在CPU或GPU上处理请求。

解决方案概述

要使用ONNX后端，请完成以下步骤：

  1. 将模型编译为ONNX格式。
  2. 配置模型。
  3. 创建SageMaker端点。

先决条件

确保您拥有访问带有足够AWS身份验证和访问管理IAM权限的AWS帐户，以创建笔记本电脑，访问Amazon Simple Storage Service（Amazon S3）存储桶并将模型部署到SageMaker端点。有关更多信息，请参见创建执行角色。

将模型编译为ONNX格式

transformers库提供了一种方便的方法，将PyTorch模型编译为ONNX格式。以下代码实现了NLP模型的转换：

onnx_inputs, onnx_outputs = transformers.onnx.export(
    preprocessor=tokenizer,
    model=model,
    config=onnx_config,
    opset=12,
    output=save_path
 )

通过提供Hugging Face transformers存储库的转换工具，轻松导出模型（无论是PyTorch还是TensorFlow）。

以下是内部发生的情况：

  1. 从transformers（PyTorch或TensorFlow）中分配模型。
  2. 通过模型向前传递虚拟输入。这样，ONNX就可以记录运行的操作集。
  3. 在导出模型时，transformers固有地处理动态轴。
  4. 保存图形以及网络参数。

计算机视觉案例从 torchvision 模型库也遵循类似的机制：

torch.onnx.export(
        resnet50,
        dummy_input,
        args.save,
        export_params=True,
        opset_version=11,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
    )

配置模型

在本节中，我们配置计算机视觉和 NLP 模型。我们展示如何创建一个已经预训练用于在 SageMaker MME 上部署的 ResNet50 和 RoBERTA 大型模型，通过 Triton 推理服务器模型配置。ResNet50 笔记本可在 GitHub 上找到。RoBERTA 笔记本也可在 GitHub 上找到。对于 ResNet50，我们使用 Docker 方法创建一个已经拥有构建 ONNX 模型所需的所有依赖项和生成此练习所需的模型工件的环境。这种方法使得分享依赖项和创建需要完成此任务的确切环境变得更加容易。

第一步是按照 ONNX 模型中指定的目录结构创建 ONNX 模型包。我们的目标是使用一个包含在单个文件中的 ONNX 模型的最小模型存储库，如下所示：

<model-repository-path> / 
    Model_name
    ├── 1
    │   └── model.onnx
    └── config.pbtxt

接下来，我们创建描述 Triton Server 选择并调用适当的 ONNX 内核的输入、输出和后端配置的模型配置文件。该文件称为 config.pbtxt，RoBERTA 案例的代码如下所示。请注意，config.pbtxt 中省略了 BATCH 维度。但是，在将数据发送到模型时，我们会包含批处理维度。以下代码还显示了如何使用模型配置文件添加此功能，以设置动态批处理，实际推理的首选批大小为 5。当前设置下，当达到首选批大小 5 或自动批处理程序接收到第一个请求后经过 100 微秒的延迟时间时，模型实例将立即被调用。

name: "nlp-onnx"
platform: "onnxruntime_onnx"
backend: "onnxruntime" 
max_batch_size: 32

  input {
    name: "input_ids"
    data_type: TYPE_INT64
    dims: [512]
  }
  input {
    name: "attention_mask"
    data_type: TYPE_INT64
    dims: [512]
  }

  output {
    name: "last_hidden_state"
    data_type: TYPE_FP32
    dims: [-1, 768]
  }
  output {
    name: "1550"
    data_type: TYPE_FP32
    dims: [768]
  }
instance_group {
  count: 1
  kind: KIND_GPU
}
dynamic_batching {
    max_queue_delay_microseconds: 100
    preferred_batch_size:5
}

计算机视觉案例的配置文件如下：

name: "resenet_onnx"
platform: "onnxruntime_onnx"
max_batch_size : 128
input [
  {
    name: "input"
    data_type: TYPE_FP32
    format: FORMAT_NCHW
    dims: [ 3, 224, 224 ]
  }
]
output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [ 1000 ]
  }
]

创建 SageMaker 端点

我们使用 Boto3 APIs 创建 SageMaker 端点。在本文中，我们展示了 RoBERTA 笔记本的步骤，但这些步骤对于 ResNet50 模型也是通用的。

创建 SageMaker 模型

我们现在创建一个 SageMaker 模型。我们使用上一步的 Amazon Elastic Container Registry (Amazon ECR) 镜像和模型工件来创建 SageMaker 模型。

创建容器

要创建容器，我们从 Amazon ECR 中拉取 Triton Server 的适当映像。SageMaker 允许我们自定义和注入各种环境变量。一些关键特性是能够设置 BATCH_SIZE，我们可以在 config.pbtxt 中为每个模型设置此值，或在此处定义默认值。对于可以受益于更大共享内存大小的模型，我们可以在 SHM 变量下设置这些值。要启用日志记录，将日志 verbose 级别设置为 true。我们使用以下代码创建模型，以在我们的端点中使用：

mme_triton_image_uri = (
    f"{account_id_map[region]}.dkr.ecr.{region}.{base}" + "/sagemaker-tritonserver:22.12-py3"
)
container = {
    "Image": mme_triton_image_uri,
    "ModelDataUrl": mme_path,
    "Mode": "MultiModel",
    "Environment": {
        "SAGEMAKER_TRITON_SHM_DEFAULT_BYTE_SIZE": "16777216000", # "16777216", #"16777216000",
        "SAGEMAKER_TRITON_SHM_GROWTH_BYTE_SIZE": "10485760",
    },
}
from sagemaker.utils import name_from_base
model_name = name_from_base(f"flan-xxl-fastertransformer")
print(model_name)
create_model_response = sm_client.create_model(
    ModelName=model_name,
    ExecutionRoleArn=role,
    PrimaryContainer={
        "Image": inference_image_uri, 
        "ModelDataUrl": s3_code_artifact
    },
)
model_arn = create_model_response["ModelArn"]
print(f"Created Model: {model_arn}")

创建SageMaker端点

您可以使用任何带有多个GPU的实例进行测试。在本文中，我们使用g4dn.4xlarge实例。我们不设置VolumeSizeInGB参数，因为该实例带有本地实例存储。 VolumeSizeInGB参数适用于支持亚马逊弹性块存储（Amazon EBS）卷附加的GPU实例。我们可以将模型下载超时和容器启动健康检查保持默认值。有关更多详细信息，请参见CreateEndpointConfig。

endpoint_config_response = sm_client.create_endpoint_config(
EndpointConfigName=endpoint_config_name,
    ProductionVariants=[{
            "VariantName": "AllTraffic",
            "ModelName": model_name,
            "InstanceType": "ml.g4dn.4xlarge",
            "InitialInstanceCount": 1,
            #"VolumeSizeInGB" : 200,
            #"ModelDataDownloadTimeoutInSeconds": 600,
            #"ContainerStartupHealthCheckTimeoutInSeconds": 600,
        },
    ],)'

最后，我们创建一个SageMaker端点：

create_endpoint_response = sm_client.create_endpoint(
EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name)

调用模型端点

这是一个生成模型，因此我们将input_ids和attention_mask作为负载的一部分传递给模型。以下代码展示了如何创建张量：

tokenizer("This is a sample", padding="max_length", max_length=max_seq_len)

我们现在通过确保数据类型与config.pbtxt中配置的一致来创建适当的负载。这也给了我们包含批处理维度的张量，这是Triton所期望的。我们使用JSON格式调用模型。Triton还为该模型提供了一种本地二进制调用方法。

response = runtime_sm_client.invoke_endpoint(
    EndpointName=endpoint_name,
    ContentType="application/octet-stream",
    Body=json.dumps(payload),
    TargetModel=f"{tar_file_name}",
    # TargetModel=f"roberta-large-v0.tar.gz",
)

请注意上述代码中的TargetModel参数。我们将要调用的模型名称作为请求头发送，因为这是一个多模型端点，因此我们可以通过更改此参数在已部署的推理端点上运行多个模型。这展示了多模型端点的强大功能！

要输出响应，我们可以使用以下代码：

import numpy as np

resp_bin = response["Body"].read().decode("utf8")
# -- keys are -- "outputs":[{"name":"1550","datatype":"FP32","shape":[1,768],"data": [0.0013,0,3433...]}]
for data in json.loads(resp_bin)["outputs"]:
    shape_1 = list(data["shape"])
    dat_1 = np.array(data["data"])
    dat_1.resize(shape_1)
    print(f"Data Outputs recieved back :Shape:{dat_1.shape}")

ONNX用于性能调优

ONNX后端使用C++竞技场内存分配。竞技场分配是一项仅适用于C ++的功能，可帮助您优化内存使用并提高性能。内存分配和释放在协议缓冲区代码中占据了CPU时间的相当大的一部分。默认情况下，新对象创建为每个对象，其每个子对象以及多个字段类型，例如字符串，执行堆分配。这些分配在解析消息和在内存中构建新消息时批量发生，并在释放消息及其子对象树时发生关联的释放。

基于Arena的分配旨在降低这种性能成本。使用Arena分配，新的对象将从一个称为Arena的大块预先分配的内存中分配。对象可以通过丢弃整个Arena一次性释放，理想情况下不运行任何包含对象的析构函数（尽管在必要时Arena仍然可以维护析构函数列表）。这通过将对象分配减少到简单的指针增量使对象分配更快，并几乎免费地进行解除分配。Arena分配还提供了更大的缓存效率：当解析消息时，它们更有可能在连续内存中分配，这使得遍历消息更有可能击中热的缓存行。基于Arena的分配的缺点是C ++堆内存将被过度分配，并在对象被解除分配后保持分配状态。这可能导致内存不足或高CPU内存使用率。为了实现最佳效果，我们使用Triton和ONNX提供的以下配置：

  • arena_extend_strategy – 此参数是指与模型大小相关的内存Arena增长策略。我们建议将值设置为1（= kSameAsRequested），这不是默认值。原因如下：默认Arena扩展策略（kNextPowerOfTwo）的缺点是它可能分配比所需更多的内存，这可能是浪费的。正如名称所示，kNextPowerOfTwo（默认值）通过2的幂扩展Arena，而kSameAsRequested扩展一个与每次分配请求相同大小的大小。当您事先知道预期的内存使用情况时，kSameAsRequested适用于高级配置。在我们的测试中，因为我们知道模型的大小是一个常量值，所以我们可以安全地选择kSameAsRequested。

  • gpu_mem_limit – 我们将值设置为CUDA内存限制。要使用所有可能的内存，请传递最大的size_t。如果未指定任何内容，则默认为SIZE_MAX。我们建议将其保持为默认值。

  • enable_cpu_mem_arena – 这使能了CPU上的内存Arena。Arena可以为未来使用预先分配内存。如果您不希望使用它，请将此选项设置为false。默认值为True。如果禁用Arena，堆内存分配将花费时间，因此推断延迟将增加。在我们的测试中，我们将其保留为默认值。

  • enable_mem_pattern – 此参数是指基于输入形状的内部内存分配策略。如果形状是常量，则可以启用此参数以为未来生成内存模式并节省一些分配时间，使其更快。使用1启用内存模式，使用0禁用。建议在预计输入特征相同时将其设置为1。默认值为1。

  • do_copy_in_default_stream – 在ONNX中的CUDA执行提供程序的上下文中，计算流是在GPU上异步运行的CUDA操作序列。ONNX运行时根据它们的依赖关系在不同的流中调度操作，这有助于最小化GPU的空闲时间并实现更好的性能。我们建议使用默认设置1来使用相同的流进行复制和计算；但是，您可以使用0来使用不同的流进行复制和计算，这可能会导致设备流水线化两个活动。在我们对ResNet50模型的测试中，我们使用了0和1，但在性能和GPU设备内存消耗方面两者之间没有任何明显的差异。

  • 图形优化 – Triton的ONNX后端支持几个参数，可帮助微调部署模型的模型大小以及运行时性能。当将模型转换为ONNX表示（下图中IR阶段中的第一个框）时，ONNX运行时提供了三个级别的图形优化：基本，扩展和布局优化。您可以通过在模型配置文件中添加以下参数来激活所有级别的图形优化：

    optimization {
  graph : {
    level : 1
}}

  • cudnn_conv_algo_search – 因为我们在测试中使用基于CUDA的Nvidia GPU，对于ResNet50模型的计算机视觉用例，我们可以在下图中第四层使用基于CUDA的执行提供程序优化，使用 cudnn_conv_algo_search 参数。默认选项是详尽（0），但是当我们将此配置更改为1 - HEURISTIC时，我们看到模型稳定状态下的延迟减少到160毫秒。发生这种情况的原因是因为ONNX运行时调用了更轻量级的cudnnGetConvolutionForwardAlgorithm_v7正向传递，因此在足够的性能下减少了延迟。

  • 运行模式 – 下一步是在下图的第5层中选择正确的execution_mode。此参数控制您是否要按顺序或并行运行图中的运算符。通常，当模型有许多分支时，将此选项设置为ExecutionMode.ORT_PARALLEL（1）将为您提供更好的性能。在图中有许多分支的情况下，将运行模式设置为并行将有助于提高性能。默认模式是顺序的，因此您可以启用此选项以满足您的需求。

    parameters { key: "execution_mode" value: { string_value: "1" } }

要深入了解ONNX中进行性能调优的机会，请参考以下图片。

来源：https://static.linaro.org/connect/san19/presentations/san19-211.pdf

基准测试数字和性能调优

通过在ResNet50模型的测试中打开图优化、cudnn_conv_algo_search和并行运行模式参数，我们看到ONNX模型图的冷启动时间从4.4秒降至1.61秒。下面是完整模型配置文件的示例，位于以下笔记本的ONNX配置部分。

测试基准结果如下：

  • PyTorch – 176 毫秒，冷启动6秒
  • TensorRT – 174 毫秒，冷启动4.5秒
  • ONNX – 168 毫秒，冷启动4.4秒

以下图表可视化这些指标。

此外，在我们的计算机视觉用例测试中，考虑使用Triton提供的HTTP客户端以二进制格式发送请求有效负载，因为它显着提高了模型调用延迟。

SageMaker在Triton上为ONNX公开的其他参数如下：

  • 动态批处理 – 动态批处理是Triton的一个功能，允许服务器通过动态创建批处理请求。创建请求的批处理通常会增加吞吐量。动态批处理应用于无状态模型。动态创建的批次分配给为该模型配置的所有模型实例。
  • 最大批处理大小 – max_batch_size属性表示模型支持的Triton可利用的批处理类型的最大批处理大小。如果模型的批处理维度是第一维，并且模型的所有输入和输出都具有此批处理维度，则Triton可以使用其动态批处理程序或序列批处理程序自动使用模型进行批处理。在这种情况下，max_batch_size应设置为大于或等于1的值，这表示Triton应使用的模型的最大批处理大小。
  • 默认最大批处理大小 – 默认最大批处理大小值用于当找不到其他值时自动完成时的max_batch_size。如果自动完成已确定模型能够批处理请求，并且模型配置中的max_batch_size为0或省略max_batch_size，则onnxruntime后端将将模型的max_batch_size设置为此默认值。如果max_batch_size大于1且未提供调度程序，则将使用动态批处理程序。默认最大批处理大小为4。

清理

确保在运行笔记本后删除模型、模型配置和模型端点。在GitHub repo的示例笔记本的末尾提供了执行此操作的步骤。

结论

在本文中，我们深入探讨了Triton Inference Server在SageMaker上支持的ONNX后端。此后端提供了GPU加速的ONNX模型。有许多选项可考虑以获取最佳推理性能，例如批处理大小、数据输入格式和其他可调整的因素。SageMaker允许您使用单一模型和多个模型端点使用此功能。MME允许更好地平衡性能和成本节省。有关使用GPU的MME支持的入门，请参见在一个端点后面的一个容器中托管多个模型。

我们邀请您在SageMaker中尝试Triton推理服务器容器，并在评论中分享您的反馈和问题。

Checks，谷歌基于人工智能技术打造的隐私平台
在实验室中测试谷歌的功能和产品
介绍谷歌的安全人工智能框架
认识 Instagram 的 AI 聊天机器人-您的新朋友
苹果首席执行官蒂姆·库克支持ChatGPT并讨论人...
“我们会解决幻觉问题，” Sam Altman说道
AI 新闻

AI 新闻

获取人工智能领域的最新新闻和更新