如何在Vertex AI中调用自定义预测逻辑部署推理服务

核心实现思路 #

要实现自定义推理逻辑响应Vertex AI Endpoint请求，你需要使用Vertex AI自定义预测例程功能，而非直接使用平台提供的预构建容器部署，该方案支持完全自定义模型加载、推理前后处理、日志上报逻辑，可以完全匹配你的需求场景。

分步实现步骤 #

1. 准备自定义预测代码包 #

先搭建预测代码的标准目录结构，核心是实现predictor.py文件，自定义两类核心逻辑：

• 模型加载逻辑：实现load方法，从GCS拉取你上传的image_quality.pb文件，完成TensorFlow模型的反序列化
• 推理执行逻辑：实现predict方法，处理传入的GCS图像路径列表，拉取图像做预处理、执行推理、返回结果
  标准目录结构如下：

custom_predictor/
├── predictor.py
└── requirements.txt

requirements.txt中添加你依赖的第三方库，比如tensorflow、google-cloud-storage、tensorboard等。

2. 构建自定义预测镜像 #

你可以基于Vertex AI官方提供的预构建基础镜像打包自定义镜像，也可以自行编写Dockerfile完成镜像构建。镜像中需要配置好模型拉取路径、预测入口，同时添加TensorBoard日志写入的相关配置，将生成的日志输出到你指定的GCS路径即可。
镜像构建完成后推送到Google Container Registry或者Artifact Registry，供Vertex AI部署时调用。

3. 导入模型并部署到Endpoint #

在Vertex AI控制台或者用SDK创建模型资源时，指定你刚才构建的自定义镜像地址，以及模型文件在GCS的存储路径。模型创建完成后即可部署到Endpoint，部署时可以按需选择机器配置、弹性扩缩容规则。

4. 适配GCS图像列表推理逻辑 #

在predictor.py的predict方法中添加GCS客户端逻辑，收到请求传入的GCS图像路径列表后，逐个从GCS拉取图像，完成resize、归一化等你需要的预处理操作，再输入到反序列化后的模型中执行推理，推理完成后将结果格式化返回即可。

5. 配置TensorBoard日志采集 #

在预测代码中初始化TensorBoard SummaryWriter，把你需要记录的指标、中间结果写入到本地路径，同时配置Vertex AI部署日志同步规则，将本地生成的TensorBoard日志文件自动同步到指定的GCS路径，后续直接用Vertex AI的TensorBoard实例关联该GCS路径，即可查看所有推理日志和指标。

核心代码示例 #

predictor.py 核心逻辑示例 #

import tensorflow as tf
from google.cloud import storage
from torch.utils.tensorboard import SummaryWriter
import os

class ImageQualityPredictor:
    def load(self, artifacts_uri: str):
        # 从GCS拉取pb模型并反序列化
        storage_client = storage.Client()
        bucket_name = artifacts_uri.split("/")[2]
        blob_path = "/".join(artifacts_uri.split("/")[3:]) + "image_quality.pb"
        blob = storage_client.bucket(bucket_name).blob(blob_path)
        blob.download_to_filename("/tmp/image_quality.pb")
        
        # 加载TensorFlow pb模型
        self.model = tf.saved_model.load("/tmp/")
        # 初始化TensorBoard writer，日志会自动同步到绑定的GCS路径
        self.writer = SummaryWriter(log_dir=os.environ.get("AIP_TENSORBOARD_LOG_DIR", "/tmp/tb_logs"))

    def predict(self, instances):
        # instances参数即为传入的GCS图像路径列表
        results = []
        storage_client = storage.Client()
        for idx, img_gcs_path in enumerate(instances):
            # 拉取GCS存储的图像
            bucket_name = img_gcs_path.split("/")[2]
            img_blob_path = "/".join(img_gcs_path.split("/")[3:])
            img_blob = storage_client.bucket(bucket_name).blob(img_blob_path)
            img_bytes = img_blob.download_as_bytes()
            # 图像预处理逻辑可根据你的需求调整
            img = tf.io.decode_jpeg(img_bytes, channels=3)
            img = tf.image.resize(img, (224, 224)) / 255.0
            img = tf.expand_dims(img, 0)
            # 执行模型推理
            score = self.model(img).numpy()[0][0]
            results.append({"image_path": img_gcs_path, "quality_score": float(score)})
            # 写入指标到TensorBoard
            self.writer.add_scalar("quality_score", float(score), idx)
        self.writer.flush()
        return results

Endpoint 调用示例 #

from google.cloud import aiplatform

aiplatform.init(project="你的GCP项目ID", location="部署区域")
endpoint = aiplatform.Endpoint("你的Endpoint资源ID")
# 传入GCS图像路径列表发起推理请求
predictions = endpoint.predict(
    instances=[
        "gs://图像存储桶名称/img1.jpg",
        "gs://图像存储桶名称/img2.jpg"
    ]
)
print(predictions.predictions)

内容的提问来源于stack exchange，提问作者kaniska Mandal