用 SageMaker SDK v3 在边缘端管理车队级计算机视觉模型生命周期
Fleet-wide computer vision lifecycle at the edge with SageMaker SDK v3
基于 Amazon SageMaker Python SDK v3 的车队级计算机视觉模型生命周期方案发布,用 SageMaker Neo 替代已弃用的 SageMaker Edge Manager 做边缘优化。
全车队范围的计算机视觉模型生命周期管理,是一个在质量检测概念验证成功的那一刻就会浮现的问题。一台相机和一个模型运行良好。接着工厂要求在40条产线上取得同样的结果,每条产线有多台相机,而每台相机看到的世界都不一样。每台相机角度不同,同一不锈钢表面的反光不同,环境光也因哪扇车间门打开而不同。在我们参与的生产部署中,用汇总素材训练的单一模型在足够多的单台相机上表现不佳,以至于每台相机最终都需要自己的微调模型。现在真正的问题不再是训练好一个模型,而是训练、评估、版本化并发布数千个模型,且无需为每一个模型都安排人工介入。
本文针对当前的 Amazon SageMaker Python SDK v3,更新了2023年系列文章Build an end-to-end MLOps pipeline for visual quality inspection at the edge中引入的模式。它用 Amazon SageMaker Neo 进行边缘优化,取代了此后已弃用的 Amazon SageMaker Edge Manager。我详细介绍了一个双分支 Amazon SageMaker Pipeline,它为每台相机训练一个异常检测器和一个缺陷分类器,并设有质量门禁来拦截表现不佳的模型。流水线内的 SageMaker Neo 编译步骤确保每个注册模型都附带设备优化后的产物,而 AWS Step Functions 编排层则控制全车队训练的节奏。你使用 AWS Cloud Development Kit(AWS CDK)部署所有内容,在合成数据上完整运行一次流水线的计算成本约为 $0.07。
完整示例可在 aws-samples GitHub 仓库中获取。
解决方案概述
下图展示了该示例的全车队模型生命周期。编排层将再训练请求排队并限制并发工作。每台相机的训练流水线准备数据,训练并评估异常检测器和分类器,并为边缘设备生成带版本的模型产物。
图1:解决方案架构
训练部分是一个由相机参数化的单一 SageMaker Pipeline,每次执行运行13个步骤。三个入口步骤并行运行:Preprocess、Fetch Latest Anomaly 和 Fetch Latest Classifier。流水线预处理该相机的图像,并为每个分支获取该相机最近一次获批的模型,然后训练两个模型。第一个是一个小型卷积自编码器,它学习正常产品的外观,并通过重建误差标记异常。第二个是一个经过微调的 ResNet18 分类器,用于确认缺陷。每个分支将其模型导出为 ONNX,用留出数据评估导出的产物,并通过质量门禁。通过门禁后,模型会针对配置的边缘目标用 SageMaker Neo 编译,只有到这时才会注册到 SageMaker Model Registry。因此,每个获批的包都带有设备就绪的编译产物。包以相机 ID 作为元数据,每个分支的热启动获取会据此过滤,因此再训练始终从该相机自身的历史记录继续。
编排层负责控制整个集群的节奏。重新训练请求会进入 Amazon Simple Queue Service(Amazon SQS)FIFO 队列,由 AWS Lambda 函数为每个请求启动一次 Step Functions 执行。每次执行在启动其流水线之前,都会通过 DynamoDB 信号量获取一个槽位,因此突发请求会以稳定速率被处理,而不是全部同时启动。完成事件会发布到 Amazon Simple Notification Service(Amazon SNS),而运行故障则通过一条 Amazon EventBridge 规则到达同一主题,并保留原始错误信息。
系统的输出是经过版本管理、质量门控的每摄像头模型产物。Model Registry 保存可移植的 ONNX,S3 保存该模型的 Neo 编译构建,这是一个由 DLR 运行时加载的设备优化共享库,而不是 ONNX。两者都已准备好,可供你的工厂所使用的任何交付机制使用。
先决条件
要部署此示例,你需要以下内容:
- 一个具有创建 AWS Identity and Access Management(IAM)角色权限的 AWS 账户。
- 一个 Amazon Simple Storage Service(Amazon S3)存储桶。
- 一个 Amazon Elastic Container Registry(Amazon ECR)存储库。
- 各堆栈中的 SageMaker AI、AWS Step Functions、AWS Lambda、Amazon Simple Queue Service(Amazon SQS)、Amazon Simple Notification Service(Amazon SNS)、Amazon EventBridge 和 Amazon DynamoDB 资源。
- Python 3.10、3.11 或 3.12,以及 Docker。
- AWS CDK CLI 版本 2.1131.0 或更高版本。
较早的 CLI 版本在与近期版本的 aws-cdk-lib 配合时会因 cloud assembly schema 不匹配而失败。无需增加服务配额。编排层默认最多同时接纳五个流水线执行,这使并发的 SageMaker 训练作业保持在标准账户配额之内。
如果你要部署到非 CLI 默认区域的 AWS 区域,请在部署命令本身设置 AWS_REGION、AWS_DEFAULT_REGION 和 CDK_DEFAULT_REGION 这三项,而不是在 shell 配置文件中设置,例如 AWS_REGION=us-east-2 AWS_DEFAULT_REGION=us-east-2 CDK_DEFAULT_REGION=us-east-2 cdk deploy --all。CDK CLI、Boto3 和 AWS Command Line Interface(AWS CLI)从相互重叠但不完全相同的变量集合中解析区域,不匹配会导致令人困惑的部分部署。在命令前添加前缀只会为那次调用设置这三项,因此你自己的默认区域保持不变。
使用 SDK v3 定义训练流水线
本节中的代码摘自仓库中的 CDK 应用程序,而不是独立脚本,它在 cdk deploy 期间运行,此时堆栈会合成流水线定义。环境设置和部署命令将在后面的“部署和运行示例”一节中介绍。SDK v3 将 SageMaker Python SDK 重组为模块化包,导入语句展示了该库的新形态:
from sagemaker.train import ModelTrainer
from sagemaker.train.configs import InputData, Compute, OutputDataConfig
from sagemaker.core.training.configs import MetricDefinition
from sagemaker.serve.model_builder import ModelBuilder
from sagemaker.mlops.workflow.pipeline import Pipeline
from sagemaker.mlops.workflow.steps import ProcessingStep, TrainingStep
from sagemaker.mlops.workflow.model_step import ModelStep
from sagemaker.mlops.workflow.condition_step import ConditionStep
from sagemaker.core.workflow.conditions import ConditionGreaterThanOrEqualTo
from sagemaker.core.workflow.functions import JsonGet安装 sagemaker>=3.0.0,<4 会拉入 sagemaker-core、sagemaker-train、sagemaker-serve 和 sagemaker-mlops 作为单独的包。v2 中分散在 Estimator 构造函数参数中的训练配置,现在位于传递给 ModelTrainer 的类型化配置对象中。以下是异常分支的训练器,逐字摘自示例:
anomalib_trainer = ModelTrainer(
training_image=custom_training_image_uri,
compute=Compute(
instance_type=training_instance_type.to_string(),
instance_count=1,
),
base_job_name="cv-pipeline-anomalib-train",
sagemaker_session=pipeline_session,
role=role_arn,
hyperparameters={
"program": "train_anomalib.py",
"model_name": "autoencoder",
"image_size": "256",
"num_epochs": "3",
"batch_size": "16",
"learning_rate": "0.001",
},
input_data_config=anomalib_input_data,
output_data_config=OutputDataConfig(
s3_output_path=run_prefix("anomalib-training")
),
).with_metric_definitions(
[
MetricDefinition(
name="image_auroc",
regex=r"validation:image_auroc=([0-9]*\.?[0-9]+)",
),
]
)anomalib_input_data 从上游步骤属性连接四个通道(train、validation、test、pretrained_model),而 run_prefix 将每个 S3 路径限定到流水线执行 ID,因此并发的集群运行永远不会冲突。两者都在仓库中。
当流水线位于 CDK 应用程序内部时,有两个 v3 行为值得注意,而这两个行为都没有在变更日志中作为破坏性变更出现。首先,SDK v3 在构建流水线定义时会针对实时 IAM 验证流水线的执行角色。此验证发生在合成时,而不是执行时。CDK 应用无法在一次部署中同时创建角色和流水线,因为当 SDK 验证该角色时,角色尚不存在。这与 cdk bootstrap 不同,后者为每个账户和区域一次性配置 CDK 工具包资源,并且可能已经在你的账户中完成。示例自身的两阶段部署由一个恰好同名的上下文标志驱动。使用 -c bootstrap=true 的初始阶段创建角色、存储桶和 ECR 仓库,第二阶段则针对现已存在的角色合成流水线。嵌入流水线定义中的角色 Amazon Resource Name (ARN) 也必须是一个具体字符串,而不是未解析的 CDK 令牌。
其次,v3 角色验证器在没有资源 ARN 的情况下模拟 ECR 拉取权限,因此,即使执行角色在运行时可以正常工作,限定到单个 ECR 仓库的执行角色也无法通过验证。示例中的角色因此携带了更宽泛的 ECR 授权。
最后,示例采用的一个 v3 习惯是:使用 PipelineSession(default_bucket=your_bucket) 构建流水线会将 SDK 管理的上传(例如步骤代码)路由到流水线自己的存储桶中,因此每个构件都存放在一个地方,并且执行角色的权限覆盖单个存储桶。(授予角色访问 SDK 默认存储桶的权限也可以。这种方法更整洁。)
与缺陷数据行为相匹配的质量门禁
每个训练好的模型在注册前都会进行评估,并且评估针对导出的 ONNX 构件运行,而不是内存中的 PyTorch 模型。这一区别很重要:你进行门禁的构件就是你交付的构件,因此导出问题绝不会让一个有问题的模型通过看起来健康的训练运行。
异常分支的门禁要求 AUROC 至少为 0.8。分类器分支同时以两个指标作为门禁,而指标的选择由缺陷数据的决定性属性塑造:严重的类别不平衡。在此示例中,大约 90% 的图像是良品,这对生产线来说是现实的。在这种分布下,单纯的准确率毫无意义,因为一个对每张图像都预测“良品”的模型得分约为 0.9,却捕获零缺陷。平衡准确率,即各类别召回率的平均值,可以纠正这一点,但它有自己的盲点。一个在良品上召回率完美、在缺陷上召回率仅为 20% 的模型仍然得分为 (1.0 + 0.2) / 2 = 0.6。因此,门禁要求平衡准确率至少为 0.6,并且缺陷召回率至少为 0.5,两者同时满足:
step_check_classifier = ConditionStep(
name="CheckClassifierQuality",
conditions=[
ConditionGreaterThanOrEqualTo(
left=JsonGet(
step_name=step_eval_classifier.name,
property_file="evaluation",
json_path="balanced_accuracy",
),
right=min_accuracy,
),
ConditionGreaterThanOrEqualTo(
left=JsonGet(
step_name=step_eval_classifier.name,
property_file="evaluation",
json_path="defect_recall",
),
right=min_defect_recall,
),
],
if_steps=[step_compile_classifier, step_register_classifier],
else_steps=[step_fail_classifier],
depends_on=[step_eval_classifier],
)ConditionStep 会对其条件进行 AND 运算,因此两者都必须通过。阈值是流水线参数,可在每次执行时调整而无需重新部署。如果门禁失败,FailStep 会将执行标记为失败,并且该分支的模型永远不会到达注册表(两个分支独立进行门禁,因此另一个分支的模型仍可能注册)。设备群会继续为该摄像头提供之前已批准的模型。
训练代码从另一侧解决了同样的不平衡问题。分类器是一个微调过的resnet18,使用 ImageNet 权重和双类别头。在每摄像头数据集规模下,微调预训练权重正是防止模型坍缩到多数类上的关键。训练代码使用逆类别频率的WeightedRandomSampler来平衡训练批次,而损失保持不加权。将加权损失与加权采样结合会过度校正,偏向少数类。
在流水线内使用 SageMaker Neo 为边缘编译
注意前面代码中门控的if_steps:通过门控后,编译先运行,注册依赖于它。两个训练脚本都导出具有动态批次轴的 ONNX。质量门控在 ONNX Runtime 下对这些导出的产物进行评分,然后CompileAnomalibNeo或CompileClassifierNeo步骤将门控通过的产物交给 SageMaker Neo。不变式是:门控、编译、注册。只有当模型也为机群的边缘目标编译过,它才会成为 Approved 包,因此注册表永远不会公布设备无法运行的模型。
编译步骤封装了三个容易出错的细节。Neo 要求输入 tarball 只包含模型文件,但训练作业会生成model.tar.gz,其中包含检查点、指标以及 ONNX 模型。该步骤会先重新打包一个仅含 ONNX 的 tarball。Neo 还要求具体的输入形状,而导出声明了动态批次轴,因此在编译请求中固定批次大小为 1。而且由于流水线步骤可以重试,编译作业名称从执行 ID 确定性派生,因此重试会继续轮询它已经启动的作业,而不是将其遗弃。
目标是流水线参数:NeoTargetDevice选择一个命名设备,例如jetson_xavier或rasp4b,其默认值为NeoTargetOs/NeoTargetArch平台编译,默认为 LINUX/X86_64。编译后的包包含优化后的compiled.so以及用于加载它的 DLR 运行时。默认输出可在任何 x86_64 Linux 主机上运行,因此你可以在确定设备特定目标之前进行验证。编译本身不产生费用。你需要为驱动它的处理实例和输出的 S3 存储付费。为了以后在不重新训练的情况下为其他设备类型编译已注册的模型,仓库包含一个临时脚本,对 Model Registry 运行相同的流程。
Model Registry 中的每摄像头模型
每个分支注册到自己的模型包组中,每个包都携带摄像头 ID 及其评估指标。从示例中,省略了指标接线:
classifier_model_builder = ModelBuilder(
s3_model_data_url=classifier_model_data,
image_uri=custom_training_image_uri,
sagemaker_session=pipeline_session,
role_arn=role_arn,
)
step_register_classifier = ModelStep(
name="RegisterClassifierModel",
depends_on=[step_compile_classifier],
step_args=classifier_model_builder.register(
model_package_group_name="sm1k-classifier-model-group",
approval_status="Approved",
model_metrics=classifier_metrics,
customer_metadata_properties={"CameraId": camera_id},
),
)customer_metadata_properties中的摄像头身份是将两个共享包组转变为机群目录的关键。来自所有摄像头的包都位于相同的组中,通过按元数据过滤可以恢复每摄像头的历史记录。每个分支都有一个获取步骤,按最新优先分页浏览已批准的包,直到找到请求摄像头的最新模型,因为ListModelPackages无法在服务器端按客户元数据过滤。在摄像头的首次运行时,获取步骤找不到任何内容,训练回退到其冷启动初始化(分类器使用 ImageNet 权重,自编码器使用随机初始化)。之后的每次运行都从该摄像头自己的先前模型恢复,这正是使频繁、廉价、增量重新训练在整个机群中可行的原因。
请注意这些注册表条目是什么、不是什么:它们编目的是带版本、经过质量门控的工件,但其容器引用指向的是训练镜像,该镜像不提供任何推理端点。它们服务于一种车队工作流——由下游系统拉取工件并将其分发到设备——而不是用于CreateEndpoint。
编排车队训练请求
每个 AWS 账户对并发的 SageMaker 训练作业都有服务配额,而这里的每次管道执行都会启动多个作业,因此车队重新训练需要节奏控制。必须有某种机制防止 500 个请求的突发变成 500 个同时启动的管道。示例使用信号量,即一个可用槽位计数器,其工作方式类似于有五个车位的停车场。在 Step Functions 状态机启动管道之前,它通过在 DynamoDB 中原子递增计数器来占用一个槽位,并附带一个条件:一旦达到上限就拒绝递增:
acquire_semaphore = tasks.DynamoUpdateItem(
self, "AcquireSemaphore",
table=semaphore_table,
key={"LockName": tasks.DynamoAttributeValue.from_string("pipeline-concurrency")},
update_expression="SET ActiveCount = if_not_exists(ActiveCount, :zero) + :one",
condition_expression="attribute_not_exists(ActiveCount) OR ActiveCount < :limit",
)如果所有槽位都被占用,执行会等待并重试。当其管道完成后,它会递减计数器以释放槽位。有一个细节值得带入你自己的设计中。DynamoDB 条件表达式在单个原子操作中强制执行上限,而不是由状态机读取计数器再做出决定,后者会在读取和写入之间留下竞态窗口。该仓库实现了其余运维加固措施,包括在每条失败路径上释放槽位,以及通过 Amazon EventBridge 将失败的执行转发到 Amazon SNS。其 README 涵盖了等待边界的规模设定。
上限为五是一个堆栈参数,设置得较低,以便示例能在任何标准账户中部署和运行,而无需申请配额增加。请根据你账户的训练作业配额支持情况调整它。超出单个账户的配额后,扩展车队训练就是另一个话题了。标准手段包括提高服务配额、将训练分散到不同实例类型或区域,以及将车队分布到多个账户。你也可以将积压工作保留在 SQS 中,仅在容量释放时接纳工作。这些选择取决于你车队的规模和形态。示例为你提供了可扩展的工作骨架。
部署和运行示例
仓库 README 中的演练是权威版本。以下简短形式从干净的检出开始,并在首次部署前设置环境:
git clone https://github.com/aws-samples/sample-sagemaker-fleet-cv-pipeline
cd sample-sagemaker-fleet-cv-pipeline && python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
cdk bootstrap # once per account/region
cdk deploy SageMakerCVPipelineStack -c bootstrap=true # roles, bucket, ECR
cdk deploy --all # pipeline + orchestration
./scripts/push_ecr_image.sh # build and push training image默认情况下,管道生成合成生产线图像,因此你可以在没有数据集、没有数据审批的情况下端到端验证整个系统。一次完整运行会执行全部 13 个步骤,包括两次 Neo 编译,耗时约 15 分钟,计费 ml.m5.xlarge 时间成本约为 $0.07(美国东部(俄亥俄)区域(us-east-2),2026 年 7 月,按需)。在示例规模下,编排资源合计每月成本远低于 $0.10。
为了在真实图像上运行该流水线,仓库中包含一个用于下载 BCCD 数据集 的下载器。该数据集由 364 张带标注的血涂片显微镜图像组成,最初在 GitHub 上开源,并以 MIT 许可证重新整理为 Pascal VOC 格式。同一数据集也由 Roboflow 以公有领域形式分发。该脚本将血小板的存在映射到缺陷类别,可选地按摄像头拆分图像以模拟设备集群,并将训练 tar 包上传到流水线存储桶。随附的 notebook 在真实数据上运行完整流水线,并对一个正常样本和一个异常样本进行并行的 ONNX 推理。将血小板视为缺陷是对制造缺陷分布的一种示意性替代,并非临床结论。
清理
模型包在流水线运行时创建,因此 AWS CloudFormation 无法删除它们。请先用 AWS CLI 删除它们,然后使用 cdk destroy --all -c bootstrap=true --force 销毁堆栈。README 列出了确切的命令。
结论
在本文中,我展示了一个可运行的 SageMaker SDK v3 参考架构,用于设备集群规模的计算机视觉生命周期。该架构包含一个双分支训练流水线,其中每个模型都要在其导出的 ONNX 产物上通过质量门禁,使用 SageMaker Neo 编译,然后才被注册。模型注册表布局利用摄像头元数据将共享的包组转变为可按摄像头热启动的历史记录,而 Step Functions 编排层则在账户配额内调节集群训练节奏。对于需求可由 Amazon Lookout for Vision 等托管选项满足的团队来说,那是更简单的路径。此模式针对的是托管选项未覆盖的按摄像头、设备集群规模的约束。这些设计规则可迁移到计算机视觉之外:以平衡准确率加少数类召回率而非准确率作为门禁,对你交付的产物进行评估和编译,并让运维故障发出通知而不是静默终止。
该示例可在几分钟内部署完成,一次完整验证运行的计算成本约为七美分,因此低成本的下一步是克隆仓库并运行它。在此基础上,将 BCCD 下载器指向真实存储桶,为你的硬件设置 NeoTargetDevice,或根据你的设备集群调整编排边界。至于 SDK v3 迁移本身,SageMaker Python SDK 文档以及宣布新 SDK 的两部分系列文章是我全程保持打开的参考资料。
如果你需要帮助将此模式适配到你的生产环境,请联系你的 AWS 客户团队。
关于作者
Kiowa Jackson
Kiowa 是 AWS ProServe 的高级机器学习工程师。他专注于将计算机视觉与智能体 AI 相结合的视觉智能系统,并在企业工业领域将这些系统投入生产。
来源:AWS Machine Learning Blog · aws.amazon.com
