使用自定义Docker镜像时，无法向SageMaker Estimator传递超参数

当在Amazon SageMaker训练任务中使用自定义Docker镜像时，超参数无法正常传递；移除image_uri参数后超参数能正常工作，但业务必须使用自定义镜像。尝试过普通Estimator和PyTorch Estimator两种方式，相关代码如下：

普通Estimator初始化代码 #

Estimator(
            image_uri=image_uri,
            role=role,
            instance_count=1,
            instance_type="ml.m5.large",
            hyperparameters=hyperparameters,
            output_path=f"s3://{ML_S3_BUCKET}/models/{training_job_name}",
        ) 

PyTorch Estimator初始化代码 #

hyperparameters={
        "epochs": 3000,
        "learning_rate": 0.0005,
        "input_size": 57,
        "hidden_size": 40,
        "output_size": 1,
        "patience": 1000,
        "hidden_layer_count": 4,
        "is_alpha": False,
    }
estimator = PyTorch(
            entry_point="train/train_illuvitar_neural_net.py",
            image_uri=image_uri,
            role=role,
            py_version="py310",
            framework_version="2.0",
            instance_count=1,
            instance_type="ml.m5.large",
            hyperparameters=hyperparameters,
            output_path=f"s3://{ML_S3_BUCKET}/models/{training_job_name}",
        )

Dockerfile内容 #

# Start with an official PyTorch image as a base
FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime

# Install any additional libraries and tools needed
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgomp1
RUN apt-get install -y libpq-dev gcc

# &&
# rm -rf /var/lib/apt/lists/*

# Install the necessary Python packages
COPY requirements_sm_training.txt /opt/ml/code/requirements.txt
RUN pip install --no-cache-dir -r /opt/ml/code/requirements.txt

# Probably don't need this
RUN pip3 install sagemaker-training

# Copy the training script
COPY train/train_illuvitar_neural_net.py /opt/ml/code/train
RUN chmod +x /opt/ml/code/train
COPY shared/ /opt/ml/code/shared/
COPY helper/ /opt/ml/code/helper/

# Set the working directory
WORKDIR /opt/ml/code
ENV PATH="/opt/ml/code:${PATH}"

训练脚本参数解析代码 #

def get_args():
    parser = argparse.ArgumentParser()

    is_local = os.environ.get("ISLOCAL")
    default_model_dir = (
        "./local_model_storage" if is_local else "/opt/ml/model"
    )  # os.environ.get("SM_MODEL_DIR")

    parser.add_argument("--epochs", type=int, default=500, help="Number of training epochs")
    parser.add_argument("--learning_rate", type=float, help="Learning rate for the optimizer")
    parser.add_argument("--batch_size", type=int, default=1024, help="Batch size for training")
    parser.add_argument("--model_dir", type=str, default=default_model_dir)
    parser.add_argument("--input_size", type=int, default=57, help="Input size for the model")
    parser.add_argument(
        "--hidden_size", type=int, default=40, help="Number of neurons in the hidden layer"
    )
    parser.add_argument("--output_size", type=int, default=1, help="Output size for the model")
    parser.add_argument("--patience", type=int, default=200, help="Patience for early stopping")
    parser.add_argument("--min_delta_for_early_stopping", type=float, default=0.03)
    parser.add_argument("--hidden_layer_count", type=int, default=3, help="Number of hidden layers")
    parser.add_argument(
        "--is_alpha", type=bool, default=False, help="Is this an alpha model?"
    )  # sagemaker doesn't support actions lik "store_true"
    parser.add_argument("--training_job_name", type=str, default="sagemaker-run-1684176842-32dd")
    parser.add_argument("--experiment_name", type=str, default="firstexperiment")

    parser.add_argument("--run_number", type=int, default=1)

    args = parser.parse_args()

    return args

解决方案 #

核心问题 #

自定义镜像缺少SageMaker训练容器的启动约定：SageMaker不会自动将超参数传递给脚本，需要通过镜像的启动命令配合sagemaker-training库实现超参数的读取和传递。

1. 修正Dockerfile的启动逻辑 #

修改Dockerfile，添加ENTRYPOINT指定训练启动命令，利用sagemaker-training库自动处理超参数：

# Start with an official PyTorch image as a base
FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime

# Install any additional libraries and tools needed
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgomp1 libpq-dev gcc \
 && rm -rf /var/lib/apt/lists/*  # 清理缓存，减少镜像体积

# Install the necessary Python packages
COPY requirements_sm_training.txt /opt/ml/code/requirements.txt
RUN pip install --no-cache-dir -r /opt/ml/code/requirements.txt

# 必须保留sagemaker-training，用于处理超参数传递
RUN pip3 install sagemaker-training

# 调整脚本命名为train.py，方便统一调用
COPY train/train_illuvitar_neural_net.py /opt/ml/code/train.py
COPY shared/ /opt/ml/code/shared/
COPY helper/ /opt/ml/code/helper/

# Set the working directory
WORKDIR /opt/ml/code

# 使用sagemaker-training的入口，自动读取超参数并传递给训练脚本
ENTRYPOINT ["python", "-m", "sagemaker_training.entry_point", "--entry-point", "train.py"]

2. 优化训练脚本的参数解析 #

当前--is_alpha的type=bool存在逻辑问题：SageMaker传递的超参数是字符串类型，"False"会被解析为True。修改为字符串后手动转换：

# 修改参数定义
parser.add_argument(
    "--is_alpha", type=str, default="False", help="Is this an alpha model?"
)

# 在获取参数后转换为bool类型
args = parser.parse_args()
args.is_alpha = args.is_alpha.lower() == "true"

3. 选择合适的Estimator使用方式 #

• 普通Estimator：无需指定entry_point，因为自定义镜像已包含完整训练逻辑，直接使用原初始化代码即可。
• PyTorch Estimator：指定image_uri后，SageMaker会上传本地entry_point脚本覆盖镜像内的文件，可能导致依赖目录（shared/、helper/）不匹配。更推荐使用普通Estimator，或确保本地目录结构与镜像完全一致。

关键理解 #

使用自定义镜像时，你需要对容器训练流程完全负责：

• SageMaker会将超参数写入容器内的/opt/ml/input/config/hyperparameters.json文件
• sagemaker-training库的作用是自动读取该文件，将超参数转换为命令行参数传递给训练脚本
• 若不用该库，也可手动编写启动脚本读取hyperparameters.json，自行构造参数调用训练脚本

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