SmolVLA ONNX (LIBERO-10)

转换为 ONNX 格式的 SmolVLA 模型，拆分为 3 个模块化组件，适用于 CPU 和边缘设备的高效推理。

模型概述

原始模型： HuggingFaceVLA/smolvla_libero

SmolVLA 是一个专为机器人操作任务设计的紧凑型视觉-语言-动作模型。本仓库包含针对部署优化的 ONNX 转换版本。

模型架构

该模型拆分为 3 个 ONNX 组件：

1. 视觉编码器 (vision_encoder.onnx)

将视觉观察编码为嵌入向量。

输入：

• pixel_values：形状 [1, 3, 512, 512] - RGB 图像（归一化至 [-1, 1]）

输出：

• image_embeddings：形状 [1, 64, 960] - 视觉嵌入向量

2. LLM 主干网络 (llm_backbone.onnx)

通过语言模型处理视觉嵌入向量和语言标记。

输入：

• image_embs_1：形状 [1, 64, 960] - 视觉嵌入向量（第一份副本）
• image_embs_2：形状 [1, 64, 960] - 视觉嵌入向量（第二份副本）
• lang_tokens：形状 [1, 48] - 语言标记 ID
• lang_masks：形状 [1, 48] - 语言注意力掩码
• state：形状 [1, 32] - 初始状态

输出：

• kv_keys：形状 [32, 1, 177, 5, 64] - 注意力的键缓存
• kv_values：形状 [32, 1, 177, 5, 64] - 注意力的值缓存
• prefix_pad_masks：形状 [1, 177] - 前缀填充掩码

3. 动作头 (action_head.onnx)

根据 LLM 主干网络的输出预测动作速度。

输入：

• noisy_actions：形状 [1, 50, 32] - 带噪声的动作样本
• timestep：形状 [1] - 扩散时间步
• prefix_pad_masks：形状 [1, 177] - 填充掩码
• kv_keys：形状 [32, 1, 177, 5, 64] - 键缓存
• kv_values：形状 [32, 1, 177, 5, 64] - 值缓存

输出：

• velocity：形状 [1, 50, 32] - 预测的动作速度

性能

LIBERO-10 评估： 约 40% 成功率（基础 SFT 模型，与原始 PyTorch 版本一致）

文件

• vision_encoder.onnx + vision_encoder.onnx.data - 视觉编码器模型
• llm_backbone.onnx + llm_backbone.onnx.data - LLM 主干网络模型
• action_head.onnx + action_head.onnx.data - 动作头模型
• infer_libero_episode_onnx_cpu.py - LIBERO 场景推理脚本
• config.json - 模型配置
• policy_preprocessor.json - 预处理器配置
• policy_postprocessor.json - 后处理器配置
• policy_preprocessor_step_5_normalizer_processor.safetensors - 输入归一化权重
• policy_postprocessor_step_1_unnormalizer_processor.safetensors - 输出反归一化权重

使用方法

基础推理

from infer_libero_episode_onnx_cpu import SmolVLAONNXInference

# 初始化推理引擎
inference = SmolVLAONNXInference(
    vision_encoder_path=\"vision_encoder.onnx\",
    llm_backbone_path=\"llm_backbone.onnx\",
    action_head_path=\"action_head.onnx\",
    config_path=\"config.json\",
    preprocessor_config_path=\"policy_preprocessor.json\",
    postprocessor_config_path=\"policy_postprocessor.json\",
    normalizer_path=\"policy_preprocessor_step_5_normalizer_processor.safetensors\",
    unnormalizer_path=\"policy_postprocessor_step_1_unnormalizer_processor.safetensors\"
)

# 对图像和语言指令运行推理
image = ...  # PIL 图像或 numpy 数组 [H, W, 3]
instruction = \"pick up the red cube\"
actions = inference.infer(image, instruction)

完整实现细节请参见 infer_libero_episode_onnx_cpu.py。

依赖要求

• onnxruntime - 用于 ONNX 模型推理
• transformers - 用于分词
• safetensors - 用于加载归一化权重
• numpy - 用于数值运算
• PIL - 用于图像处理

安装

pip install onnxruntime transformers safetensors numpy pillow

模型转换

这些 ONNX 模型是从原始 PyTorch SmolVLA 模型转换而来，使用以下设置进行 ONNX 导出：

• Opset 版本：17
• 输入/输出形状如上所述
• 针对 CPU 推理优化

许可证

本模型遵循与原始 SmolVLA 模型相同的许可证。许可证详情请参阅原始仓库。

引用

如果您使用本模型，请引用原始 SmolVLA 论文：

@article{smolvla,
  title={SmolVLA: Compact Vision Language Action Models for Robotic Manipulation},
  author={...},
  year={2024}
}

参考资料

• 原始模型：https://huggingface.co/HuggingFaceVLA/smolvla_libero
• LIBERO 基准测试：https://libero-project.github.io/