跳到正文
AWS Machine Learning Blog· Kiowa Jackson·· 2 小时前精选AI 评分70

AWS SageMaker SDK v3 构建边缘端多摄像头计算机视觉模型生命周期流程

Fleet-wide computer vision lifecycle at the edge with SageMaker SDK v3

AI 导读

AWS Machine Learning Blog 给出基于 Amazon SageMaker Python SDK v3 的多摄像头计算机视觉模型生命周期架构,为每个摄像头训练异常检测器和缺陷分类器,并对导出的 ONNX 模型执行质量门禁、SageMaker Neo 编译后再注册。

推荐理由

文章给出多摄像头模型生命周期的可复用架构,将质量门禁、边缘编译和并发编排串成完整流程。

正文 · AI 翻译

全车队计算机视觉模型生命周期管理,是质量检测概念验证一旦成功就会出现的问题。一台摄像机和一个模型运行良好。但随后工厂要求在 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 存储库 中获取。

解决方案概述

下图展示了该示例面向整个车队的模型生命周期。编排层负责排队重新训练请求并限制并发任务数。每台摄像机对应一个训练管道,该管道准备数据,训练并评估异常检测器和分类器,并为边缘设备生成带版本的模型构件。

Solution architecture: SQS FIFO queue triggers AWS Lambda, which starts a Step Functions state machine with an Amazon DynamoDB semaphore for concurrency control. The state machine launches a 13-step SageMaker Pipeline with parallel entry steps (Preprocess, Fetch Latest Anomaly, Fetch Latest Classifier), dual training branches with ONNX evaluation, quality gates, Neo compilation, and Model Registry registration. Compiled artifacts land in S3. Completions and failures converge on an SNS topic.

图 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 Jackson

Kiowa 是 AWS ProServe 的高级机器学习工程师。他专注于将计算机视觉与智能体 AI 相结合的视觉智能系统,并推动这些系统在企业工业领域投入生产。

来源:AWS Machine Learning Blog · aws.amazon.com