feat(actpred): 新增 Hailo DFC 编译链路,ONNX → HEF 全流程打通
训练态 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)。
This commit is contained in:
1 parent
0cfb3a7581
commit
0d16832612
5 files changed
+451
-21
No files matched your search
@@ -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()
|
||||
Reference in new issue
Block a user