如何在Amazon SageMaker单个端点中按顺序集成两个模型？

完全可以通过Amazon SageMaker实现你需要的单端点串联模型推理需求，你之前对model_fn、prediction_fn的理解有局限——这些函数支持在自定义推理脚本中加载并调用多个模型，完成"模型A输出→模型B输入"的串联逻辑。

核心实现思路 #

• 自定义推理脚本：在推理脚本中同时加载两个模型，在预测环节按顺序调用模型完成串联推理
• 打包模型资源：将两个模型文件、推理脚本及依赖打包后上传至S3
• 部署单端点：基于打包的模型创建SageMaker模型并部署为实时端点

示例推理脚本（inference.py） #

# 以PyTorch为例，其他框架逻辑类似
import torch

def model_fn(model_dir):
    # 从模型目录加载两个预训练模型
    model1 = torch.jit.load(f"{model_dir}/model_v1.pt")
    model1.eval()
    
    model2 = torch.jit.load(f"{model_dir}/model_v2.pt")
    model2.eval()
    
    # 返回包含两个模型的字典
    return {"model1": model1, "model2": model2}

def predict_fn(input_data, model):
    # 第一步：用第一个模型处理输入
    intermediate_output = model["model1"](input_data)
    
    # 第二步：将第一个模型的输出作为第二个模型的输入
    final_output = model["model2"](intermediate_output)
    
    return final_output

额外注意事项 #

• 如果两个模型基于不同框架（比如一个用TensorFlow，一个用PyTorch），需要构建包含对应框架依赖的自定义容器，再部署端点
• 选择实例类型时要考虑两个模型的内存占用总和，避免因资源不足导致加载失败
• 模型打包时需确保两个模型文件、推理脚本及requirements.txt（若有依赖）都被包含在.tar.gz归档中

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