将PyTorch稳定扩散模型部署为Vertex AI端点¶

创建于：2023年11月14日 | 最后更新：2023年11月14日 | 最后验证：未验证

部署大型模型，如Stable Diffusion，可能具有挑战性且耗时。

在本教程中，我们将展示如何通过利用Vertex AI来简化PyTorch Stable Diffusion模型的部署。

PyTorch 是 Stability AI 在 Stable Diffusion v1.5 中使用的框架。Vertex AI 是一个完全托管的机器学习平台，提供工具和基础设施，旨在帮助机器学习从业者利用 PyTorch 等开源框架加速和扩展生产中的机器学习。

在四个步骤中，您可以部署一个PyTorch Stable Diffusion模型（v1.5）。

在Vertex AI端点上部署您的Stable Diffusion模型可以通过四个步骤完成：

创建一个自定义的TorchServe处理程序。
将模型工件上传到Google云存储（GCS）。
使用模型工件和预构建的PyTorch容器镜像创建一个Vertex AI模型。
将Vertex AI模型部署到端点上。

让我们更详细地看一下每个步骤。您可以使用Notebook示例来跟随并实现这些步骤。

注意：请记住，此配方需要可计费的Vertex AI，如笔记本示例中更详细解释的那样。

创建一个自定义的TorchServe处理程序¶

TorchServe 是一个用于服务 PyTorch 模型的简单且灵活的工具。部署到 Vertex AI 的模型使用 TorchServe 来处理请求并返回模型的响应。您必须创建一个自定义的 TorchServe 处理程序，并将其包含在上传到 Vertex AI 的模型工件中。将处理程序文件与其他模型工件一起放在目录中，例如：model_artifacts/handler.py。

创建处理程序文件后，您必须将处理程序打包为模型归档（MAR）文件。输出文件必须命名为model.mar。

!torch-model-archiver \
-f \
--model-name <your_model_name> \
--version 1.0 \
 --handler model_artifacts/handler.py \
--export-path model_artifacts

将模型工件上传到Google云存储(GCS)¶

在这一步中，我们正在上传模型工件到GCS，例如模型文件或处理程序。将工件存储在GCS上的优势是您可以在一个中央存储桶中跟踪这些工件。

BUCKET_NAME = "your-bucket-name-unique"  # @param {type:"string"}
BUCKET_URI = f"gs://{BUCKET_NAME}/"

# Will copy the artifacts into the bucket
!gsutil cp -r model_artifacts $BUCKET_URI

使用模型工件和预构建的PyTorch容器镜像创建Vertex AI模型¶

一旦你将模型工件上传到GCS存储桶中，你就可以将你的PyTorch模型上传到 Vertex AI 模型注册表。在Vertex AI 模型注册表中，你可以概览你的模型，从而更好地组织、跟踪和训练新版本。为此，你可以使用 Vertex AI SDK 和这个预构建的PyTorch容器。

from google.cloud import aiplatform as vertexai
PYTORCH_PREDICTION_IMAGE_URI = (
    "us-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.1-12:latest"
)
MODEL_DISPLAY_NAME = "stable_diffusion_1_5-unique"
MODEL_DESCRIPTION = "stable_diffusion_1_5 container"

vertexai.init(project='your_project', location='us-central1', staging_bucket=BUCKET_NAME)

model = aiplatform.Model.upload(
    display_name=MODEL_DISPLAY_NAME,
    description=MODEL_DESCRIPTION,
    serving_container_image_uri=PYTORCH_PREDICTION_IMAGE_URI,
    artifact_uri=BUCKET_URI,
)

将Vertex AI模型部署到端点¶

一旦模型上传到Vertex AI模型注册表，您就可以将其部署到Vertex AI端点。为此，您可以使用控制台或Vertex AI SDK。在此示例中，您将在NVIDIA Tesla P100 GPU和n1-standard-8机器上部署模型。您可以指定您的机器类型。

endpoint = aiplatform.Endpoint.create(display_name=ENDPOINT_DISPLAY_NAME)

model.deploy(
    endpoint=endpoint,
    deployed_model_display_name=MODEL_DISPLAY_NAME,
    machine_type="n1-standard-8",
    accelerator_type="NVIDIA_TESLA_P100",
    accelerator_count=1,
    traffic_percentage=100,
    deploy_request_timeout=1200,
    sync=True,
)

如果您按照这个 notebook 操作，您也可以使用Vertex AI SDK获取在线预测，如下面的代码片段所示。

instances = [{"prompt": "An examplePup dog with a baseball jersey."}]
response = endpoint.predict(instances=instances)

with open("img.jpg", "wb") as g:
    g.write(base64.b64decode(response.predictions[0]))

display.Image("img.jpg")

使用模型工件和预构建的PyTorch容器镜像创建一个Vertex AI模型

将PyTorch稳定扩散模型部署为Vertex AI端点¶

创建一个自定义的TorchServe处理程序¶

将模型工件上传到Google云存储(GCS)¶

使用模型工件和预构建的PyTorch容器镜像创建Vertex AI模型¶

将Vertex AI模型部署到端点¶

更多资源¶