SageMaker推理管道能否将原始输入同时传递给两个PyTorch模型？

核心实现逻辑：SageMaker推理管道默认将前一个模型的输出作为后一个模型的唯一输入，你只需要在第一个模型的输出阶段，把原始输入和自身预测结果打包成统一结构传递给第二个模型即可，不需要修改管道的基础配置。

具体实现步骤 #

步骤1：修改model_1的处理逻辑，打包原始输入+预测结果 #

调整model_1的处理函数，将原始请求体和自身预测结果一起序列化后传给下一个模型，示例代码如下：

import base64
import json

def input_fn(request_body, request_content_type):
    parsed_input = None
    if request_content_type == "application/x-npy":
        parsed_input = some_function_to_parse_input(request_body)
    # 返回解析后的输入 + 原始二进制请求体
    return parsed_input, request_body

def predict_fn(input_tuple, predictor):
    parsed_input, raw_payload = input_tuple
    predictions = predictor(parsed_input)
    # 返回预测结果 + 原始请求体
    return predictions, raw_payload

def output_fn(output_tuple, response_content_type):
    predictions, raw_payload = output_tuple
    # 二进制原始payload转base64编码，适配JSON序列化
    payload_b64 = base64.b64encode(raw_payload).decode('utf-8')
    # 打包成统一结构传给下一个模型
    output = {
        "model_1_predictions": predictions,
        "raw_payload_b64": payload_b64
    }
    return json.dumps(output), "application/json"

步骤2：修改model_2的处理逻辑，拆分两种输入 #

model_2的输入处理函数直接解析model_1传递的打包内容，即可同时拿到原始输入和model_1的预测结果，示例代码如下：

import base64
import json

def input_fn(request_body, request_content_type):
    if request_content_type == "application/json":
        input_data = json.loads(request_body)
        # 拆分model1预测结果和原始payload
        model_1_pred = input_data["model_1_predictions"]
        # 把base64编码的原始payload转回二进制
        raw_payload = base64.b64decode(input_data["raw_payload_b64"])
        # 可按需进一步解析原始payload
        parsed_raw_input = some_function_to_parse_input(raw_payload)
        return parsed_raw_input, model_1_pred

def predict_fn(input_tuple, predictor):
    parsed_raw_input, model_1_pred = input_tuple
    # 同时使用原始输入和model1的输出完成推理
    final_outputs = predictor(parsed_raw_input, model_1_pred)
    return final_outputs

def output_fn(predictions, response_content_type):
    # 按业务需求返回最终结果即可
    return json.dumps(predictions)

注意事项 #

• 两个模型之间传递的Content-Type要匹配，上述示例中model1输出格式为application/json，所以model2的input_fn要对应处理该类型
• 如果原始payload体积较大，base64编码会增加约33%的数据体积，对延迟敏感的场景可以自定义二进制序列化格式代替JSON，比如使用MessagePack或者自定义协议打包两类内容
• 不需要修改推理管道的创建逻辑，按原有流程把两个模型按顺序加入管道即可

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