AWS SageMaker SDK v3 构建边缘端多摄像头计算机视觉模型生命周期流程
Fleet-wide computer vision lifecycle at the edge with SageMaker SDK v3
AWS Machine Learning Blog 给出基于 Amazon SageMaker Python SDK v3 的多摄像头计算机视觉模型生命周期架构,为每个摄像头训练异常检测器和缺陷分类器,并对导出的 ONNX 模型执行质量门禁、SageMaker Neo 编译后再注册。
文章给出多摄像头模型生命周期的可复用架构,将质量门禁、边缘编译和并发编排串成完整流程。
全车队计算机视觉模型生命周期管理,是质量检测概念验证一旦成功就会出现的问题。一台摄像机和一个模型运行良好。但随后工厂要求在 40 条产线上获得相同结果,而每条产线又配有多台摄像机,并且每台摄像机看到的世界都不一样。每台摄像机的角度不同,同一不锈钢表面反射的眩光不同,还会根据哪个舱门处于打开状态而受到不同环境光的影响。在我们参与过的生产部署中,使用汇总视频训练的单个模型,在足够多的摄像机上表现不佳,最终每台摄像机都需要拥有自己的微调模型。现在,真正的问题不再是如何训练好一个模型,而是在没有人为参与每个模型的情况下,训练、评估、版本控制并交付成千上万个模型。
本文针对当前的 Amazon SageMaker Python SDK v3,更新了 2023 年系列文章 在边缘为视觉质量检测构建端到端 MLOps 管道 中介绍的模式。本文用 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。
- 版本为 2.1131.0 或更高版本的 AWS CDK CLI。
旧版 CLI 会因与最新版 aws-cdk-lib 的云程序集架构不匹配而失败。无需提高任何服务配额。默认情况下,编排层最多同时允许五次流水线执行,这可以使并发的 SageMaker 训练作业保持在账户的标准配额内。
如果你要部署到不同于 CLI 默认值的 AWS Region,请在 deploy 命令本身中设置 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) 会从相互重叠但并不完全相同的变量集合中解析 Region,而不一致会导致令人困惑的部分部署。为命令添加前缀只会对本次调用设置这三个值,因此不会改变你自己的默认 Region。
使用 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 不同,后者为每个账户和 Region 配置一次 CDK 工具包资源,而你的账户可能已经完成了这一步。示例自身的两阶段部署由一个上下文标志驱动,该标志恰好也叫这个名字。第一次使用 -c bootstrap=true 部署会创建角色、存储桶和 ECR 存储库,第二次则根据此时已存在的角色合成流水线。嵌入流水线定义中的角色 Amazon 资源名称(ARN)也必须是具体的字符串,而不能是尚未解析的 CDK 令牌。
其次,v3 角色验证器在模拟 ECR 拉取权限时不指定资源 ARN,因此,权限范围限定为单个 ECR 存储库的执行角色会无法通过验证,即使它在运行时本可以正常工作。正因如此,示例中的角色获得了范围更广的 ECR 权限。
最后,示例采用了 v3 的一种做法:使用 PipelineSession(default_bucket=your_bucket) 构建流水线,会将 SDK 管理的上传内容(例如步骤代码)存入流水线自己的存储桶,使所有工件都集中在一处,执行角色的权限也只需覆盖单个存储桶。(授予该角色访问 SDK 默认存储桶的权限也可以。这里的做法更整洁。)
符合缺陷数据特性的质量门禁
每个训练完成的模型都会在注册前接受评估,而且评估针对的是导出的 ONNX 工件,而不是内存中的 PyTorch 模型。这一区别很重要:接受门禁检查的工件就是最终交付的工件,因此,即使训练过程看似正常,导出问题也绝不会让损坏的模型蒙混过关。
异常检测分支要求 AUROC 至少达到 0.8。分类器分支则同时检查两项指标,而指标的选择取决于缺陷数据的典型特征:严重的类别不平衡。在此示例中,大约 90 percent 的图像是良品,这符合生产线的实际情况。在这种分布下,普通准确率毫无意义,因为将每张图像都预测为“良品”的模型能获得约 0.9 的分数,却检测不出任何缺陷。平衡准确率,即各类别召回率的平均值,可以纠正这一问题,但它也有自己的盲点。一个对良品的召回率达到满分、对缺陷的召回率仅为 20 percent 的模型,仍能获得 (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 会对其条件执行逻辑与运算,因此两项条件都必须通过。这些阈值是流水线参数,每次执行时均可调整,无需重新部署。如果未通过门禁,FailStep 会将此次执行标记为失败,该分支的模型也就无法进入注册表(两个分支独立进行门禁检查,因此另一分支的模型仍可能完成注册)。设备群会继续为该摄像头提供先前已获批准的模型。
训练代码从另一个角度应对同样的类别不平衡问题。分类器是经过微调的 resnet18,采用 ImageNet 权重和二分类头。在单摄像头数据集的规模下,微调预训练权重是防止模型退化为只预测多数类的关键。训练代码使用基于类别频率倒数的 WeightedRandomSampler 来平衡训练批次,而损失函数不加权。将加权损失与加权采样结合使用,会导致模型过度偏向少数类。
在流水线内使用 SageMaker Neo 为边缘端编译
注意前面代码中门控步骤的 if_steps:通过门控后,先执行编译,注册则依赖编译完成。两个训练脚本都导出具有动态批次轴的 ONNX 模型。质量门控在 ONNX Runtime 下对这些导出的产物进行评分,然后由 CompileAnomalibNeo 或 CompileClassifierNeo 步骤将通过门控的产物交给 SageMaker Neo。固定顺序是先门控,再编译,最后注册。模型只有在针对设备群的边缘目标编译成功后,才能成为 Approved 状态的模型包,因此注册表绝不会列出设备无法运行的模型。
编译步骤封装了三个容易出错的细节。Neo 要求输入 tar 包仅包含模型文件,但训练作业输出的 model.tar.gz 除了 ONNX 模型,还包含检查点和指标。该步骤首先重新打包出仅含 ONNX 的 tar 包。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 分钟,并产生约 0.07 美元的计费ml.m5.xlarge时间(美国东部(俄亥俄)区域(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 编译,最后才注册。Model Registry 的组织方式利用摄像头元数据,将共享模型包组转化为各摄像头可用于热启动的历史记录,而 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
