使用 Amazon SageMaker 多模型终端节点和 TorchServe 在 GPU 上加速 AI 模型，可将推理成本降低高达 75%

作者：James Wu、Ankith Gunapal、Li Ning、Subhash Talluri 和 Saurabh Trikande

多模型终端节点 (MMEs) 是 Amazon SageMaker 的强大功能之一，旨在简化机器学习 (ML) 模型的部署和操作。借助 MMEs，您可以在单个服务容器上托管多个模型，并将所有模型托管在单个终端节点之后。SageMaker 平台会根据流量模式自动管理模型的加载和卸载并扩展资源，从而减轻管理大量模型的操作负担。此功能对于需要加速计算的深度学习和生成式 AI 模型尤其有利。通过资源共享和简化模型管理实现的成本节省，使 SageMaker MMEs 成为您在 AWS 上大规模托管模型的绝佳选择。

最近，生成式 AI 应用引起了广泛关注和想象。客户希望在 GPU 上部署生成式 AI 模型，但同时又对成本很在意。SageMaker MMEs 支持 GPU 实例，是这些类型应用的一个很好的选择。今天，我们很高兴宣布 SageMaker MMEs 支持 TorchServe。这项新的模型服务器支持让您能够充分利用 MMEs 的所有优势，同时仍使用 TorchServe 用户最熟悉的服务堆栈。在本文中，我们将演示如何使用 TorchServe 在 SageMaker MMEs 上托管 Stable Diffusion 和 Segment Anything Model 等生成式 AI 模型，并构建一个语言引导编辑解决方案，帮助艺术家和内容创作者更快地开发和迭代他们的作品。

解决方案概述

语言引导编辑是一个常见的跨行业生成式 AI 用例。通过自动化重复性任务、优化活动和提供超个性化体验，它可以帮助艺术家和内容创作者更高效地工作，满足内容需求。企业可以从提高内容输出量、节省成本、改进个性化和增强客户体验中获益。在本文中，我们将演示如何使用 MME TorchServe 构建语言辅助编辑功能，让您能够擦除图像中任何不需要的对象，并通过提供文本指令修改或替换图像中的任何对象。

每个用例的用户体验流程如下

要移除不需要的对象，请从图像中选择对象以将其高亮显示。此操作会将像素坐标和原始图像发送到生成式 AI 模型，该模型会为对象生成一个分割掩码。确认对象选择正确后，您可以将原始图像和掩码图像发送到第二个模型进行移除。此用户流程的详细说明如下所示。


步骤 1: 从图像中选择一个对象（“狗”）	步骤 2: 确认高亮显示的对象正确	步骤 3: 从图像中擦除对象

要修改或替换对象，请选择并高亮显示所需对象，遵循与上述相同的过程。确认对象选择正确后，您可以通过提供原始图像、掩码和文本提示来修改对象。然后，模型将根据提供的指令更改高亮显示的对象。此第二个用户流程的详细说明如下。


步骤 1: 从图像中选择一个对象（“花瓶”）	步骤 2: 确认高亮显示的对象正确	步骤 3: 提供文本提示（“未来主义花瓶”）以修改对象

为支持此解决方案，我们使用了三个生成式 AI 模型：Segment Anything Model (SAM)、Large Mask Inpainting Model (LaMa) 和 Stable Diffusion Inpaint (SD)。以下是这些模型在用户体验工作流程中的使用方式

移除不需要的对象	修改或替换对象

Segment Anything Model (SAM) 用于生成感兴趣对象的分割掩码。SAM 由 Meta Research 开发，是一个开源模型，可以分割图像中的任何对象。该模型已在名为 SA-1B 的大型数据集上训练，该数据集包含超过 1100 万张图像和 11 亿个分割掩码。有关 SAM 的更多信息，请参阅其网站和研究论文。
LaMa 用于从图像中移除任何不需要的对象。LaMa 是一种生成对抗网络 (GAN) 模型，专门使用不规则掩码填充图像缺失部分。模型架构融合了图像范围的全局上下文以及使用傅里叶卷积的单步架构，使其能够以更快的速度获得最先进的结果。有关 LaMa 的更多详细信息，请访问其网站和研究论文。
Stability AI 的 SD 2 inpaint 模型用于修改或替换图像中的对象。该模型允许我们通过提供文本提示来编辑掩码区域中的对象。inpaint 模型基于文本到图像的 SD 模型，该模型可以使用简单的文本提示创建高质量图像。它提供额外的参数，例如原始图像和掩码图像，支持快速修改和恢复现有内容。要了解更多关于 AWS 上的 Stable Diffusion 模型的信息，请参阅使用 Stable Diffusion 模型创建高质量图像，并使用 Amazon SageMaker 以经济高效的方式进行部署。

所有三个模型都托管在 SageMaker MMEs 上，这减轻了管理多个终端节点的操作负担。除此之外，使用 MME 消除了对某些模型利用不足的担忧，因为资源是共享的。您可以从改进的实例利用率中受益，这最终会节省成本。以下架构图说明了如何使用带有 TorchServe 的 SageMaker MMEs 提供所有三个模型。

flow diagram

我们已在GitHub 仓库中发布了实现此解决方案架构的代码。要继续阅读本文的其余部分，请使用笔记本文件。建议在 SageMaker 笔记本实例上使用 conda_python3 (Python 3.10.10) 内核运行此示例。

扩展 TorchServe 容器

第一步是准备模型托管容器。SageMaker 提供托管的 PyTorch 深度学习容器 (DLC)，您可以使用以下代码片段检索

# Use SageMaker PyTorch DLC as base image
baseimage = sagemaker.image_uris.retrieve(
    framework="pytorch",
    region=region,
    py_version="py310",
    image_scope="inference",
    version="2.0.0",
    instance_type="ml.g5.2xlarge",
)
print(baseimage)

由于模型需要基础 PyTorch DLC 中不包含的资源和额外的软件包，因此需要构建一个 Docker 镜像。然后将此镜像上传到 Amazon Elastic Container Registry (Amazon ECR)，以便我们可以直接从 SageMaker 访问。定制安装的库列在 Docker 文件中

ARG BASE_IMAGE

FROM $BASE_IMAGE

#Install any additional libraries
RUN pip install segment-anything-py==1.0
RUN pip install opencv-python-headless==4.7.0.68
RUN pip install matplotlib==3.6.3
RUN pip install diffusers
RUN pip install tqdm
RUN pip install easydict
RUN pip install scikit-image
RUN pip install xformers
RUN pip install tensorflow
RUN pip install joblib
RUN pip install matplotlib
RUN pip install albumentations==0.5.2
RUN pip install hydra-core==1.1.0
RUN pip install pytorch-lightning
RUN pip install tabulate
RUN pip install kornia==0.5.0
RUN pip install webdataset
RUN pip install omegaconf==2.1.2
RUN pip install transformers==4.28.1
RUN pip install accelerate
RUN pip install ftfy

运行 shell 命令文件在本地构建定制镜像并将其推送到 Amazon ECR

%%capture build_output

reponame = "torchserve-mme-demo"
versiontag = "genai-0.1"

# Build our own docker image
!cd workspace/docker && ./build_and_push.sh {reponame} {versiontag} {baseimage} {region} {account}

准备模型工件

支持 TorchServe 的新 MMEs 的主要区别在于如何准备模型工件。代码仓库为每个模型（models 文件夹）提供一个骨架文件夹，用于存放 TorchServe 所需的文件。我们遵循相同的四步过程来准备每个模型的 .tar 文件。以下代码是 SD 模型骨架文件夹的示例

workspace
|--sd
   |-- custom_handler.py
   |-- model-config.yaml

第一步是在 models 文件夹中下载预训练模型检查点

import diffusers
import torch
import transformers

pipeline = diffusers.StableDiffusionInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-inpainting", torch_dtype=torch.float16
)

sd_dir = "workspace/sd/model"
pipeline.save_pretrained(sd_dir)

下一步是定义一个 custom_handler.py 文件。这是定义模型接收请求时的行为所必需的，例如加载模型、预处理输入和后处理输出。handle 方法是请求的主入口点，它接受一个请求对象并返回一个响应对象。它加载预训练模型检查点，并将 preprocess 和 postprocess 方法应用于输入和输出数据。以下代码片段展示了 custom_handler.py 文件的简单结构。有关更多详细信息，请参阅TorchServe handler API。

def initialize(self, ctx: Context):

def preprocess(self, data):

def inference(self, data):

def handle(self, data, context):
    requests = self.preprocess(data)
    responses = self.inference(requests)

    return responses

TorchServe 所需的最后一个文件是 model-config.yaml。该文件定义了模型服务器的配置，例如 worker 数量和批量大小。配置是针对每个模型，以下代码显示了一个示例配置文件。完整参数列表，请参阅GitHub 仓库。

minWorkers: 1
maxWorkers: 1
batchSize: 1
maxBatchDelay: 200
responseTimeout: 300

最后一步是使用 torch-model-archiver 模块将所有模型工件打包成一个 .tar.gz 文件

!torch-model-archiver --model-name sd --version 1.0 --handler workspace/sd/custom_handler.py --extra-files workspace/sd/model --config-file workspace/sam/model-config.yaml --archive-format no-archive!cd sd && tar cvzf sd.tar.gz .

创建多模型终端节点

创建 SageMaker MME 的步骤与之前相同。在此特定示例中，您使用 SageMaker SDK 启动终端节点。首先定义一个 Amazon Simple Storage Service (Amazon S3) 位置和托管容器。SageMaker 将根据调用模式在此 S3 位置动态加载模型。托管容器是您在之前步骤中构建并推送到 Amazon ECR 的定制容器。见以下代码

# This is where our MME will read models from on S3.
multi_model_s3uri = output_path

然后需要定义一个 MultiDataModel，它捕获模型位置、托管容器和权限访问等所有属性

print(multi_model_s3uri)
model = Model(
    model_data=f"{multi_model_s3uri}/sam.tar.gz",
    image_uri=container,
    role=role,
    sagemaker_session=smsess,
    env={"TF_ENABLE_ONEDNN_OPTS": "0"},
)

mme = MultiDataModel(
    name="torchserve-mme-genai-" + datetime.now().strftime("%Y-%m-%d-%H-%M-%S"),
    model_data_prefix=multi_model_s3uri,
    model=model,
    sagemaker_session=smsess,
)
print(mme)

deploy() 函数创建终端节点配置并托管终端节点

mme.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.2xlarge",
    serializer=sagemaker.serializers.JSONSerializer(),
    deserializer=sagemaker.deserializers.JSONDeserializer(),
)

在我们提供的示例中，我们还展示了如何使用 SDK 列出模型并动态添加新模型。add_model() 函数将您的本地模型 .tar 文件复制到 MME S3 位置

# Only sam.tar.gz visible!
list(mme.list_models())

models = ["sd/sd.tar.gz", "lama/lama.tar.gz"]
for model in models:
    mme.add_model(model_data_source=model)

调用模型

现在我们已将所有三个模型托管在 MME 上，我们可以按顺序调用每个模型来构建我们的语言辅助编辑功能。调用每个模型时，请在 predictor.predict() 函数中提供 target_model 参数。模型名称就是我们上传的模型 .tar 文件的文件名。以下是 SAM 模型的一个示例代码片段，它接收像素坐标、点标签和膨胀核大小，并生成像素位置处对象的分割掩码

img_file = "workspace/test_data/sample1.png"
img_bytes = None

with Image.open(img_file) as f:
    img_bytes = encode_image(f)

gen_args = json.dumps(dict(point_coords=[750, 500], point_labels=1, dilate_kernel_size=15))

payload = json.dumps({"image": img_bytes, "gen_args": gen_args}).encode("utf-8")

response = predictor.predict(data=payload, target_model="/sam.tar.gz")
encoded_masks_string = json.loads(response.decode("utf-8"))["generated_image"]
base64_bytes_masks = base64.b64decode(encoded_masks_string)

with Image.open(io.BytesIO(base64_bytes_masks)) as f:
    generated_image_rgb = f.convert("RGB")
    generated_image_rgb.show()

要从图像中移除不需要的对象，请获取由 SAM 生成的分割掩码，并将其与原始图像一起馈送到 LaMa 模型。以下图像显示了一个示例。


样本图像	SAM 生成的分割掩码	使用 LaMa 擦除狗

要使用文本提示修改或替换图像中的任何对象，请获取 SAM 生成的分割掩码，并将其与原始图像和文本提示一起馈送到 SD 模型，如下例所示。


样本图像	SAM 生成的分割掩码	使用 SD 模型和文本提示进行替换 “长凳上的一只仓鼠”

成本节省

SageMaker MMEs 的优势随着模型整合规模的增加而增加。下表显示了本文中三个模型的 GPU 内存使用情况。它们通过一个 SageMaker MME 部署在一个 g5.2xlarge 实例上。

模型	GPU 内存 (MiB)
Segment Anything Model	3,362
Stable Diffusion Inpaint	3,910
LaMa	852

将三个模型托管在一个终端节点上可以看到成本节省，对于拥有数百或数千个模型的用例，节省的成本会大得多。

例如，考虑 100 个 Stable Diffusion 模型。每个模型自身都可以由一个 ml.g5.2xlarge 终端节点（4 GiB 内存）提供服务，在美国东部（弗吉尼亚北部）区域每实例小时花费 $1.52 美元。使用各自的终端节点提供所有 100 个模型的成本将达到每月 $218,880 美元。使用 SageMaker MME，单个使用 ml.g5.2xlarge 实例的终端节点可以同时托管四个模型。这可将生产推理成本降低 75%，降至每月仅 $54,720 美元。下表总结了此示例中单模型和多模型终端节点之间的差异。假定终端节点配置具有足够的目标模型内存，所有模型加载后的稳态调用延迟将与单模型终端节点的延迟相似。

	单模型终端节点	多模型终端节点
每月终端节点总费用	$218,880	$54,720
终端节点实例类型	ml.g5.2xlarge	ml.g5.2xlarge
CPU 内存容量 (GiB)	32	32
GPU 内存容量 (GiB)	24	24
每小时终端节点费用	$1.52	$1.52
每个终端节点的实例数量	2	2
100 个模型所需的终端节点数量	100	25

清理

完成后，请按照笔记本清理部分中的说明进行操作，删除本文中配置的资源，以避免产生不必要的费用。有关推理实例成本的详细信息，请参阅Amazon SageMaker 定价。

结论

本文展示了如何使用带有 TorchServe 的 SageMaker MMEs 上托管的生成式 AI 模型实现语言辅助编辑功能。我们分享的示例说明了如何利用 SageMaker MMEs 的资源共享和简化模型管理功能，同时仍将 TorchServe 用作我们的模型服务堆栈。我们利用了三个深度学习基础模型：SAM、SD 2 Inpainting 和 LaMa。这些模型使我们能够构建强大的功能，例如擦除图像中任何不需要的对象，并通过提供文本指令修改或替换图像中的任何对象。这些功能可以帮助艺术家和内容创作者更高效地工作，并通过自动化重复性任务、优化活动和提供超个性化体验来满足他们的内容需求。我们邀请您探索本文提供的示例，并在 SageMaker MME 上使用 TorchServe 构建您自己的 UI 体验。

要开始使用，请参阅使用 GPU 支持实例的多模型终端节点支持的算法、框架和实例。

关于作者

	James Wu 是 AWS 的高级 AI/ML 专家解决方案架构师，帮助客户设计和构建 AI/ML 解决方案。James 的工作涵盖广泛的 ML 用例，主要关注计算机视觉、深度学习和在企业范围内扩展 ML。在加入 AWS 之前，James 曾担任超过 10 年的架构师、开发人员和技术领导者，其中包括 6 年的工程领域经验和 4 年的营销与广告行业经验。
	Li Ning 是 AWS 的高级软件工程师，专注于构建大规模 AI 解决方案。作为 TorchServe（一个由 AWS 和 Meta 共同开发的项目）的技术负责人，她热衷于利用 PyTorch 和 AWS SageMaker 帮助客户拥抱 AI，造福社会。在工作之余，Li 喜欢游泳、旅行、关注最新的技术进展以及与家人共度美好时光。
	Ankith Gunapal 是 Meta (PyTorch) 的 AI 合作伙伴工程师。他热衷于模型优化和模型服务，拥有从 RTL 验证、嵌入式软件、计算机视觉到 PyTorch 的丰富经验。他拥有数据科学硕士学位和电信硕士学位。在工作之余，Ankith 也是一位电子舞曲制作人。
	Saurabh Trikande 是 Amazon SageMaker 推理高级产品经理。他热衷于与客户合作，并以普及机器学习为目标。他专注于与部署复杂的 ML 应用、多租户 ML 模型、成本优化以及使深度学习模型部署更易于使用相关的核心挑战。在业余时间，Saurabh 喜欢徒步旅行、了解创新技术、关注 TechCrunch 并与家人共度时光。
	Subhash Talluri 是 Amazon Web Services 电信行业业务部门的首席 AI/ML 解决方案架构师。他一直致力于为全球电信客户和合作伙伴开发创新的 AI/ML 解决方案。他将工程和计算机科学的跨学科专业知识相结合，通过基于云优化的架构在 AWS 上帮助构建可扩展、安全且合规的 AI/ML 解决方案。