qing
a2ba797410
- pybot：recorder 逐帧落盘 imu.npy（T×8，roll,pitch,ax,ay,az,gx,gy,gz，与帧/动作同调用对齐）；run_loop 对声明 imu_dim 的 predictor 以 predict(..., imu=) 调用（鸭子类型，旧两参调用向后兼容） - actpred：新增 --imu-dim；model 序列 [img×F, act×H, (IMU), PRED]（imu_dim>0 时 F+H+2，IMU token 位于 PRED 前）；dataset/train 读 imu.npy（缺失按零向量兜底）；predict/ONNX/Hailo 三后端与 export/gen_calibration/dfc_compile/_npu_test 部署链适配 imu 输入 - 姿态按物理量程固定缩放归一化（imu_norm.py，训练/推理共用，零状态） - 测试：actpred 128 passed / pybot loop+recorder 28 passed
181 lines
7.2 KiB
Python
181 lines
7.2 KiB
Python
"""actpred 部署版 ONNX → HEF 编译脚本（Hailo DFC 5.3，仅 x86 Linux）
|
||
|
||
流程（对应 DFC 5.x CLI）：
|
||
hailo parser onnx → HAR（网络翻译）
|
||
hailo optimize → 量化优化后的 HAR（需校准集 npz）
|
||
hailo compiler → HEF
|
||
|
||
用法:
|
||
# 1. 导出部署版 ONNX（DFC 友好图结构）
|
||
python export_deploy.py --model checkpoints/pretrain_best_inference.pt
|
||
# 2. 生成校准集（首次）
|
||
python gen_calibration.py --deploy
|
||
# 3. 编译
|
||
python dfc_compile.py
|
||
|
||
关键点（踩坑记录）：
|
||
- 训练态 ONNX（5D 输入、AdaptiveAvgPool、MultiheadAttention 的 head
|
||
reshape）DFC 无法解析，必须用 export_deploy.py 的改造版图
|
||
- pred_token 必须作为图输入：常量 initializer 会被 DFC 误转置
|
||
- 必须显式 --start/--end-node-names 与 --input-format，且禁用
|
||
onnxsim（本图简化无收益，反而引入 reshape）
|
||
- 校准集键名必须匹配 HN 层名 actpred_deploy/input_layerN
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import tarfile
|
||
import tempfile
|
||
|
||
|
||
NET_NAME = "actpred_deploy"
|
||
|
||
|
||
def start_nodes(frame_count: int, with_imu: bool = False) -> list:
|
||
"""export_deploy.py 的输入命名约定：frame{i} × frame_count + action_history
|
||
+ pred_token（R2 imu_dim>0 时末尾追加 imu）"""
|
||
nodes = ([f"frame{i}" for i in range(frame_count)]
|
||
+ ["action_history", "pred_token"])
|
||
if with_imu:
|
||
nodes.append("imu")
|
||
return nodes
|
||
|
||
|
||
def hailo(*args, venv_bin: str):
|
||
cmd = [os.path.join(venv_bin, "hailo"), *args]
|
||
env = dict(os.environ, PATH=venv_bin + os.pathsep + os.environ["PATH"])
|
||
return subprocess.run(cmd, env=env, check=True)
|
||
|
||
|
||
def _har_hn_layers(har_path: str) -> dict:
|
||
"""解包 HAR 并读取主 hn 文件的 layers（优化产物含 fp/native 变体，取主图）"""
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
with tarfile.open(har_path) as tf:
|
||
tf.extractall(tmp)
|
||
hn_files = [f for f in os.listdir(tmp) if f.endswith(".hn")]
|
||
if not hn_files:
|
||
raise ValueError(f"{har_path} 中无 .hn 文件")
|
||
main = [f for f in hn_files
|
||
if not f.endswith(".fp.hn") and not f.endswith(".native.hn")]
|
||
hn_file = main[0] if main else hn_files[0]
|
||
with open(os.path.join(tmp, hn_file)) as f:
|
||
return json.load(f)["layers"]
|
||
|
||
|
||
def _input_layer_index(name: str) -> int:
|
||
m = re.search(r"input_layer(\d+)$", name)
|
||
return int(m.group(1)) if m else 0
|
||
|
||
|
||
def har_input_layer_names(har_path: str) -> list:
|
||
"""从 HAR 主 hn 图读输入层名，按 input_layerN 数字后缀排序"""
|
||
layers = _har_hn_layers(har_path)
|
||
names = [n for n, l in layers.items() if l["type"] == "input_layer"]
|
||
names.sort(key=_input_layer_index)
|
||
return names
|
||
|
||
|
||
def check_input_connectivity(har_path: str, input_layers: list):
|
||
"""检查 HAR 图中每个输入层是否可达任一输出层
|
||
|
||
输入层被优化断开时编译照常成功，但 HEF 输出恒定值且无任何报错
|
||
（2026-08-13 实测踩坑），必须在编译前拦截。DFC 5.3 的 HAR 不含
|
||
.edges 文件，连通性改从主 hn 图的 layer output 边 BFS 判定。
|
||
"""
|
||
layers = _har_hn_layers(har_path)
|
||
adj = {n: l.get("output", []) for n, l in layers.items()}
|
||
targets = {n for n, l in layers.items() if l["type"] == "output_layer"}
|
||
if not targets:
|
||
print(f"WARNING: {har_path} 图中无输出层，跳过连通性检查")
|
||
return
|
||
disconnected = []
|
||
for name in input_layers:
|
||
seen, stack = set(), [name]
|
||
while stack:
|
||
cur = stack.pop()
|
||
if cur in seen:
|
||
continue
|
||
seen.add(cur)
|
||
stack.extend(adj.get(cur, []))
|
||
if not (seen & targets):
|
||
disconnected.append(name)
|
||
if disconnected:
|
||
sys.exit(
|
||
f"错误: {har_path} 中输入层不可达输出层（输入被优化断开，"
|
||
f"编译出的 HEF 将输出恒定值）: {disconnected}")
|
||
print(f"输入连通性检查通过: {len(input_layers)} 个输入层均可达输出层")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Compile actpred deploy ONNX to HEF")
|
||
parser.add_argument("--onnx", default="checkpoints/actpred_deploy.onnx")
|
||
parser.add_argument("--calib-npz", default="data/calibration_deploy.npz")
|
||
parser.add_argument("--output", default="checkpoints/actpred_pretrain.hef")
|
||
parser.add_argument("--hw-arch", default="hailo10h",
|
||
choices=["hailo15h", "hailo15m", "hailo15l",
|
||
"hailo10h", "hailo10p"])
|
||
parser.add_argument("--frame-count", type=int, default=3,
|
||
help="帧输入个数（与导出 ONNX 的 frame_len 一致，"
|
||
"默认 3）")
|
||
parser.add_argument("--with-imu", action="store_true",
|
||
help="R2：导出 ONNX 含 imu 输入时开启（start nodes "
|
||
"末尾追加 imu，与 export_deploy.py 命名一致）")
|
||
parser.add_argument("--venv-bin", default="venv_dfc/bin",
|
||
help="安装了 hailo-sdk 的 venv bin 目录")
|
||
args = parser.parse_args()
|
||
|
||
venv_bin = os.path.abspath(args.venv_bin)
|
||
if not os.path.exists(os.path.join(venv_bin, "hailo")):
|
||
sys.exit(f"错误: {venv_bin}/hailo 不存在，请先安装 hailo-sdk")
|
||
|
||
nodes = start_nodes(args.frame_count, args.with_imu)
|
||
input_formats = ["NHWC"] * len(nodes)
|
||
|
||
workdir = tempfile.mkdtemp(prefix="dfc_")
|
||
har_raw = os.path.join(workdir, "parsed.har")
|
||
har_opt = os.path.join(workdir, "optimized.har")
|
||
hef_dir = os.path.join(workdir, "hef")
|
||
os.makedirs(hef_dir)
|
||
|
||
try:
|
||
print("[1/3] parser: ONNX → HAR")
|
||
hailo("parser", "onnx", args.onnx, "--hw-arch", args.hw_arch, "-y",
|
||
"--har-path", har_raw,
|
||
"--start-node-names", *nodes,
|
||
"--end-node-names", "predicted_action",
|
||
"--input-format",
|
||
*[f"{n}={f}" for n, f in zip(nodes, input_formats)],
|
||
"--disable-onnx-simplifier", venv_bin=venv_bin)
|
||
|
||
# 校验未被推荐机制截断：输出层应来自预测头（conv 而非 ew_add）
|
||
names = har_input_layer_names(har_raw)
|
||
if len(names) != len(nodes):
|
||
sys.exit(f"错误: HAR 输入层数 {len(names)} != {len(nodes)}，解析可能被截断")
|
||
check_input_connectivity(har_raw, names)
|
||
|
||
print("[2/3] optimize: 量化校准")
|
||
hailo("optimize", har_raw, "--hw-arch", args.hw_arch,
|
||
"--calib-set-path", os.path.abspath(args.calib_npz),
|
||
"--output-har-path", har_opt, venv_bin=venv_bin)
|
||
check_input_connectivity(har_opt, har_input_layer_names(har_opt))
|
||
|
||
print("[3/3] compiler: HAR → HEF")
|
||
hailo("compiler", har_opt, "--hw-arch", args.hw_arch,
|
||
"--output-dir", hef_dir, venv_bin=venv_bin)
|
||
|
||
hef_src = os.path.join(hef_dir, f"{NET_NAME}.hef")
|
||
shutil.copy(hef_src, args.output)
|
||
shutil.copy(har_opt, args.output.replace(".hef", "_optimized.har"))
|
||
print(f"HEF 已保存: {args.output} ({os.path.getsize(args.output)} bytes)")
|
||
finally:
|
||
shutil.rmtree(workdir, ignore_errors=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|