调整自己的推理容器 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

调整自己的推理容器

如果任何 Amazon SageMaker 预构建推理容器都不适用于您的情况,您可以自己构建 Docker 容器,然后在 SageMaker 中使用它进行训练和推理。有关如何使用 SageMaker 构建自己的 Docker 容器进行训练和推理的更多信息以及示例,请参阅构建自己的算法容器

您还可以使用 SageMaker 推理工具包来调整您的容器,以便将其用于 SageMaker 托管。以下指南向您展示了如何使用 SageMaker 推理工具包来调整推理容器。

使用 SageMaker 推理工具包

要调整容器以将其用于 SageMaker 托管,请在一个或多个 Python 脚本文件中创建推理代码,并使用 Dockerfile 导入推理工具包。

用于工具包的推理代码应包含推理处理程序、处理程序服务和入口点。在此示例中,它们存储为三个独立的 Python 文件。所有这三个 Python 文件必须与 Dockerfile 位于相同目录中。

有关展示使用 SageMaker 推理工具包扩展容器的完整 Jupyter 笔记本示例,请参阅使用您自己算法容器的 Amazon SageMaker 多模型端点

步骤 1:创建推理处理程序

SageMaker 推理工具包构建在多模型服务器 (MMS) 上。MMS 需要一个实施以下功能的 Python 脚本:

  • 加载模型。

  • 对输入数据进行预处理。

  • 从模型获取预测。

  • 在模型处理程序中处理输出数据。

或者,您可以在实施的函数中使用 context 参数来访问其他服务信息,例如 GPU ID 或批次大小。如果您使用多个 GPU,则必须在 predict_fn 函数中指定 context 参数以选择预测所使用的 GPU。有关 context 参数及其用法的更多信息,请参阅 SageMaker Python SDK 中的 SageMaker PyTorch 模型服务器

model_fn 函数

model_fn 函数加载您的模型。在 SageMaker PyTorch 和 MXNet 推理工具包上,model_fn 函数有名为 default_model_fn 的默认实施。默认实施会加载使用 torchscript 保存的模型,其格式为 .pt.pth。如果您的模型需要自定义方法来加载,或者您要在加载模型时执行额外的步骤,则必须实施 model_fn 函数。下面的简单示例展示了 model_fn 函数的实施,该函数加载 PyTorch 模型:

def model_fn(self, model_dir): import torch logger.info('model_fn') device = "cuda" if torch.cuda.is_available() else "cpu" with open(os.path.join(model_dir, 'model.pth'), 'rb') as f: model = torch.jit.load(f) return model.to(device)

输入函数

input_fn 预处理输入数据。具体而言,input_fn 函数负责反序列化输入数据,以便它可以传递给模型。它获取输入数据和内容类型作为参数,并返回反序列化的数据。SageMaker 推理工具包提供了一个默认实施,用于反序列化以下内容类型:

  • JSON

  • CSV

  • Numpy 数组

  • NPZ

如果您的模型需要不同的内容类型,或者希望先预处理输入数据,然后再发送到模型,则必须实施 input_fn 函数。下面的示例展示了 input_fn 函数的简单实施。

from sagemaker_inference import content_types, decoder def input_fn(self, input_data, content_type): """A default input_fn that can handle JSON, CSV and NPZ formats. Args: input_data: the request payload serialized in the content_type format content_type: the request content_type Returns: input_data deserialized into torch.FloatTensor or torch.cuda.FloatTensor depending if cuda is available. """ return decoder.decode(input_data, content_type)

predict_fn 函数

predict_fn 函数负责从模型中获取预测。它获取从 input_fn 返回的模型和数据作为参数,然后返回预测。predict_fn 没有默认实施。您必须自行实施。以下是 PyTorch 模型的 predict_fn 函数的简单实施。

def predict_fn(self, data, model): """A default predict_fn for PyTorch. Calls a model on data deserialized in input_fn. Runs prediction on GPU if cuda is available. Args: data: input data (torch.Tensor) for prediction deserialized by input_fn model: PyTorch model loaded in memory by model_fn Returns: a prediction """ return model(data)

output_fn 函数

output_fn 函数在模型处理程序中处理输出数据。模型处理程序包含在模型级别定义的函数。output_fn 函数负责序列化 predict_fn 函数作为预测返回的数据。SageMaker 推理工具包实施序列化 Numpy 数组、JSON 和 CSV 的默认 output_fn 函数。如果模型输出任何其他内容类型,或者您希望在将数据发送给用户之前对其进行其他后处理,则必须实施自己的 output_fn 函数。以下显示了用于 PyTorch 模型的简单 output_fn 函数。

from sagemaker_inference import encoder def output_fn(self, prediction, accept): """A default output_fn for PyTorch. Serializes predictions from predict_fn to JSON, CSV or NPY format. Args: prediction: a prediction result from predict_fn accept: type which the output data needs to be serialized Returns: output data serialized """ return encoder.encode(prediction, accept)

步骤 2:实施处理程序服务

处理程序服务实施 initializehandle 方法,由模型服务器运行。模型服务器在启动时调用 initialize 方法,并对所有传入模型服务器的推理请求调用 handle 方法。有关更多信息,请参阅多模型服务器文档中的定制服务。以下是 PyTorch 模型服务器的处理程序服务示例。

from sagemaker_inference.default_handler_service import DefaultHandlerService from sagemaker_inference.transformer import Transformer from sagemaker_pytorch_serving_container.default_inference_handler import DefaultPytorchInferenceHandler class HandlerService(DefaultHandlerService): """Handler service that is executed by the model server. Determines specific default inference handlers to use based on model being used. This class extends ``DefaultHandlerService``, which define the following: - The ``handle`` method is invoked for all incoming inference requests to the model server. - The ``initialize`` method is invoked at model server start up. Based on: https://github.com/awslabs/mxnet-model-server/blob/master/docs/custom_service.md """ def __init__(self): transformer = Transformer(default_inference_handler=DefaultPytorchInferenceHandler()) super(HandlerService, self).__init__(transformer=transformer)

步骤 3:实施入口点

入口点通过调用处理程序服务来启动模型服务器。在您的 Dockerfile 中指定入口点的位置。以下是入口点的示例:

from sagemaker_inference import model_server model_server.start_model_server(handler_service=HANDLER_SERVICE)

步骤 4:编写 Dockerfile

编写 Dockerfile。有关推理容器的 Dockerfile 的完整示例,请参阅 Dockerfile。在 Dockerfile 中,复制步骤 2 中的模型处理程序,并在 Dockerfile 中将上一步中的 Python 文件指定为入口点。以下是代码行示例,您可将其添加到 Dockerfile 中以复制模型处理程序并指定入口点。

# Copy the default custom service file to handle incoming data and inference requests COPY model_handler.py /home/model-server/model_handler.py # Define an entrypoint script for the docker image ENTRYPOINT ["python", "/usr/local/bin/entrypoint.py"]

步骤 5:构建并注册您的容器

最后,构建您的容器并注册。示例笔记本中的以下 Shell 脚本构建容器,并将其上传到您的 Amazon 账户中的 Amazon ECR 存储库:

注意

SageMaker 托管支持使用存储在 Amazon ECR 以外的存储库中的推理容器。有关信息,请参阅 为实时推理容器使用私有 Docker 注册表

以下 Shell 脚本展示了如何构建和注册容器。

%%sh # The name of our algorithm algorithm_name=demo-sagemaker-multimodel cd container account=$(aws sts get-caller-identity --query Account --output text) # Get the region defined in the current configuration (default to us-west-2 if none defined) region=$(aws configure get region) region=${region:-us-west-2} fullname="${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest" # If the repository doesn't exist in ECR, create it. aws ecr describe-repositories --repository-names "${algorithm_name}" > /dev/null 2>&1 if [ $? -ne 0 ] then aws ecr create-repository --repository-name "${algorithm_name}" > /dev/null # Get the login command from ECR and execute it directly aws ecr get-login-password --region ${region}|docker login --username AWS --password-stdin ${fullname} # Build the Docker image locally with the image name and then push it to ECR # with the full name. docker build -q -t ${algorithm_name} . docker tag ${algorithm_name} ${fullname} docker push ${fullname}

现在,您可以使用此容器在 SageMaker 中部署端点。

附加信息

有关如何在 SageMaker 中部署端点的示例,请参阅使用您自己的算法容器的 Amazon SageMaker 多模型端点

有关将模型部署到实时推理端点的介绍性教程,请参阅将机器学习模型部署到实时推理端点

有关如何通过 API 网关,在 SageMaker 中部署自定义 Docker 映像的信息,请参阅使用 Amazon API Gateway 映射模板和 Amazon SageMaker 创建运用机器学习的 REST API