模型导出与验证

相关源文件

生成此 wiki 页面时使用了以下文件作为上下文:

model_utils 包提供一组工具,用于将 LeRobot 策略从 PyTorch 训练环境迁移到面向专用 NPU(Neural Processing Unit)硬件的高性能部署环境。它支持将模型导出为 ONNX 和专用二进制格式(Ascend OM、Rockchip RKNN),并提供严格的跨平台精度验证工具,用于检测量化或转换过程中的数值漂移。

架构与数据流

模型部署流水线遵循结构化路径,从高层 PyTorch 模型转换为适合嵌入式执行的优化二进制产物。

部署流水线

下图展示从已训练 LeRobot checkpoint 到已验证硬件专用模型的流程。

模型转换与验证流程

        graph TD
    subgraph "Training Environment (GPU/CPU)"
        A["LeRobot Checkpoint (model.safetensors)"] --> B["export_onnx_atc.py"]
        A --> C["export_onnx_rknn.py"]
        A --> D["export_onnx_3403.py"]
        B --> E["General ONNX"]
        C --> F["RKNN-Optimized ONNX"]
        D --> G["3403-Optimized ONNX"]
    end

    subgraph "NPU Toolchains"
        E -->|"atc tool"| H["Ascend OM Model"]
        F -->|"rknn-toolkit2"| I["RKNN Model"]
        G -->|"3403 SDK"| J["3403 OM Model"]
    end

    subgraph "Accuracy Validation"
        K["loss_compare.py (Generate Target)"] -->|"JSON Baseline"| L["loss_compare.py (Compute Loss)"]
        A --> K
        H --> L
        I --> L
        L --> M["L1 / Cosine Sim Report"]
    end
    

来源:src/model_utils/model_utils/export_onnx_atc.py:106-152src/model_utils/model_utils/export_onnx_rknn.py:117-129src/model_utils/model_utils/loss_compare.py:17-26

模型导出工具

该包为不同硬件后端提供专用导出脚本,处理输入 tensor 重构等模型专属需求。

export_onnx_atc.py (General Ascend)

该脚本自动执行标准 Ascend 硬件(例如 Ascend 310P3)的端到端转换。

export_onnx_rknn.py (Rockchip RK3588)

该脚本面向 Rockchip NPU 定制,实现特定优化以降低开销。

export_onnx_3403.py (Ascend 3403)

该脚本面向 Ascend 3403 平台,重点生成简化 ONNX 图。

来源:src/model_utils/model_utils/export_onnx_atc.py:93-103src/model_utils/model_utils/export_onnx_rknn.py:15-30src/model_utils/model_utils/export_onnx_3403.py:21-31

跨平台验证 (loss_compare)

模型转换期间数值精度可能下降,例如 FP32 转 FP16 或 INT8。loss_compare.py 提供一个框架,用于将 NPU 输出与 PyTorch baseline 比较。

精度指标

该工具计算 L1 Loss 和 Cosine Similarity。

PI05 分布式评估

对于 PI05 等 flow-matching 模型,由于 ODE solver 中的混沌放大,逐点 L1 比较通常会产生误导 src/model_utils/model_utils/pi05_dist_metrics.py:6-11。该包包含 pi05_dist_metrics.py,提供对混沌更稳定的指标:

PI05 验证逻辑

        graph TD
    subgraph "loss_compare.py"
        A["Load Preds (NPU)"] --> B{"Policy Type?"}
        B -->|"Standard (ACT)"| C["Pointwise L1 / Cosine"]
        B -->|"PI05"| D["pi05_dist_metrics.py"]
    end

    subgraph "pi05_dist_metrics.py"
        D --> E["Method A: Wasserstein-1"]
        D --> F["Method C: First-frame Cosine"]
        E --> G["Scale-invariant Ratio (W1/std)"]
        G --> H["Verdict: Excellent/Good/Poor"]
    end
    

来源:src/model_utils/model_utils/loss_compare.py:65-71src/model_utils/model_utils/pi05_dist_metrics.py:25-33src/model_utils/model_utils/pi05_dist_metrics.py:116-124

专用 NPU Runtime:PI05OMModel

对于 Ascend 上的 PI05 部署,PI05OMModel 实现了高性能 runtime,用于最小化 VLM(Vision Language Model)与 Action Expert 之间的数据移动。

来源:src/inference_service/inference_service/core/ascend_om/pi05/PI05OMModel.py:71-180src/inference_service/inference_service/core/ascend_om/pi05/policy_wrapper.py:11-15