lerobot/act_aloha_sim_transfer_cube_human · 昇腾 NPU 适配

ACT (Action Chunking with Transformers) 策略模型, 在 ALOHA 仿真器上针对立方体转移任务训练。 本仓库提供昇腾 Ascend NPU (torch_npu) 推理适配, 模型、输入张量全部运行在 npu:0 设备, 不含任何 CUDA 调用。

模型架构

• Vision Backbone: ResNet18
• Encoder: Transformer Encoder (4 层, dim_model=512, n_heads=8, ff=3200)
• Decoder: Transformer Decoder (1 层) + 可学习 action queries
• VAE: 训练时使用, 推理时 latent 置零
• 输出: action chunk, 100 步 × 14 维

权重

• 原始权重: https://ai.gitcode.com/hf_mirrors/lerobot/act_aloha_sim_transfer_cube_human
• 本地缓存: ./model_weights/lerobot/act_aloha_sim_transfer_cube_human/ (config.json + model.safetensors)
• 缺失时自动从镜像下载, 优先复用本地缓存

环境

• Python 3.11, torch + torch_npu (CANN)
• 依赖: pip install fastapi uvicorn python-multipart pillow safetensors

用法

命令行单图推理

python inference.py --img <图片路径> [--state "0,0,...,0"]

FastAPI 服务

python inference.py --server
# 浏览器访问 http://localhost:8000 上传图片, POST /predict 返回推理 JSON

验证

• 模型权重全部加载成功 (234 个张量)
• 输出 action chunk: (100, 14)
• 设备: npu:0