如何将Vertex AI自定义训练的模型保存至Cloud Storage？

现有一段可生成训练好的wine_mdl模型的Python代码：

from sklearn import datasets, svm
import logging

logging.info('Loading a Scikit-learn toy dataset')
wine_df=datasets.load_wine(as_frame=True)

logging.info('Model training')
wine_mdl=svm.SVC(probability=True)
wine_mdl.fit(X=wine_df.data, y=wine_df.target)

# logging.info('Saving the trained model')

需要在Vertex AI Training上运行自定义任务，满足以下要求：

• 不使用托管数据集
• 使用scikit-learn预构建容器
• 代码符合Google Cloud文档规范
• 模型输出目录为Google Storage路径gs://my-bucket-for-vertexai/my-output-directory

请问需要在上述代码中添加哪些代码行，才能将训练好的模型导出至指定的Google Storage存储库？

需要添加以下内容完成模型导出：

1. 导入必要依赖库：添加joblib（用于序列化scikit-learn模型）和os（读取环境变量）的导入语句

import joblib
import os

2. 获取Vertex AI提供的模型输出路径：从环境变量中读取预构建容器自动注入的模型存储路径，无需硬编码云存储地址

model_dir = os.environ.get('AIP_MODEL_DIR')

3. 序列化并保存模型：使用joblib将模型写入指定路径，预构建容器会自动将本地文件同步到对应的Google Storage位置

logging.info('Saving the trained model')
joblib.dump(wine_mdl, os.path.join(model_dir, 'model.joblib'))

修改后的完整代码 #

from sklearn import datasets, svm
import logging
import joblib
import os

logging.info('Loading a Scikit-learn toy dataset')
wine_df=datasets.load_wine(as_frame=True)

logging.info('Model training')
wine_mdl=svm.SVC(probability=True)
wine_mdl.fit(X=wine_df.data, y=wine_df.target)

logging.info('Saving the trained model')
model_dir = os.environ.get('AIP_MODEL_DIR')
joblib.dump(wine_mdl, os.path.join(model_dir, 'model.joblib'))

关键说明 #

• Vertex AI的scikit-learn预构建容器会自动设置AIP_MODEL_DIR环境变量，其值就是任务配置中指定的Google Storage路径，无需手动处理云存储的文件上传逻辑
• joblib是scikit-learn官方推荐的模型序列化工具，相比pickle更适配sklearn模型的存储需求

内容的提问来源于stack exchange，提问作者Maxime Oriol