cyy/go1_pro_deploy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

bpu

Browse Source
This commit is contained in:
cyy_mac
2026-07-27 23:17:55 +08:00
parent ebbe9d5432
commit e017bce6e2
20 changed files with 3062 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
deploy_45dim_rl_gym/bpu_quantization/README.md Normal file
View File

@@ -0,0 +1,237 @@
# RDK X5 BPU 量化流程
这个目录用于把 `../policy_robotlab_15000.onnx` 转成 RDK X5 可运行的
Horizon runtime `.bin`。量化在 Mac 上用 CPU Docker 完成，板端只做离线
`hrt_model_exec` 验证，暂时不要直接接入实机控制。
参考资料：
- D-Robotics 官方 X5 工具链说明：`https://developer.d-robotics.cc/rdk_x_doc/Advanced_development/toolchain_development/overview?v=3.5.0&p=RDK+X5`
- D-Robotics 论坛 WTW/Go2/X5 流程：`https://forum.d-robotics.cc/t/topic/28338`
官方工具链对 ONNX 的关键限制是：`ir_version <= 7`、`opset10/11`、固定
4 维输入，且 N 维只能为 1。因此这里不能直接拿原始 RobotLab ONNX 编译，
需要先降级 opset，再把 `[1, 450]` 输入包成固定 4D `NCHW`：
`[1, 1, 1, 450]`。
## 当前状态
已经完成 `policy_robotlab_15000.onnx` 和 `policy_robotlab_6500.onnx` 的 int16
量化。当前 BPU 部署默认使用 6500 版本：
- 原始模型：`../policy_robotlab_6500.onnx`
- 原始输入：`obs [1, 450]`
- BPU 编译输入：`obs_4d [1, 1, 1, 450]`
- BPU 输出：`actions [1, 12, 1, 1]`
- Docker 镜像：`openexplorer/ai_toolchain_ubuntu_20_x5_cpu:v1.2.8`
- `hb_mapper`：`1.24.3`
- `hbdk`：`3.49.15`
- 当前产物：`mapper_output_6500_gemm/policy_robotlab_6500_int16_gemm.bin`
`mapper_output*/`、`.hb_check/`、校准数据、中间 ONNX、`hb_mapper` 日志都被
`.gitignore` 忽略；需要时按下面步骤重新生成。仓库里只保留脚本和 YAML 配置。
6500 量化使用 `calibration_data_fast64/` 的 64 个真实样本。原因是
`hb_mapper` 会先尝试 calibration batch 8，但当前 4D featuremap 包装会被工具链
内部改成固定 batch 的 reshape，batch 8 失败后会退回 batch 1；用 64 样本可以把
校准时间从 512 次 batch1 显著降下来。
## 1. 生成校准数据
在本机 `deploy_go1_pro` 根目录执行：
```bash
/Users/chenyouyuan/miniconda/envs/free_dog_sdk/bin/python3.10 \
deploy_45dim_rl_gym/bpu_quantization/make_calibration_data.py \
--max-samples 512 \
--overwrite
```
脚本会从 `logs/robotlab_go1_deploy_*/steps.jsonl` 中提取 RL 状态下的真实
`obs_single`，按部署代码相同的 10 帧历史顺序重建 450 维输入，并写成
`float32` feature-map `.bin`。
本次校准数据：
- 候选 RL 输入：`62606`
- 选中样本：`512`
- 单个样本形状：`1x1x1x450`
- 元数据：`deploy_45dim_rl_gym/bpu_quantization/calibration_data_metadata.json`
## 2. 生成 opset11 和 4D ONNX
启动 Docker Desktop 后，在 `deploy_go1_pro` 根目录执行：
```bash
docker run --rm --platform linux/amd64 \
-v "$PWD:/workspace/deploy_go1_pro" \
openexplorer/ai_toolchain_ubuntu_20_x5_cpu:v1.2.8 \
bash -lc 'cd /workspace/deploy_go1_pro/deploy_45dim_rl_gym/bpu_quantization &&
python3 downgrade_policy_to_opset11.py \
--input ../policy_robotlab_15000.onnx \
--output policy_robotlab_15000_opset11.onnx &&
python3 make_bpu_4d_onnx.py \
--input policy_robotlab_15000_opset11.onnx \
--output policy_robotlab_15000_bpu4d.onnx &&
python3 compare_4d_onnx.py \
--flat-onnx ../policy_robotlab_15000.onnx \
--bpu4d-onnx policy_robotlab_15000_bpu4d.onnx \
--limit 256'
```
本次对比结果：256 个真实输入上 `max_abs_diff = 0`，说明降级和 4D 包装没有
改变浮点 ONNX 输出。
## 3. 可选：替换 MoE grouped Conv
原始导出的 MoE expert 末端有一个 `group=8, kernel=1` 的 `Conv1d`：
```text
[1, 2048] -> Unsqueeze -> Conv(group=8, kernel=1) -> Squeeze -> [1, 256]
```
在 X5 checker 里，这段会切到 CPU float，导致 BPU 子图被拆开。这个 Conv
等价于一个块对角 `Gemm`，可以不重训直接改图：
```bash
docker run --rm --platform linux/amd64 \
-v "$PWD:/workspace/deploy_go1_pro" \
openexplorer/ai_toolchain_ubuntu_20_x5_cpu:v1.2.8 \
bash -lc 'cd /workspace/deploy_go1_pro/deploy_45dim_rl_gym/bpu_quantization &&
python3 replace_group_conv_with_gemm.py \
--input policy_robotlab_15000_bpu4d.onnx \
--output policy_robotlab_15000_bpu4d_gemm.onnx &&
python3 compare_4d_onnx.py \
--flat-onnx ../policy_robotlab_15000.onnx \
--bpu4d-onnx policy_robotlab_15000_bpu4d_gemm.onnx \
--limit 512'
```
本次浮点对比结果：
- `max_abs_diff = 1.43e-6`
- `max_mean_abs_diff = 3.53e-7`
这是浮点舍入级误差，可以认为图变换等价。
注意：这里没有把 `ReduceL2/Pow/Reshape/Sqrt` 这类 norm 算子挪到后处理，
因为它们在 actor 之前，属于策略中间计算，不能当成输出后处理裁掉。当前
`hb_mapper makertbin` 已经能把这些 norm 节点放到 BPU 上。
## 4. checker 检查
```bash
docker run --rm --platform linux/amd64 \
-v "$PWD:/workspace/deploy_go1_pro" \
openexplorer/ai_toolchain_ubuntu_20_x5_cpu:v1.2.8 \
bash -lc 'cd /workspace/deploy_go1_pro/deploy_45dim_rl_gym/bpu_quantization &&
hb_mapper checker \
--model policy_robotlab_15000_bpu4d_gemm.onnx \
--model-type onnx \
--march bayes-e \
--input-shape obs_4d 1x1x1x450'
```
checker 可以通过。原始模型中 MoE expert 的 `Elu/Reshape/Conv` CPU fallback
会被消掉；checker 阶段仍可能显示 gating `Softmax` 或输出 reshape 的 CPU
适配，正式 int16 编译时 `Softmax` 会被量化拆成 BPU 子算子。
## 5. 编译 int16 `.bin`
```bash
docker run --rm --platform linux/amd64 \
-v "$PWD:/workspace/deploy_go1_pro" \
openexplorer/ai_toolchain_ubuntu_20_x5_cpu:v1.2.8 \
bash -lc 'cd /workspace/deploy_go1_pro/deploy_45dim_rl_gym/bpu_quantization &&
hb_mapper makertbin \
--config policy_robotlab_15000_int16_gemm.yaml \
--model-type onnx'
```
优化版 `hb_mapper` 输出的量化精度：
```text
Output Cosine Similarity L1 Distance L2 Distance Chebyshev Distance
actions 0.999910 0.009500 0.003393 0.024229
```
优化版正式 int16 编译后，日志中列出的策略节点全部在同一个 BPU 子图
`id(0)` 上，原始模型的 CPU `Elu/Reshape/Conv` fallback 已消失。
优化版 BPU 子图编译估计延迟：
- subgraph0：`424.5 us`
原始版 BPU 子图估计延迟是 `284.3 us + 80.8 us`，但带 CPU/hybrid 切换。
优化版单子图的编译估计延迟略高，实际是否更快要以板端 `hrt_model_exec perf`
为准。
## 6. 板端离线验证
把 `.bin` 和一个校准输入传到板端，例如：
```bash
ssh root@192.168.150.167 'mkdir -p /root/go1_pro_deploy/bpu_quant_test'
scp \
deploy_45dim_rl_gym/bpu_quantization/mapper_output_gemm/policy_robotlab_15000_int16_gemm.bin \
deploy_45dim_rl_gym/bpu_quantization/calibration_data/00000.bin \
root@192.168.150.167:/root/go1_pro_deploy/bpu_quant_test/
```
板端查看模型：
```bash
cd /root/go1_pro_deploy/bpu_quant_test
hrt_model_exec model_info --model_file policy_robotlab_15000_int16_gemm.bin
```
板端测速：
```bash
hrt_model_exec perf \
--model_file policy_robotlab_15000_int16_gemm.bin \
--model_name policy_robotlab_15000_int16_gemm \
--input_file 00000.bin \
--frame_count 1000 \
--thread_num 1
```
原始版板端结果：
- 平均延迟：`1.498257 ms`
- 最大延迟：`2.942 ms`
- FPS：`664.77`
优化版板端结果：
- 平均延迟：`0.872819 ms`
- 最大延迟：`1.627 ms`
- 最小延迟：`0.613 ms`
- FPS：`1133.28`
同一板端之前测 CPU ONNX 大约是 `1.63 ms`。原始 BPU hybrid `.bin` 只快了一点；
Gemm 优化版移掉 MoE 中间 CPU fallback 后，板端离线平均延迟比 CPU ONNX 快约
46%，比原始 BPU hybrid 快约 42%。
原始版板端单样本输出与本机浮点 ONNX 对比：
- `max_abs = 0.02545`
- `mean_abs = 0.01040`
- `l2 = 0.04584`
优化版板端单样本输出与本机浮点 ONNX 对比：
- `max_abs = 0.02528`
- `mean_abs = 0.00962`
- `l2 = 0.04187`
这个误差对离线验证是可接受的，但还不足以直接上实机。
## 7. 安全结论
当前建议只做离线推理验证，不要把 `.bin` 接入真实机器人控制循环。原因：
- 原始 BPU 端到端速度提升很小，不能解决目前右前腿掉线、力矩保护、上下坡不稳这些核心问题。
- Gemm 优化版已经改善离线推理速度，但还没有做部署代码适配和悬空状态机测试。
- 输出 shape 从 `[1, 12]` 变成 `[1, 12, 1, 1]`，部署代码需要单独适配。
- 接入前至少要做更多 held-out 真实日志对比、悬空状态机测试，再进入地面低速测试。

76
deploy_45dim_rl_gym/bpu_quantization/calibration_data_metadata.json Normal file
View File

@@ -0,0 +1,76 @@
{
"format": "raw float32 feature-map",
"flat_shape": [
1,
450
],
"mapper_shape": [
1,
1,
1,
450
],
"history_len": 10,
"num_obs": 45,
"term_dims": [
3,
3,
3,
12,
12,
12
],
"selected_samples": 512,
"candidate_rl_inputs": 62606,
"seed": 20260727,
"runs": [
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260726_221258/steps.jsonl",
"samples": 6644
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260726_222533/steps.jsonl",
"samples": 292
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260726_230656/steps.jsonl",
"samples": 620
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260726_230816/steps.jsonl",
"samples": 13705
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260727_171014/steps.jsonl",
"samples": 275
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260727_171147/steps.jsonl",
"samples": 8928
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260727_171522/steps.jsonl",
"samples": 5737
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260727_172753/steps.jsonl",
"samples": 1759
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260727_173528/steps.jsonl",
"samples": 6049
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260727_175350/steps.jsonl",
"samples": 2498
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260727_180352/steps.jsonl",
"samples": 345
},
{
"steps": "/Users/chenyouyuan/cyy_ws/deploy_go1_pro/logs/robotlab_go1_deploy_20260727_180434/steps.jsonl",
"samples": 15754
}
]
}

46
deploy_45dim_rl_gym/bpu_quantization/compare_4d_onnx.py Normal file
View File

@@ -0,0 +1,46 @@
#!/usr/bin/env python3
"""Compare original flat ONNX outputs with the 4D BPU-wrapper ONNX."""
import argparse
from pathlib import Path
import numpy as np
import onnxruntime as ort
def load_samples(calibration_dir, limit):
paths = sorted(Path(calibration_dir).glob("*.bin"))[:limit]
if not paths:
raise FileNotFoundError(f"No calibration .bin files found in {calibration_dir}")
return [np.fromfile(path, dtype=np.float32).reshape(1, 450) for path in paths]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--flat-onnx", type=Path, required=True)
parser.add_argument("--bpu4d-onnx", type=Path, required=True)
parser.add_argument("--calibration-dir", type=Path, default=Path("calibration_data"))
parser.add_argument("--limit", type=int, default=64)
args = parser.parse_args()
flat = ort.InferenceSession(str(args.flat_onnx), providers=["CPUExecutionProvider"])
wrapped = ort.InferenceSession(str(args.bpu4d_onnx), providers=["CPUExecutionProvider"])
flat_input = flat.get_inputs()[0].name
wrapped_input = wrapped.get_inputs()[0].name
max_abs = 0.0
max_mean_abs = 0.0
for sample in load_samples(args.calibration_dir, args.limit):
out_flat = flat.run(None, {flat_input: sample})[0]
out_wrapped = wrapped.run(None, {wrapped_input: sample.reshape(1, 1, 1, 450)})[0]
diff = np.abs(out_flat - out_wrapped)
max_abs = max(max_abs, float(diff.max()))
max_mean_abs = max(max_mean_abs, float(diff.mean()))
print(f"checked={min(args.limit, len(list(args.calibration_dir.glob('*.bin'))))}")
print(f"max_abs_diff={max_abs:.9g}")
print(f"max_mean_abs_diff={max_mean_abs:.9g}")
if __name__ == "__main__":
main()

30
deploy_45dim_rl_gym/bpu_quantization/convert_opset.py Normal file
View File

@@ -0,0 +1,30 @@
#!/usr/bin/env python3
"""Convert policy ONNX to an older opset accepted by X5 hb_mapper."""
import argparse
from pathlib import Path
import onnx
from onnx import version_converter
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--opset", type=int, default=11)
parser.add_argument("--ir-version", type=int, default=7)
args = parser.parse_args()
model = onnx.load(str(args.input))
converted = version_converter.convert_version(model, args.opset)
converted.ir_version = min(converted.ir_version, args.ir_version)
args.output.parent.mkdir(parents=True, exist_ok=True)
onnx.save(converted, str(args.output))
print(f"Wrote {args.output}")
print(f"ir={converted.ir_version}")
print(f"opsets={[(op.domain, op.version) for op in converted.opset_import]}")
if __name__ == "__main__":
main()

91
deploy_45dim_rl_gym/bpu_quantization/downgrade_policy_to_opset11.py Normal file
View File

@@ -0,0 +1,91 @@
#!/usr/bin/env python3
"""Downgrade the current RobotLab policy ONNX from opset18 to opset11.
This is intentionally narrow: it only rewrites ops whose axes changed from an
attribute in opset11 to an input tensor in newer opsets. The policy graph uses
otherwise opset11-compatible operators.
"""
import argparse
from pathlib import Path
import numpy as np
import onnx
from onnx import helper, numpy_helper
AXES_INPUT_OPS = {"Unsqueeze", "Squeeze", "ReduceL2", "ReduceSum"}
REMOVE_ATTRS = {
"Reshape": {"allowzero"},
}
def constant_values(model):
values = {}
for initializer in model.graph.initializer:
values[initializer.name] = numpy_helper.to_array(initializer)
for node in model.graph.node:
if node.op_type != "Constant" or len(node.output) != 1:
continue
for attr in node.attribute:
if attr.name == "value":
values[node.output[0]] = numpy_helper.to_array(attr.t)
break
return values
def axes_from_value(values, name):
if name not in values:
raise KeyError(f"Cannot find constant axes tensor {name!r}")
axes = np.asarray(values[name]).astype(np.int64).reshape(-1).tolist()
return [int(axis) for axis in axes]
def set_opset(model, target):
for opset in model.opset_import:
if opset.domain in ("", "ai.onnx"):
opset.version = target
return
model.opset_import.append(helper.make_operatorsetid("", target))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--target-opset", type=int, default=11)
parser.add_argument("--ir-version", type=int, default=7)
args = parser.parse_args()
model = onnx.load(str(args.input))
values = constant_values(model)
converted = 0
for node in model.graph.node:
if node.op_type in REMOVE_ATTRS:
remove_names = REMOVE_ATTRS[node.op_type]
kept_attrs = [attr for attr in node.attribute if attr.name not in remove_names]
del node.attribute[:]
node.attribute.extend(kept_attrs)
if node.op_type not in AXES_INPUT_OPS or len(node.input) < 2:
continue
axes = axes_from_value(values, node.input[1])
del node.input[1:]
node.attribute.extend([helper.make_attribute("axes", axes)])
converted += 1
set_opset(model, args.target_opset)
model.ir_version = min(model.ir_version, args.ir_version)
args.output.parent.mkdir(parents=True, exist_ok=True)
onnx.save(model, str(args.output))
print(f"Wrote {args.output}")
print(f"converted_axes_nodes={converted}")
print(f"ir={model.ir_version}")
print(f"opsets={[(op.domain, op.version) for op in model.opset_import]}")
if __name__ == "__main__":
main()

61
deploy_45dim_rl_gym/bpu_quantization/make_bpu_4d_onnx.py Normal file
View File

@@ -0,0 +1,61 @@
#!/usr/bin/env python3
"""Wrap a flat policy ONNX input as a fixed 4D feature-map input for X5 hb_mapper."""
import argparse
from pathlib import Path
import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper
def tensor_shape(value_info):
return [dim.dim_value for dim in value_info.type.tensor_type.shape.dim]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--flat-dim", type=int, default=450)
parser.add_argument("--input-name", default="obs_4d")
args = parser.parse_args()
model = onnx.load(str(args.input))
if len(model.graph.input) != 1:
raise ValueError(f"Expected one input, got {len(model.graph.input)}")
old_input = model.graph.input[0]
old_name = old_input.name
old_shape = tensor_shape(old_input)
if old_shape != [1, args.flat_dim]:
raise ValueError(f"Expected input shape [1, {args.flat_dim}], got {old_shape}")
reshape_out_name = "__bpu_flat_input"
flatten_node = helper.make_node(
"Flatten",
inputs=[args.input_name],
outputs=[reshape_out_name],
name="Bpu4dInputFlatten",
axis=1,
)
for node in model.graph.node:
for index, name in enumerate(node.input):
if name == old_name:
node.input[index] = reshape_out_name
new_input = helper.make_tensor_value_info(
args.input_name, TensorProto.FLOAT, [1, 1, 1, args.flat_dim]
)
model.graph.input.remove(old_input)
model.graph.input.insert(0, new_input)
model.graph.node.insert(0, flatten_node)
args.output.parent.mkdir(parents=True, exist_ok=True)
onnx.save(model, str(args.output))
print(f"Wrote {args.output}")
if __name__ == "__main__":
main()

160
deploy_45dim_rl_gym/bpu_quantization/make_calibration_data.py Normal file
View File

@@ -0,0 +1,160 @@
#!/usr/bin/env python3
"""Build float32 BPU calibration inputs from recorded RobotLab deployment logs."""
import argparse
import json
import random
from collections import deque
from pathlib import Path
import numpy as np
NUM_OBS = 45
HISTORY_LEN = 10
ONNX_INPUT_DIM = NUM_OBS * HISTORY_LEN
TERM_DIMS = (3, 3, 3, 12, 12, 12)
HERE = Path(__file__).resolve().parent
DEFAULT_LOG_ROOT = HERE.parents[1] / "logs"
def build_onnx_input(history):
frames = list(history)
while len(frames) < HISTORY_LEN:
frames.insert(0, np.zeros(NUM_OBS, dtype=np.float32))
chunks = []
offset = 0
for dim in TERM_DIMS:
for frame in frames:
chunks.append(frame[offset:offset + dim])
offset += dim
return np.concatenate(chunks, dtype=np.float32)
def reservoir_add(samples, value, seen, max_samples, rng):
if len(samples) < max_samples:
samples.append(value)
return
replace_index = rng.randrange(seen)
if replace_index < max_samples:
samples[replace_index] = value
def collect_samples(log_paths, max_samples, seed):
rng = random.Random(seed)
samples = []
seen = 0
usable_runs = []
for steps_path in log_paths:
history = deque(maxlen=HISTORY_LEN)
run_seen = 0
was_rl = False
with steps_path.open("r", encoding="utf-8") as fp:
for line in fp:
try:
record = json.loads(line)
except json.JSONDecodeError:
continue
if record.get("mode") != "RL":
if was_rl:
history.clear()
was_rl = False
continue
was_rl = True
obs = np.asarray(record.get("obs_single"), dtype=np.float32)
if obs.shape != (NUM_OBS,) or not np.all(np.isfinite(obs)):
continue
history.append(obs)
onnx_input = build_onnx_input(history)
if not np.all(np.isfinite(onnx_input)):
continue
seen += 1
run_seen += 1
reservoir_add(samples, onnx_input, seen, max_samples, rng)
if run_seen:
usable_runs.append({"steps": str(steps_path), "samples": run_seen})
return samples, seen, usable_runs
def main():
parser = argparse.ArgumentParser(
description="Create BPU float32 calibration .bin files from RobotLab JSONL logs."
)
parser.add_argument(
"--logs-root",
type=Path,
default=DEFAULT_LOG_ROOT,
help="Directory containing robotlab_go1_deploy_*/steps.jsonl.",
)
parser.add_argument(
"--output-dir",
type=Path,
default=HERE / "calibration_data",
help="Output directory for raw float32 feature-map .bin files.",
)
parser.add_argument("--max-samples", type=int, default=512)
parser.add_argument("--seed", type=int, default=20260727)
parser.add_argument("--overwrite", action="store_true")
args = parser.parse_args()
if args.max_samples < 32:
raise ValueError("--max-samples must be at least 32")
log_paths = sorted(args.logs_root.glob("robotlab_go1_deploy_*/steps.jsonl"))
if not log_paths:
raise FileNotFoundError(f"No RobotLab step logs found below {args.logs_root}")
samples, total_seen, usable_runs = collect_samples(
log_paths, args.max_samples, args.seed
)
if len(samples) < 32:
raise RuntimeError(
f"Only {len(samples)} valid RL inputs found; need at least 32 calibration samples."
)
output_dir = args.output_dir.resolve()
existing = list(output_dir.glob("*.bin")) if output_dir.exists() else []
if existing and not args.overwrite:
raise FileExistsError(
f"{output_dir} already contains calibration files; pass --overwrite to replace them."
)
output_dir.mkdir(parents=True, exist_ok=True)
for path in existing:
path.unlink()
for index, sample in enumerate(samples):
sample.astype(np.float32, copy=False).tofile(output_dir / f"{index:05d}.bin")
metadata = {
"format": "raw float32 feature-map",
"flat_shape": [1, ONNX_INPUT_DIM],
"mapper_shape": [1, 1, 1, ONNX_INPUT_DIM],
"history_len": HISTORY_LEN,
"num_obs": NUM_OBS,
"term_dims": list(TERM_DIMS),
"selected_samples": len(samples),
"candidate_rl_inputs": total_seen,
"seed": args.seed,
"runs": usable_runs,
}
(output_dir.parent / f"{output_dir.name}_metadata.json").write_text(
json.dumps(metadata, indent=2, ensure_ascii=True) + "\n",
encoding="utf-8",
)
print(
f"Wrote {len(samples)} calibration inputs from {total_seen} RL observations "
f"to {output_dir}"
)
if __name__ == "__main__":
main()

27
deploy_45dim_rl_gym/bpu_quantization/policy_robotlab_15000_int16.yaml Normal file
View File

@@ -0,0 +1,27 @@
model_parameters:
onnx_model: "./policy_robotlab_15000_bpu4d.onnx"
march: "bayes-e"
layer_out_dump: false
working_dir: "mapper_output"
output_model_file_prefix: "policy_robotlab_15000_int16"
input_parameters:
input_name: "obs_4d"
input_shape: "1x1x1x450"
input_type_rt: "featuremap"
input_layout_rt: "NCHW"
input_type_train: "featuremap"
input_layout_train: "NCHW"
norm_type: "no_preprocess"
calibration_parameters:
cal_data_dir: "./calibration_data"
cal_data_type: "float32"
calibration_type: "default"
optimization: "set_all_nodes_int16"
per_channel: true
compiler_parameters:
compile_mode: "latency"
debug: false
optimize_level: "O3"

27
deploy_45dim_rl_gym/bpu_quantization/policy_robotlab_15000_int16_gemm.yaml Normal file
View File

@@ -0,0 +1,27 @@
model_parameters:
onnx_model: "./policy_robotlab_15000_bpu4d_gemm.onnx"
march: "bayes-e"
layer_out_dump: false
working_dir: "mapper_output_gemm"
output_model_file_prefix: "policy_robotlab_15000_int16_gemm"
input_parameters:
input_name: "obs_4d"
input_shape: "1x1x1x450"
input_type_rt: "featuremap"
input_layout_rt: "NCHW"
input_type_train: "featuremap"
input_layout_train: "NCHW"
norm_type: "no_preprocess"
calibration_parameters:
cal_data_dir: "./calibration_data"
cal_data_type: "float32"
calibration_type: "default"
optimization: "set_all_nodes_int16"
per_channel: true
compiler_parameters:
compile_mode: "latency"
debug: false
optimize_level: "O3"

27
deploy_45dim_rl_gym/bpu_quantization/policy_robotlab_6500_int16_gemm.yaml Normal file
View File

@@ -0,0 +1,27 @@
model_parameters:
onnx_model: "./policy_robotlab_6500_bpu4d_gemm.onnx"
march: "bayes-e"
layer_out_dump: false
working_dir: "mapper_output_6500_gemm"
output_model_file_prefix: "policy_robotlab_6500_int16_gemm"
input_parameters:
input_name: "obs_4d"
input_shape: "1x1x1x450"
input_type_rt: "featuremap"
input_layout_rt: "NCHW"
input_type_train: "featuremap"
input_layout_train: "NCHW"
norm_type: "no_preprocess"
calibration_parameters:
cal_data_dir: "./calibration_data_fast64"
cal_data_type: "float32"
calibration_type: "default"
optimization: "set_all_nodes_int16"
per_channel: true
compiler_parameters:
compile_mode: "latency"
debug: false
optimize_level: "O3"

149
deploy_45dim_rl_gym/bpu_quantization/replace_group_conv_with_gemm.py Normal file
View File

@@ -0,0 +1,149 @@
#!/usr/bin/env python3
"""Replace the policy's grouped 1x1 Conv expert layer with an equivalent Gemm.
The RobotLab MoE expert layer is exported as:
[1, 2048] -> Unsqueeze -> grouped Conv1d(group=8, kernel=1)
-> Squeeze -> [1, 256]
For this fixed shape, the grouped Conv is exactly a block-diagonal linear layer.
Using Gemm lets hb_mapper keep the expert path on BPU instead of cutting through
CPU float execution.
"""
import argparse
from pathlib import Path
import numpy as np
import onnx
from onnx import helper, numpy_helper
CONV_NAME = "/student_moe_encoder/moe/experts/experts/Conv"
def attr_value(node, name, default=None):
for attr in node.attribute:
if attr.name == name:
return helper.get_attribute_value(attr)
return default
def initializer_map(model):
return {initializer.name: initializer for initializer in model.graph.initializer}
def producer_map(model):
producers = {}
for node in model.graph.node:
for output in node.output:
producers[output] = node
return producers
def replace_node(model):
initializers = initializer_map(model)
producers = producer_map(model)
nodes = list(model.graph.node)
conv = next((node for node in nodes if node.name == CONV_NAME), None)
if conv is None:
raise ValueError(f"Cannot find node {CONV_NAME!r}")
if conv.op_type != "Conv":
raise ValueError(f"{CONV_NAME!r} is {conv.op_type}, expected Conv")
if attr_value(conv, "kernel_shape") != [1]:
raise ValueError("Only kernel_shape=[1] is supported")
if attr_value(conv, "strides") not in (None, [1]):
raise ValueError("Only stride=1 is supported")
if attr_value(conv, "pads") not in (None, [0, 0]):
raise ValueError("Only no padding is supported")
groups = int(attr_value(conv, "group", 1))
weight = numpy_helper.to_array(initializers[conv.input[1]])
bias = numpy_helper.to_array(initializers[conv.input[2]]) if len(conv.input) >= 3 else None
if weight.ndim != 3 or weight.shape[2] != 1:
raise ValueError(f"Expected Conv weight [out, in_per_group, 1], got {weight.shape}")
out_channels, in_per_group, _ = weight.shape
if out_channels % groups != 0:
raise ValueError(f"out_channels={out_channels} is not divisible by groups={groups}")
out_per_group = out_channels // groups
in_channels = in_per_group * groups
full_weight = np.zeros((out_channels, in_channels), dtype=np.float32)
flat_weight = weight[:, :, 0].astype(np.float32, copy=False)
for group in range(groups):
out_begin = group * out_per_group
out_end = out_begin + out_per_group
in_begin = group * in_per_group
in_end = in_begin + in_per_group
full_weight[out_begin:out_end, in_begin:in_end] = flat_weight[out_begin:out_end]
unsqueeze = producers.get(conv.input[0])
if unsqueeze is None or unsqueeze.op_type != "Unsqueeze":
raise ValueError("Expected Conv input to be produced by Unsqueeze")
squeeze_users = [node for node in nodes if conv.output[0] in node.input]
if len(squeeze_users) != 1 or squeeze_users[0].op_type != "Squeeze":
raise ValueError("Expected Conv output to feed exactly one Squeeze")
squeeze = squeeze_users[0]
gemm_weight_name = conv.input[1] + "_as_gemm"
model.graph.initializer.append(
numpy_helper.from_array(full_weight, name=gemm_weight_name)
)
gemm_inputs = [unsqueeze.input[0], gemm_weight_name]
if bias is not None:
gemm_inputs.append(conv.input[2])
gemm = helper.make_node(
"Gemm",
inputs=gemm_inputs,
outputs=[squeeze.output[0]],
name=CONV_NAME + "_as_Gemm",
alpha=1.0,
beta=1.0,
transB=1,
)
conv_index = nodes.index(conv)
remove_names = {unsqueeze.name, conv.name, squeeze.name}
new_nodes = []
inserted = False
for index, node in enumerate(nodes):
if node.name in remove_names:
if index >= conv_index and not inserted:
new_nodes.append(gemm)
inserted = True
continue
new_nodes.append(node)
del model.graph.node[:]
model.graph.node.extend(new_nodes)
return {
"groups": groups,
"conv_weight_shape": list(weight.shape),
"gemm_weight_shape": list(full_weight.shape),
"removed": sorted(remove_names),
"added": gemm.name,
}
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
model = onnx.load(str(args.input))
info = replace_node(model)
args.output.parent.mkdir(parents=True, exist_ok=True)
onnx.save(model, str(args.output))
print(f"Wrote {args.output}")
for key, value in info.items():
print(f"{key}={value}")
if __name__ == "__main__":
main()