如何在Sagemaker中将SKLearnProcessor拟合后的转换器保存到S3复用

SKLearnProcessor本身支持自动同步容器指定目录下的文件到S3，不需要自行编写S3上传逻辑，你只需要按以下步骤配置即可：

步骤1：定义Processor时配置输出映射 #

初始化SKLearnProcessor后，在提交任务时通过ProcessingOutput指定容器内的存储路径和对应的S3目标路径，示例代码如下：

from sagemaker.sklearn.processing import SKLearnProcessor
from sagemaker.processing import ProcessingInput, ProcessingOutput

# 初始化SKLearnProcessor，替换为你实际使用的版本、角色等参数
sklearn_processor = SKLearnProcessor(
    framework_version="1.2-1",
    role=your_sagemaker_execution_role,
    instance_type="ml.m5.xlarge",
    instance_count=1,
)

# 提交处理任务
sklearn_processor.run(
    code="your_processing_script.py", # 你的数据处理脚本路径
    inputs=[
        ProcessingInput(
            source="s3://你的输入数据集S3路径/",
            destination="/opt/ml/processing/input"
        )
    ],
    outputs=[
        ProcessingOutput(
            output_name="fitted_transformer",
            # 容器内你用来存放转换器的本地路径，可自定义
            source="/opt/ml/processing/transformer",
            # 转换器最终要存储到的S3路径，可自定义
            destination="s3://你的S3桶名/转换器存储路径/"
        )
    ]
)

步骤2：在处理脚本中保存转换器到指定路径 #

在你的数据处理脚本your_processing_script.py中，拟合完成转换器后，直接将其序列化保存到上一步指定的容器本地路径即可，SageMaker会在任务结束后自动同步该路径下的所有文件到你配置的S3路径，示例代码如下：

import os
import joblib
import pandas as pd
from sklearn.preprocessing import StandardScaler # 替换为你实际用的转换器类

# 读取输入数据
input_df = pd.read_csv("/opt/ml/processing/input/input_data.csv")

# 拟合转换器
transformer = StandardScaler()
transformer.fit(input_df[["数值列1", "数值列2"]])

# 保存转换器到指定输出路径
save_dir = "/opt/ml/processing/transformer"
os.makedirs(save_dir, exist_ok=True)
joblib.dump(transformer, os.path.join(save_dir, "fitted_transformer.joblib"))

步骤3：其他脚本复用转换器 #

后续需要使用转换器时，直接从对应的S3路径下载加载即可，示例代码如下：

import boto3
import joblib
from io import BytesIO

s3_client = boto3.client("s3")
resp = s3_client.get_object(
    Bucket="你的S3桶名",
    Key="转换器存储路径/fitted_transformer.joblib"
)
transformer = joblib.load(BytesIO(resp["Body"].read()))

# 直接使用转换器处理新数据
new_data = pd.DataFrame({"数值列1": [1,2,3], "数值列2": [4,5,6]})
transformed_data = transformer.transform(new_data)

• 运行SKLearnProcessor的sklearn版本需要和后续加载转换器的环境sklearn版本保持一致，避免序列化兼容性问题
• 如果使用自定义转换器，需要保证自定义类的代码在加载环境中也存在，否则会加载失败

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