qing
0d16832612
训练态 ONNX（5D 输入/AdaptiveAvgPool/注意力 head reshape）无法被 DFC 解析， 新增 export_deploy.py 做权重不变的等价图改造：帧分支展开 CNN、img_proj 折叠为 1x1 卷积、注意力按头展开、全程 (1,1,S,d) 无 transpose 布局。 dfc_compile.py 一键走 parser→optimize→compiler（DFC 5.3）， gen_calibration.py --deploy 生成部署布局校准集。 已验证 110 层完整网络编译为 hailo10h HEF（1.8MB，4 contexts）。
116 lines
4.6 KiB
Python
116 lines
4.6 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 shutil
|
||
import subprocess
|
||
import sys
|
||
import tarfile
|
||
import tempfile
|
||
|
||
|
||
NET_NAME = "actpred_deploy"
|
||
START_NODES = ["frame0", "frame1", "frame2", "action_history", "pred_token"]
|
||
INPUT_FORMATS = ["NHWC", "NHWC", "NHWC", "NHWC", "NHWC"]
|
||
|
||
|
||
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_input_layer_names(har_path: str) -> list:
|
||
"""从 HAR（tar 包）中读 hn 文件，按 input_layerN 排序返回层名"""
|
||
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")]
|
||
with open(os.path.join(tmp, hn_files[0])) as f:
|
||
hn = json.load(f)
|
||
layers = [(name, l) for name, l in hn["layers"].items()
|
||
if l["type"] == "input_layer"]
|
||
layers.sort(key=lambda x: x[0])
|
||
return [name for name, _ in 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("--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")
|
||
|
||
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", *START_NODES,
|
||
"--end-node-names", "predicted_action",
|
||
"--input-format",
|
||
*[f"{n}={f}" for n, f in zip(START_NODES, INPUT_FORMATS)],
|
||
"--disable-onnx-simplifier", venv_bin=venv_bin)
|
||
|
||
# 校验未被推荐机制截断：输出层应来自预测头（conv 而非 ew_add）
|
||
names = har_input_layer_names(har_raw)
|
||
if len(names) != len(START_NODES):
|
||
sys.exit(f"错误: HAR 输入层数 {len(names)} != {len(START_NODES)}，解析可能被截断")
|
||
|
||
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)
|
||
|
||
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()
|