如何将Jupyter Notebook训练代码转换为SageMaker Pipeline步骤？

核心思路 #

SageMaker Pipeline的训练步骤本质是把Notebook中手动执行的训练逻辑，封装成可复用、可编排的Pipeline组件。核心是用TrainingStep包装你的HuggingFace Estimator，同时通过Pipeline参数实现配置的灵活调整。

具体转换步骤 #

1. 梳理Notebook中的核心训练逻辑

先提取你已验证通过的关键部分：

• 训练脚本（train.py）及依赖文件
• HuggingFace Estimator的配置（实例类型、框架版本、超参数）
• 数据输入路径

2. 导入Pipeline相关依赖库

import sagemaker
from sagemaker.huggingface import HuggingFace
from sagemaker.workflow.pipeline import Pipeline
from sagemaker.workflow.steps import TrainingStep
from sagemaker.workflow.parameters import ParameterString, ParameterInteger

3. 定义Pipeline可配置参数

把需要灵活调整的配置定义为Pipeline参数，方便后续修改（比如切换实例类型、调整训练轮次）：

# 数据输入路径参数
input_data = ParameterString(
    name="TrainingDataPath",
    default_value="s3://your-bucket/training-data"
)
# 训练实例类型参数
training_instance = ParameterString(
    name="TrainingInstanceType",
    default_value="ml.p3.2xlarge"
)
# 训练轮次参数
epochs = ParameterInteger(
    name="TrainingEpochs",
    default_value=3
)

4. 创建适配Pipeline的HuggingFace Estimator

和Notebook中的Estimator定义逻辑一致，但不要调用.fit()，而是留待Pipeline执行：

huggingface_estimator = HuggingFace(
    entry_point="train.py",  # 复用你已有的训练脚本
    source_dir="./training-scripts",  # 脚本所在本地目录
    instance_type=training_instance,
    instance_count=1,
    role=sagemaker.get_execution_role(),
    transformers_version="4.26.0",
    pytorch_version="1.13.1",
    py_version="py39",
    hyperparameters={
        "epochs": epochs,
        "train_batch_size": 32,
        "model_name": "bert-base-uncased"
    }
)

5. 封装为TrainingStep并组装Pipeline

用TrainingStep包装Estimator，指定训练数据输入，最后组装成完整Pipeline：

# 创建训练步骤
training_step = TrainingStep(
    name="HuggingFaceModelTraining",
    estimator=huggingface_estimator,
    inputs={
        "train": sagemaker.inputs.TrainingInput(
            s3_data=input_data,
            content_type="text/csv"  # 根据你的数据格式调整
        )
    }
)

# 组装Pipeline
pipeline = Pipeline(
    name="HuggingFaceTrainingPipeline",
    parameters=[input_data, training_instance, epochs],
    steps=[training_step]
)

6. 验证并启动Pipeline

# 验证Pipeline定义并上传到SageMaker
pipeline.upsert(role_arn=sagemaker.get_execution_role())

# 启动Pipeline执行
execution = pipeline.start()
# 等待执行完成（可选）
execution.wait()

关键注意事项 #

• 训练脚本train.py无需大幅修改，只需确保通过SageMaker环境变量读取数据（比如os.environ["SM_CHANNEL_TRAIN"]获取训练数据路径）
• 若需要扩展Pipeline（比如添加数据预处理、模型评估、模型注册），可以继续添加ProcessingStep、ModelStep等组件
• 所有参数化配置都通过Parameter类定义，便于后续在SageMaker Studio界面直接调整

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