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

Created On: Nov 14, 2023 | Last Updated: Nov 14, 2023 | Last Verified: Not Verified

部署像稳定扩散这样的大型模型可能具有挑战性且耗时。

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

PyTorch是Stability AI在稳定扩散v1.5中使用的框架。Vertex AI是一个完全托管的机器学习平台，提供工具和基础设施，旨在帮助ML从业者加速和扩展生产中的机器学习，同时受益于像PyTorch这样的开源框架。

通过四个步骤，您可以部署PyTorch稳定扩散模型(v1.5)。

在Vertex AI端点上部署稳定扩散模型可以通过四个步骤完成:

创建一个自定义TorchServe处理器。
将模型工件上传到Google Cloud Storage (GCS)。
使用模型工件和预构建的PyTorch容器镜像创建一个Vertex AI模型。
将Vertex AI模型部署到端点上。

让我们详细了解每个步骤。您可以通过 Notebook示例进行跟随和实施。

注意：请记住，此教程需要可收费的Vertex AI，如Notebook示例中更详细地阐述。

创建一个自定义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 Cloud Storage (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 模型注册表 <https://cloud.google.com/vertex-ai/docs/model-registry/introduction>`__。通过Vertex AI模型注册表，您可以一览您的模型，以便更好地组织、追踪和训练新版本。为此，您可以使用`Vertex AI SDK <https://cloud.google.com/vertex-ai/docs/python-sdk/use-vertex-ai-python-sdk>`__以及这个`预构建PyTorch容器 <https://cloud.google.com/blog/products/ai-machine-learning/prebuilt-containers-with-pytorch-and-vertex-ai>`__。

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,
)

如果您遵循这篇`笔记本 <https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/torchserve/dreambooth_stablediffusion.ipynb>`__，还可以像以下代码片段展示的那样，使用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 Cloud Storage (GCS)¶

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

在端点上部署Vertex AI模型¶

更多资源¶

文档

教程

资源