SageMaker 训练工具包

使用 Amazon SageMaker 在 Docker 容器中训练机器学习模型。

:books: 背景

Amazon SageMaker 是一项完全托管的数据科学和机器学习 (ML) 工作流服务。 您可以使用 Amazon SageMaker 来简化构建、训练和部署 ML 模型的过程。

要训练模型，您可以将训练脚本和依赖项包含在一个运行您的训练代码的 Docker 容器 中。 容器提供了一个有效的隔离环境，确保一致的运行时和可靠的训练过程。

SageMaker 训练工具包 可以轻松添加到任何 Docker 容器中，使其与 SageMaker 兼容，用于 训练模型。 如果您使用 SageMaker 预建训练用 Docker 镜像，则此库可能已包含在内。

有关更多信息，请参阅 Amazon SageMaker 开发人员指南中关于 使用 Docker 容器进行训练 的部分。

:hammer_and_wrench: 安装

要在您的 Docker 镜像中安装此库，请将以下行添加到您的 Dockerfile：

RUN pip3 install sagemaker-training

:computer: 使用

以下是简要的操作指南。 有关使用 SageMaker 训练工具包构建的自定义训练容器的完整可运行示例，请参阅 示例笔记本。

创建 Docker 镜像并训练模型

1. 编写训练脚本（例如 train.py）。

2. 使用 Dockerfile 定义一个包含训练脚本及所有依赖项的容器。

   训练脚本必须位于 /opt/ml/code 目录中。 环境变量 SAGEMAKER_PROGRAM 定义了 /opt/ml/code 目录中哪个文件作为训练入口点。 当训练开始时，解释器会执行由 SAGEMAKER_PROGRAM 定义的入口点。 Python 和 Shell 脚本均受支持。

   FROM yourbaseimage:tag
   
   # 安装 SageMaker 训练工具包
   RUN pip3 install sagemaker-training
   
   # 将训练脚本复制到容器内
   COPY train.py /opt/ml/code/train.py
   
   # 将 train.py 定义为脚本入口点
   ENV SAGEMAKER_PROGRAM train.py
   

3. 构建并标记 Docker 镜像。

   docker build -t custom-training-container .
   

4. 使用该 Docker 镜像通过 SageMaker Python SDK 启动训练任务。

   from sagemaker.estimator import Estimator
   
   estimator = Estimator(image_name="custom-training-container",
                         role="SageMakerRole",
                         train_instance_count=1,
                         train_instance_type="local")
   
   estimator.fit()
   

   要在 SageMaker 上使用该镜像训练模型，需先将镜像 推送到 ECR，然后使用镜像 URI 启动 SageMaker 训练作业。

使用超参数向入口点传递参数

训练任务提供的任何超参数都会作为脚本参数传递给入口点。 SageMaker Python SDK 利用此功能将特殊超参数传递给训练任务，包括 sagemaker_program 和 sagemaker_submit_directory。 完整的 SageMaker 超参数列表可在 此处 查看。

1. 在入口点脚本中实现参数解析器。例如，在 Python 脚本中：

   import argparse
   
   if __name__ == "__main__":
     parser = argparse.ArgumentParser()
   
     parser.add_argument("--learning-rate", type=int, default=1)
     parser.add_argument("--batch-size", type=int, default=64)
     parser.add_argument("--communicator", type=str)
     parser.add_argument("--frequency", type=int, default=20)
   
     args = parser.parse_args()
     ...
   

2. 带超参数启动训练任务。

   {"HyperParameters": {"batch-size": 256, "learning-rate": 0.0001, "communicator": "pure_nccl"}}
   

使用环境变量读取附加信息

入口点通常需要一些无法从 hyperparameters 中获取的额外信息。 SageMaker 训练工具包会将这些信息写入环境变量，供脚本内部使用。 例如，此训练任务包含 training 和 testing 两个数据通道：

from sagemaker.pytorch import PyTorch

estimator = PyTorch(entry_point="train.py", ...)

estimator.fit({"training": "s3://bucket/path/to/training/data", 
               "testing": "s3://bucket/path/to/testing/data"})

环境变量 SM_CHANNEL_TRAINING 和 SM_CHANNEL_TESTING 提供了这些通道的路径：

import argparse
import os

if __name__ == "__main__":
  parser = argparse.ArgumentParser()

  ...

  # 从环境变量中读取输入通道 training 和 testing
  parser.add_argument("--training", type=str, default=os.environ["SM_CHANNEL_TRAINING"])
  parser.add_argument("--testing", type=str, default=os.environ["SM_CHANNEL_TESTING"])

  args = parser.parse_args()

  ...

训练开始时，SageMaker 训练工具包会打印所有可用的环境变量。有关提供的完整环境变量列表，请参阅 环境变量参考文档。

获取容器环境信息

要获取容器环境信息，可以初始化一个 Environment 对象。Environment 提供了对训练作业相关环境方面的访问权限，包括超参数、系统特性、文件系统路径、环境变量和配置设置等。它是在训练期间对容器环境的只读快照，并不包含任何状态信息。

from sagemaker_training import environment

env = environment.Environment()

# 从 `inputdataconfig.json` 文件中获取通道 "training" 的路径
training_dir = env.channel_input_dirs["training"]

# 从 `hyperparameters.json` 文件中获取超参数 "training_data_file"
file_name = env.hyperparameters["training_data_file"]

# 获取模型应保存的文件夹
model_dir = env.model_dir

# 训练模型
data = np.load(os.path.join(training_dir, file_name))
x_train, y_train = data["features"], keras.utils.to_categorical(data["labels"])
model = ResNet50(weights="imagenet")
...
model.fit(x_train, y_train)

# 在训练结束时将模型保存到 model_dir
model.save(os.path.join(model_dir, "saved_model"))

执行入口点

要执行入口点，调用 entry_point.run() 即可。

from sagemaker_training import entry_point, environment

env = environment.Environment()

# 将超参数作为脚本参数读取
args = env.to_cmd_args()

# 获取环境变量
env_vars = env.to_env_vars()

# 执行入口点
entry_point.run(uri=env.module_dir,
                user_entry_point=env.user_entry_point,
                args=args,
                env_vars=env_vars)

如果入口点执行失败，trainer.train() 会将错误信息写入 /opt/ml/output/failure。否则，它会将成功信息写入 /opt/ml/success 文件。

:scroll: 许可证

本库采用 Apache 2.0 许可证 授权。更多详情请参阅 LICENSE 文件。

:handshake: 贡献

欢迎贡献！如果您想提交问题或拉取请求，请阅读我们的 贡献指南。

SageMaker Training Toolkit 快速上手指南

SageMaker Training Toolkit 是一个用于在 Docker 容器内运行 Amazon SageMaker 训练任务的工具库。它允许你将自定义的训练脚本和依赖打包进容器，并使其兼容 SageMaker 的训练环境。

环境准备

• 系统要求：支持 Docker 的操作系统（Linux, macOS, Windows with WSL2）。
• 前置依赖：
  • Python 3.6 或更高版本。
  • Docker Engine 已安装并运行。
  • AWS CLI 已配置（用于推送镜像到 ECR）。
  • sagemaker Python SDK（用于本地测试或提交训练任务）。

安装步骤

要将此库集成到你的自定义训练容器中，请在你的 Dockerfile 中添加以下安装命令：

RUN pip3 install sagemaker-training

提示：如果你使用的是 AWS 提供的预构建 SageMaker 深度学习镜像，该库通常已预先安装，无需重复添加。

基本使用

以下是构建自定义训练容器并运行模型的最简流程：

1. 编写训练脚本

创建一个 Python 脚本（例如 train.py），作为训练的入口点。

# train.py
if __name__ == "__main__":
    print("Training started...")
    # 在此处添加你的模型训练逻辑
    # 可以通过 os.environ 读取 SageMaker 注入的环境变量
    # 可以通过 argparse 读取超参数

2. 创建 Dockerfile

编写 Dockerfile，将基础镜像、工具库和训练脚本打包在一起。注意：脚本必须位于 /opt/ml/code 目录，并通过 SAGEMAKER_PROGRAM 环境变量指定入口文件。

FROM python:3.8-slim

# 安装 SageMaker Training Toolkit
RUN pip3 install sagemaker-training

# 将训练脚本复制到容器内的指定目录
COPY train.py /opt/ml/code/train.py

# 定义入口点脚本
ENV SAGEMAKER_PROGRAM train.py

3. 构建 Docker 镜像

在包含 Dockerfile 和 train.py 的目录下执行构建命令：

docker build -t custom-training-container .

4. 运行训练任务

你可以使用 sagemaker Python SDK 在本地模式（local）下测试该镜像，或者将其推送到 Amazon ECR 后在云端运行。

本地测试示例：

from sagemaker.estimator import Estimator

estimator = Estimator(image_name="custom-training-container",
                      role="SageMakerRole",
                      train_instance_count=1,
                      train_instance_type="local")

estimator.fit()

云端运行提示： 若要在 AWS SageMaker 上运行，请先将镜像推送到 Amazon ECR：

docker tag custom-training-container <your-account-id>.dkr.ecr.<region>.amazonaws.com/custom-training-container:latest
docker push <your-account-id>.dkr.ecr.<region>.amazonaws.com/custom-training-container:latest

然后使用生成的 ECR 镜像 URI 替换上述 Python 代码中的 image_name 并将 train_instance_type 设置为实际的云实例类型（如 ml.m5.large）。