AWS 文档中描述的 AWS 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅中国的 AWS 服务入门。

本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

调整您自己的推理容器

如果Amazon SageMaker所有预构建的推理容器均不足以满足您的情况，并且您希望使用您自己的 Docker 容器，请使用SageMaker推理工具包调整您的容器以使用 SageMaker 托管。要调整您的容器以使用SageMaker托管，请在一个或多个 Python 脚本文件中创建推理代码，并创建一个导入推理工具包的 Dockerfile。

推理代码包括推理处理程序、处理程序服务和入口点。在此示例中，它们存储为三个单独的 Python 文件。所有这三个 Python 文件都必须位于您的 Dockerfile 所在的同一目录中。

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

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

推SageMaker理工具包基于多模型服务器 （MMS） 构建。MMS 需要实施 函数的 Python 脚本来加载模型、预处理输入数据、从模型中获取预测并在模型处理程序中处理输出数据。

model_fn 函数

该model_fn函数负责加载您的模型。它采用一个model_dir参数来指定存储模型的位置。加载模型的方式取决于您使用的框架。model_fn 函数没有默认实现。您必须自行实施它。以下简单示例显示加载 PyTorch 模型的 model_fn 函数的实现：

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 推理工具包提供了对以下内容类型进行反序列化的默认实现：

如果您的模型需要不同的内容类型，或者您希望 在将输入数据发送到模型之前对其进行预处理，则必须实施 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(input_data)

output_fn 函数

output_fn 函数负责序列化predict_fn函数作为预测返回的数据。SageMaker 推理工具包实施默认 output_fn 函数，该函数序列化 Numpy 数组、JSON 和 CSV。如果您的模型输出任何其他内容类型，或者您希望在将数据发送给用户之前执行其他后处理，则必须实施您自己的 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：实施处理程序服务

处理程序服务由模型服务器执行。处理程序服务实施 initialize 和 handle 方法。在模型服务器启动时调用 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，复制步骤 2 中的模型处理程序，并将上一步中的 Python 文件指定为中的入口点Dockerfile。以下是您可以添加到 Dockerfile 以复制模型处理程序并指定入口点的行的示例。有关推理容器的 Dockerfile 的完整示例，请参阅 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：构建并注册容器

现在，您可以构建容器并在 中注册它Amazon Elastic Container Registry (Amazon ECR)。来自示例笔记本的以下 Shell 脚本构建容器并将其上传到您Amazon ECR账户中的 AWS 存储库。

以下 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 --region ${region} --no-include-email)

# 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，请参阅将模型部署到 SageMaker 托管服务。