From 0d16832612dce8165b96007de9fac787b8da7bba Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Thu, 13 Aug 2026 01:31:58 +0000 Subject: [PATCH] =?UTF-8?q?feat(actpred):=20=E6=96=B0=E5=A2=9E=20Hailo=20D?= =?UTF-8?q?FC=20=E7=BC=96=E8=AF=91=E9=93=BE=E8=B7=AF=EF=BC=8CONNX=20?= =?UTF-8?q?=E2=86=92=20HEF=20=E5=85=A8=E6=B5=81=E7=A8=8B=E6=89=93=E9=80=9A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 训练态 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）。 --- actpred/.gitignore | 8 ++ actpred/README.md | 41 +++++---- actpred/dfc_compile.py | 115 ++++++++++++++++++++++++ actpred/export_deploy.py | 178 +++++++++++++++++++++++++++++++++++++ actpred/gen_calibration.py | 130 +++++++++++++++++++++++++++ 5 files changed, 451 insertions(+), 21 deletions(-) create mode 100644 actpred/dfc_compile.py create mode 100644 actpred/export_deploy.py create mode 100644 actpred/gen_calibration.py diff --git a/actpred/.gitignore b/actpred/.gitignore index b74f2a3..950586d 100644 --- a/actpred/.gitignore +++ b/actpred/.gitignore @@ -1,5 +1,6 @@ # Virtual environment venv/ +venv_dfc/ .env/ # Python bytecode @@ -12,6 +13,13 @@ checkpoints/ *.pt *.pth *.onnx +*.onnx.data +*.hef +*.har + +# DFC logs +acceleras.log +hailo_sdk.*.log # Dataset (too large for git) data/ diff --git a/actpred/README.md b/actpred/README.md index 441aa32..d513435 100644 --- a/actpred/README.md +++ b/actpred/README.md @@ -121,7 +121,7 @@ recorder.flush() | `--attn-layers` | 2 | TransformerEncoder 层数 | | `--epochs` | 100 | 训练轮数 | | `--lr` | 0.001 | 学习率 | -| `--batch-size` | 32 | 批次大小 | +| `--batch-size` | 4 | 批次大小（默认分辨率 1920×1080 下 CPU/Pi 内存受限；小分辨率可调大） | | `--normalize-actions` | False | 动作归一化（`--no-normalize-actions` 关闭） | | `--dropout` | 0.1 | Dropout 率 | | `--weight-decay` | 0.0001 | 权重衰减 | @@ -180,30 +180,29 @@ python3 onnx_infer.py --model model.onnx --images f1.jpg f2.jpg f3.jpg \ ## Hailo 部署（AI HAT 2 / Hailo-10H） -树莓派 AI HAT 2 上通过 `hailo_infer.py` 的 `HailoPredictor` 走 NPU 推理， -接口与 `ActionPredictorInference` 对齐。模型需先经 Hailo DFC 编译为 HEF： +模型需先经 Hailo DFC（仅 x86 Linux，`venv_dfc`）编译为 HEF。 +训练态 ONNX（5D 输入、AdaptiveAvgPool、MultiheadAttention 的 head +reshape）**无法被 DFC 解析**，必须用 `export_deploy.py` 的部署版导出 +（按帧展开 CNN、固定 kernel AvgPool、注意力按头展开为独立 Linear、 +全程 (1,1,S,d) 无 transpose 布局，权重数值不变）： ```bash -# 1. 导出 ONNX -bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx -# 2. 在 x86 主机（或 Hailo 云端 DFC）上编译 + 量化校准 → model.hef，拷回树莓派 -# 3. 推理（零 torch 依赖，仅需 hailo_platform + numpy + PIL） +# x86 主机（venv_dfc 内已装 hailo-sdk 5.3）： +# 1. 导出部署版 ONNX（输入 frame0~2 + action_history + pred_token） +python3 export_deploy.py --model checkpoints/pretrain_best_inference.pt +# 2. 生成校准集（合成数据；真实 episode 到位后应替换） +python3 gen_calibration.py --deploy +# 3. 一键编译（parser → optimize → compiler） +python3 dfc_compile.py # 产物 checkpoints/actpred_pretrain.hef +# 4. 将 HEF 拷回树莓派 ``` -```python -from hailo_infer import HailoPredictor +已验证：完整 110 层网络（含 matmul/softmax/layernorm 注意力）解析、 +量化、编译全部通过；HEF 1.8MB，4 contexts，hailo10h 目标。 +编译要点与踩坑记录见 `dfc_compile.py` 文档字符串。 -with HailoPredictor("model.hef") as predictor: # stats_path= 仅归一化模型需要 - action = predictor.predict(images, action_history) -``` - -```bash -python3 hailo_infer.py --hef model.hef --images f1.jpg f2.jpg f3.jpg \ - --action-history "45,0,...:45,0,..." [--stats action_stats.npy] -``` - -HEF 的 vstream 名称/形状/量化参数全部从 `HEF.get_input/output_vstream_infos()` -动态推导，不硬编码；图像输入按 Hailo 惯例为 uint8 原始像素 -（ImageNet 归一化由 DFC 编译时折叠进网络）。 +树莓派端 `hailo_infer.py` 的 `HailoPredictor` 尚需适配新的 5 输入 +HEF（frame0~2 各 (1,64,64,3) NHWC、action_history (1,1,5,18)、 +pred_token (1,1,1,128) 常量），输出 predicted_action (1,1,18)。 环境调研与 API 踩坑记录见 [docs/reports/actpred_hailo_hat2_inference_research_20260813.md](../docs/reports/actpred_hailo_hat2_inference_research_20260813.md)。 diff --git a/actpred/dfc_compile.py b/actpred/dfc_compile.py new file mode 100644 index 0000000..1f44957 --- /dev/null +++ b/actpred/dfc_compile.py @@ -0,0 +1,115 @@ +"""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() diff --git a/actpred/export_deploy.py b/actpred/export_deploy.py new file mode 100644 index 0000000..c44f3f1 --- /dev/null +++ b/actpred/export_deploy.py @@ -0,0 +1,178 @@ +"""导出 Hailo 部署版 ONNX：DFC 友好图结构 + +训练态 ONNX 对 DFC 不友好：5D 图像输入、AdaptiveAvgPool2d（被 +torch.export 展开成 Loop 节点）、MultiheadAttention 的 head 拆分 +Reshape 序列。此脚本逐项改造（权重数值完全不变）： + 1. 图像输入拆为 F 个独立 (1,H,W,3) NHWC 输入，按帧展开 CNN 分支， + 各分支持有独立权重副本（共享 initializer 会被 DFC 误融合） + 2. AdaptiveAvgPool2d 换成等价的固定 kernel AvgPool2d + 3. 自注意力手工展开为每头独立 Linear（128→32），无 head reshape； + pre-LN 残差结构与训练态 nn.TransformerEncoderLayer(norm_first=True) + 逐位等价 +""" + +import argparse + +import torch +import torch.nn as nn + +from config import parse_checkpoint_config +from model import ActionPredictor + + +class ManualAttnBlock(nn.Module): + """等价替换 nn.TransformerEncoderLayer(d=128, nhead=4, ff=256, norm_first=True)""" + + def __init__(self, layer: nn.TransformerEncoderLayer): + super().__init__() + d = layer.self_attn.embed_dim + nh = layer.self_attn.num_heads + self.head_dim = d // nh + self.nh = nh + self.scale = self.head_dim ** -0.5 + + ipw = layer.self_attn.in_proj_weight # (3d, d) = [Wq; Wk; Wv] + ipb = layer.self_attn.in_proj_bias + self.q_proj = nn.ModuleList() + self.k_proj = nn.ModuleList() + self.v_proj = nn.ModuleList() + for h in range(nh): + sl = slice(h * self.head_dim, (h + 1) * self.head_dim) + slices = (sl, slice(d + sl.start, d + sl.stop), + slice(2 * d + sl.start, 2 * d + sl.stop)) + for proj, ws in ((self.q_proj, slices[0]), (self.k_proj, slices[1]), + (self.v_proj, slices[2])): + lin = nn.Linear(d, self.head_dim) + lin.weight.data.copy_(ipw[ws]) + proj.append(lin) + for h in range(nh): + self.q_proj[h].bias.data.copy_(ipb[h * self.head_dim:(h + 1) * self.head_dim]) + self.k_proj[h].bias.data.copy_(ipb[d + h * self.head_dim:d + (h + 1) * self.head_dim]) + self.v_proj[h].bias.data.copy_(ipb[2 * d + h * self.head_dim:2 * d + (h + 1) * self.head_dim]) + + self.out_proj = layer.self_attn.out_proj + self.norm1 = layer.norm1 + self.norm2 = layer.norm2 + self.fc1 = layer.linear1 + self.fc2 = layer.linear2 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (1, 1, S, d)，seq 占 W 维，pre-LN 自注意力 + 残差 + n = self.norm1(x) + heads = [] + for h in range(self.nh): + q = self.q_proj[h](n) + k = self.k_proj[h](n) + v = self.v_proj[h](n) + s = torch.matmul(q, k.transpose(-1, -2)) * self.scale + a = torch.softmax(s, dim=-1) + heads.append(torch.matmul(a, v)) + attn = self.out_proj(torch.cat(heads, dim=3)) # feature 维 concat + x = x + attn + x = x + self.fc2(torch.relu(self.fc1(self.norm2(x)))) + return x + + +class DeployWrapper(nn.Module): + def __init__(self, model: ActionPredictor): + super().__init__() + import copy + m = model + ref_cnn = nn.Sequential(*list(m.image_encoder.net.children())[:-1]) + self.cnns = nn.ModuleList( + [copy.deepcopy(ref_cnn) for _ in range(m._frame_len)]) + # img_proj ∘ AvgPool == AvgPool ∘ Conv1x1（均为线性算子，可交换）。 + # 把 img_proj 的权重折进每分支末尾的 1×1 卷积，偏置单独用加法， + # 避免 DFC 对 3D Linear 的错误卷积融合 + self.proj_convs = nn.ModuleList() + for _ in range(m._frame_len): + c = nn.Conv2d(m.img_proj.in_features, m.img_proj.out_features, 1, bias=False) + c.weight.data.copy_(m.img_proj.weight.data.view( + m.img_proj.out_features, m.img_proj.in_features, 1, 1)) + self.proj_convs.append(c) + self.proj_bias = nn.Parameter(m.img_proj.bias.data.clone()) + self.pool = nn.AvgPool2d((m._image_height // 16, m._image_width // 16)) + # act_proj 权重注册为常量，用显式 MatMul+Add 代替 Linear， + # 避免 DFC 对 3D Linear 的错误卷积融合 + self.act_w = nn.Parameter(m.act_proj.weight.data.t().clone()) # (D, d) + self.act_b = nn.Parameter(m.act_proj.bias.data.clone()) + # pred_token 改为图输入（常量 initializer 会被 DFC 误转置）， + # 位置编码经 concat 后的 ew_add 广播可被正确推断 + self.pos_embed4 = nn.Parameter(m.pos_embed.data.unsqueeze(1).clone()) # (1, 1, S, d) + self.blocks = nn.ModuleList( + [ManualAttnBlock(l) for l in model.encoder.layers]) + self.head = m.head + self._frame_len = m._frame_len + + def _frame_feat(self, frame_nhwc: torch.Tensor, i: int) -> torch.Tensor: + x = self.cnns[i](frame_nhwc.permute(0, 3, 1, 2)) + x = self.proj_convs[i](x) # 1x1 conv ≡ 池化后的 img_proj + return self.pool(x) # (1, d, 1, 1) NCHW + + def forward(self, frame0: torch.Tensor, frame1: torch.Tensor, frame2: torch.Tensor, + action_history: torch.Tensor, pred_token: torch.Tensor) -> torch.Tensor: + # 全程 (1, 1, S, d) 布局：seq 占 W 维，避免任何 transpose + frames = [frame0, frame1, frame2] + feats = [self._frame_feat(frames[i], i).permute(0, 2, 3, 1) + for i in range(self._frame_len)] # (1, 1, 1, d) NHWC + img_tokens = torch.cat(feats, dim=2) + self.proj_bias # (1, 1, F, d) + act_tokens = torch.matmul(action_history, self.act_w) + self.act_b # (1, 1, H, d) + tokens = torch.cat([img_tokens, act_tokens, pred_token], dim=2) # (1, 1, S, d) + tokens = tokens + self.pos_embed4 + for blk in self.blocks: + tokens = blk(tokens) + return self.head(tokens[:, :, -1:, :]).squeeze(1) # (1, 1, action_dim) + + +def main(): + parser = argparse.ArgumentParser(description="Export Hailo-friendly ONNX") + parser.add_argument("--model", required=True, help="checkpoint path") + parser.add_argument("--output", default="checkpoints/actpred_deploy.onnx") + args = parser.parse_args() + + checkpoint = torch.load(args.model, map_location="cpu", weights_only=True) + cfg = parse_checkpoint_config(checkpoint["config"]) + + model = ActionPredictor(cfg.model, cfg.data, pretrained=False) + model.load_state_dict(checkpoint["model_state_dict"]) + model.eval() + wrapper = DeployWrapper(model).eval() + + h, w = cfg.data.image_height, cfg.data.image_width + f = cfg.data.frame_len + dummy_frames = [torch.randn(1, h, w, 3) for _ in range(f)] + dummy_history = torch.randn(1, 1, cfg.data.history_len, cfg.data.action_dim) + dummy_pred = model.pred_token.data.unsqueeze(1).clone() # (1, 1, 1, d) + + with torch.no_grad(): + ref = wrapper(*dummy_frames, dummy_history, dummy_pred) + + torch.onnx.export( + wrapper, (*dummy_frames, dummy_history, dummy_pred), + args.output, + input_names=[f"frame{i}" for i in range(f)] + ["action_history", "pred_token"], + output_names=["predicted_action"], + opset_version=18, + ) + print(f"部署版 ONNX 已导出: {args.output}（{f} 个帧输入 (1,{h},{w},3) NHWC + action_history + pred_token）") + + stacked = torch.stack([fr.permute(0, 3, 1, 2) for fr in dummy_frames], dim=1) + with torch.no_grad(): + orig_out = model(stacked, dummy_history.squeeze(1)) + print(f"wrapper 与原模型最大差异: {(ref.squeeze(1) - orig_out).abs().max():.6e}") + + try: + import onnxruntime as ort + import numpy as np + sess = ort.InferenceSession(args.output) + feeds = {f"frame{i}": dummy_frames[i].numpy() for i in range(f)} + feeds["action_history"] = dummy_history.numpy() + feeds["pred_token"] = dummy_pred.numpy() + ort_out = sess.run(None, feeds)[0] + print(f"ONNX 一致性验证通过（最大差异: {np.abs(ort_out - ref.numpy()).max():.6e}）") + except ImportError: + pass + + +if __name__ == "__main__": + main() diff --git a/actpred/gen_calibration.py b/actpred/gen_calibration.py new file mode 100644 index 0000000..8eba17f --- /dev/null +++ b/actpred/gen_calibration.py @@ -0,0 +1,130 @@ +"""生成 DFC 量化校准数据集 + +两种模式： + 默认模式（data/calibration.npz）：为旧版 dfc_compile.py 的 ClientModel + 流程准备，目前保留兼容。 + --deploy 模式（data/calibration_deploy.npz）：为 export_deploy.py 导出 + 的部署版 ONNX 生成校准集，键名匹配 HN 层名 actpred_deploy/input_layerN， + 形状匹配 HEF 输入：frame0~2 (N,H,W,3) 已归一化 float32、 + action_history (N,1,H_len,D)、pred_token (N,1,1,d) 常量。 + 图像分辨率与 --model checkpoint 的配置一致。 + +数据源： + - 图像：无真实 episode 时用 [0,255] 均匀采样像素 + ImageNet 归一化 + - 动作历史：从 bot4cj/scripts/*.mt 的真实 @Action 矩阵滑窗采样 + （列不足 18 补零，与 pretrain_mt.py 语义一致） + +真实数据到位后应替换为真实帧重新校准。 + +用法: + python gen_calibration.py --deploy --model checkpoints/pretrain_best_inference.pt +""" + +import argparse + +import numpy as np + +from config import DataConfig, parse_checkpoint_config +from mt_parser import parse_mt_directory + +IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) +IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) + + +def build_action_history_samples(mt_dir, n_samples, history_len, action_dim, rng): + sequences = [] + for _path, matrix in parse_mt_directory(mt_dir): + n_cols = matrix.shape[1] + if n_cols < action_dim: + matrix = np.pad(matrix, ((0, 0), (0, action_dim - n_cols))) + elif n_cols > action_dim: + matrix = matrix[:, :action_dim] + if len(matrix) >= history_len: + sequences.append(matrix) + if not sequences: + raise ValueError(f"{mt_dir} 中没有可用的 @Action 序列") + + windows = [] + for matrix in sequences: + for i in range(len(matrix) - history_len + 1): + windows.append(matrix[i:i + history_len]) + idx = rng.integers(0, len(windows), size=n_samples) + return np.stack([windows[i] for i in idx]).astype(np.float32) + + +def build_image_samples(n_samples, frame_len, height, width, rng): + # 存 uint8 原始像素值（0~255），加载时再做 ToTensor+Normalize，节省磁盘 + return rng.integers(0, 256, + size=(n_samples, frame_len, 3, height, width), + dtype=np.uint8) + + +def build_deploy_set(args): + import torch + from model import ActionPredictor + + ckpt = torch.load(args.model, map_location="cpu", weights_only=True) + cfg = parse_checkpoint_config(ckpt["config"]).data + rng = np.random.default_rng(args.seed) + + h, w = cfg.image_height, cfg.image_width + raw = rng.integers(0, 256, size=(args.samples, h, w, 3), + dtype=np.uint8).astype(np.float32) / 255.0 + norm = ((raw - IMAGENET_MEAN) / IMAGENET_STD).astype(np.float32) + frames = [norm.copy() for _ in range(cfg.frame_len)] + + history = build_action_history_samples( + args.mt_dir, args.samples, cfg.history_len, cfg.action_dim, rng) + history = history[:, None, :, :] # (N, 1, H_len, D) + + model = ActionPredictor(parse_checkpoint_config(ckpt["config"]).model, + cfg, pretrained=False) + model.load_state_dict(ckpt["model_state_dict"]) + d = model.pred_token.shape[-1] + pred = np.broadcast_to( + model.pred_token.data.unsqueeze(1).numpy(), + (args.samples, 1, 1, d)).astype(np.float32).copy() + + np.savez_compressed( + args.output, + **{f"actpred_deploy/input_layer{i + 1}": arr + for i, arr in enumerate([*frames, history, pred])}) + print(f"部署校准集已保存: {args.output}") + print(f" frame0~2: {frames[0].shape}（分辨率 {w}x{h}，与 checkpoint 一致）") + print(f" action_history: {history.shape} pred_token: {pred.shape}") + + +def main(): + parser = argparse.ArgumentParser(description="Generate DFC calibration dataset") + parser.add_argument("--deploy", action="store_true", + help="生成部署版（export_deploy.py）校准集") + parser.add_argument("--model", default="checkpoints/pretrain_best_inference.pt", + help="--deploy 模式的 checkpoint（决定分辨率与 pred_token）") + parser.add_argument("--mt-dir", default="../bot4cj/scripts") + parser.add_argument("--output", default=None) + parser.add_argument("--samples", type=int, default=32) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + if args.deploy: + args.output = args.output or "data/calibration_deploy.npz" + build_deploy_set(args) + return + + args.output = args.output or "data/calibration.npz" + cfg = DataConfig() + rng = np.random.default_rng(args.seed) + + history = build_action_history_samples( + args.mt_dir, args.samples, cfg.history_len, cfg.action_dim, rng) + images = build_image_samples( + args.samples, cfg.frame_len, cfg.image_height, cfg.image_width, rng) + + np.savez_compressed(args.output, images=images, action_history=history) + print(f"校准集已保存: {args.output}") + print(f" images: {images.shape} dtype={images.dtype}") + print(f" action_history: {history.shape} dtype={history.dtype}") + + +if __name__ == "__main__": + main()