如何在SageMaker Pipeline中为批量转换任务添加自定义推理脚本

要让批量转换任务使用自定义推理脚本，核心是在创建模型（或模型包）阶段就把脚本关联进去——因为批量转换依赖已创建的模型资源，无法在TransformStep阶段临时添加脚本。具体操作如下：

• 准备自定义推理脚本
  编写inference.py（命名可自定义，行业通常用这个名称），实现SageMaker要求的核心函数：model_fn（加载训练好的模型）、input_fn（解析批量转换的输入数据）、predict_fn（执行推理计算）、output_fn（格式化推理结果）。如果有依赖的辅助文件，可放在同一个目录下。

• 创建模型时绑定脚本
  在构建Model对象（或对应框架的Model类，比如TensorFlowModel、PyTorchModel）时，通过参数指定推理脚本和依赖目录，再基于这个Model创建Pipeline的CreateModelStep。示例代码如下：

# 导入必要的库
from sagemaker.model import Model
from sagemaker.workflow.step_collections import CreateModelStep
from sagemaker.transformer import Transformer
from sagemaker.inputs import TransformInput
from sagemaker.workflow.steps import TransformStep

# 配置包含自定义推理脚本的模型
custom_model = Model(
    image_uri="你的模型容器镜像URI",
    model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,  # 训练阶段产出的模型文件路径
    entry_point="inference.py",  # 自定义推理脚本的文件名
    source_dir="./code",  # 脚本及依赖文件所在的本地目录（单脚本可忽略该参数）
    role="你的SageMaker执行角色ARN",
    sagemaker_session=sagemaker_session
)

# 创建Pipeline中的CreateModel步骤
step_create_model = CreateModelStep(
    name="AbaloneCreateModel",
    model=custom_model,
    model_name="AbaloneModel"
)

# 后续的批量转换步骤保持原有逻辑即可
transformer = Transformer(
    model_name=step_create_model.properties.ModelName,
    instance_type="ml.m5.xlarge",
    instance_count=1,
    output_path=f"s3://{default_bucket}/AbaloneTransform",
)

step_transform = TransformStep(
    name="AbaloneTransform", transformer=transformer, inputs=TransformInput(data=batch_data)
)

• 如果使用模型包（ModelPackage）
  需在创建ModelPackage时，将推理脚本及依赖打包到模型包的资产中，通过ModelPackage类的参数或InferenceSpecification配置项，指定脚本路径、容器信息等，确保模型包包含推理所需的全部资源。

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