导出模型工件以进行预测和说明

Vertex AI 提供预构建容器，以便使用以下机器学习 (ML) 框架训练的模型提供预测和说明：

  • TensorFlow
  • PyTorch
  • XGBoost
  • scikit-learn

如需使用这些预构建容器之一，您必须将模型保存为符合预构建容器要求的一个或多个模型工件。无论您的模型工件是否是在 Vertex AI 上创建的，这些要求都适用。

如果您要使用自定义容器执行预测，则无需遵守本页面中的要求，但仍然可以将其用作准则。

导出到预构建容器时特定于框架的要求

根据您打算用于预测的机器学习框架，您必须以不同格式导出模型工件。以下部分介绍了每个机器学习框架可接受的模型格式。

TensorFlow

如果您使用 TensorFlow 训练模型，请将模型导出为 TensorFlow SavedModel 目录。

您可以通过多种方法从 TensorFlow 训练代码中导出 SavedModel。以下列表介绍了适用于各种 TensorFlow API 的多种不同方式：

如需使用此工件进行预测，请使用与您用于训练的 TensorFlow 版本匹配的用于预测的预构建容器创建 Model。

适用于 Vertex Explainable AI 的 TensorFlow

如果您想从使用 TensorFlow 预构建容器来执行预测的 Model 获取说明，请阅读导出适用于 Vertex Explainable AI 的 TensorFlow 模型的其他要求。

为 Tensorflow 启用服务器端请求批处理

如果要为使用 TensorFlow 预构建容器提供预测的 Model 启用请求批处理，请在 saved_model.pb 文件所在的 gcs 目录中添加 config/batching_parameters_config。如需配置批处理配置文件，请参阅 Tensorflow 官方文档。

PyTorch

您必须使用 Torch 模型归档程序创建归档文件，以打包模型工件（包括默认或自定义处理程序）。预构建的 PyTorch 映像要求归档命名为 model.mar，因此请确保将模型名称设置为“模型”。

如需了解如何优化使用 TorchServe 提供的 PyTorch 模型的内存用量、延迟时间或吞吐量，请参阅 PyTorch 性能指南。

XGBoost

如果使用 XGBoost 训练模型，则可以通过以下三种方式之一导出经过训练的模型：

  • 使用 xgboost.Booster 的 save_model 方法导出名为 model.bst 的文件。
  • 使用 joblib 库导出名为 model.joblib 的文件。
  • 使用 Python 的 pickle 模块导出名为 model.pkl 的文件。

模型工件的文件名必须与其中一个选项完全匹配。

以下标签式示例分别演示了如何使用这三种方式来训练和导出模型：

xgboost.Booster

import os

from google.cloud import storage
from sklearn import datasets
import xgboost as xgb

digits = datasets.load_digits()
dtrain = xgb.DMatrix(digits.data, label=digits.target)
bst = xgb.train({}, dtrain, 20)

artifact_filename = 'model.bst'

# Save model artifact to local filesystem (doesn't persist)
local_path = artifact_filename
bst.save_model(local_path)

# Upload model artifact to Cloud Storage
model_directory = os.environ['AIP_MODEL_DIR']
storage_path = os.path.join(model_directory, artifact_filename)
blob = storage.blob.Blob.from_string(storage_path, client=storage.Client())
blob.upload_from_filename(local_path)

joblib

import os

from google.cloud import storage
from sklearn import datasets
import joblib
import xgboost as xgb

digits = datasets.load_digits()
dtrain = xgb.DMatrix(digits.data, label=digits.target)
bst = xgb.train({}, dtrain, 20)

artifact_filename = 'model.joblib'

# Save model artifact to local filesystem (doesn't persist)
local_path = artifact_filename
joblib.dump(bst, local_path)

# Upload model artifact to Cloud Storage
model_directory = os.environ['AIP_MODEL_DIR']
storage_path = os.path.join(model_directory, artifact_filename)
blob = storage.blob.Blob.from_string(storage_path, client=storage.Client())
blob.upload_from_filename(local_path)

pickle

import os
import pickle

from google.cloud import storage
from sklearn import datasets
import xgboost as xgb

digits = datasets.load_digits()
dtrain = xgb.DMatrix(digits.data, label=digits.target)
bst = xgb.train({}, dtrain, 20)

artifact_filename = 'model.pkl'

# Save model artifact to local filesystem (doesn't persist)
local_path = artifact_filename
with open(local_path, 'wb') as model_file:
  pickle.dump(bst, model_file)

# Upload model artifact to Cloud Storage
model_directory = os.environ['AIP_MODEL_DIR']
storage_path = os.path.join(model_directory, artifact_filename)
blob = storage.blob.Blob.from_string(storage_path, client=storage.Client())
blob.upload_from_filename(local_path)

如需使用此工件进行预测，请使用与您用于训练的 XGBoost 版本匹配的用于预测的预构建容器创建 Model。

scikit-learn

如果使用 scikit-learn 训练模型，则可以通过以下两种方式之一导出模型：

  • 使用 joblib 库导出名为 model.joblib 的文件。
  • 使用 Python 的 pickle 模块导出名为 model.pkl 的文件。

模型工件的文件名必须与其中一个选项完全匹配。 您可以导出标准 scikit-learn Estimator 或 scikit-learn 流水线。

以下标签式示例分别演示了如何使用这两种方式来训练和导出模型：

joblib

import os

from google.cloud import storage
from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier
import joblib

digits = datasets.load_digits()
classifier = RandomForestClassifier()
classifier.fit(digits.data, digits.target)

artifact_filename = 'model.joblib'

# Save model artifact to local filesystem (doesn't persist)
local_path = artifact_filename
joblib.dump(classifier, local_path)

# Upload model artifact to Cloud Storage
model_directory = os.environ['AIP_MODEL_DIR']
storage_path = os.path.join(model_directory, artifact_filename)
blob = storage.blob.Blob.from_string(storage_path, client=storage.Client())
blob.upload_from_filename(local_path)

pickle

import os
import pickle

from google.cloud import storage
from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier

digits = datasets.load_digits()
classifier = RandomForestClassifier()
classifier.fit(digits.data, digits.target)

artifact_filename = 'model.pkl'

# Save model artifact to local filesystem (doesn't persist)
local_path = artifact_filename
with open(local_path, 'wb') as model_file:
  pickle.dump(classifier, model_file)

# Upload model artifact to Cloud Storage
model_directory = os.environ['AIP_MODEL_DIR']
storage_path = os.path.join(model_directory, artifact_filename)
blob = storage.blob.Blob.from_string(storage_path, client=storage.Client())
blob.upload_from_filename(local_path)

如需使用此工件进行预测，请使用与您用于训练的 scikit-learn 版本匹配的用于预测的预构建容器创建 Model。

后续步骤