qing
aa271e395a
P1-1: export_deploy.py forward() 从硬编码 3 帧改为 *inputs 动态解析
P1-2: hailo_infer.py 补 predict_delta 支持（残差模式 + --ckpt 自动检测）
P1-3: pybot _load_predictor 自动查找 pred_token.npy/action_stats.npy/
companion checkpoint，恢复 Hailo 推理训练语义
P1-4: pretrain_mt.py random_images 每 batch 重新生成噪声，防 BN 退化
P2-1: predict.py/hailo_infer.py/robot.py 三层 NaN/Inf 防线
P2-3: recorder.py imu_dim 默认 8→0，避免外部直接使用时形状错误
P2-5: _npu_test.py imu_dim 移到函数定义前，消除执行顺序依赖
296 lines
14 KiB
Python
296 lines
14 KiB
Python
"""导出 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 os
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
from config import parse_checkpoint_config
|
||
from model import ActionPredictor
|
||
|
||
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)
|
||
|
||
MAX_TOKEN_ABS = 100.0
|
||
MERGE_RATIO_WARN = 32.0
|
||
|
||
|
||
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
|
||
if not hasattr(m.image_encoder, "net"):
|
||
raise ValueError(
|
||
"部署版导出仅支持 backbone=simple（MobileNet 的 DFC 改造未实现）")
|
||
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())
|
||
# 固定 kernel AvgPool 须与训练态 AdaptiveAvgPool2d(1) 逐值等价：
|
||
# kernel 取 ref_cnn 的真实输出尺寸。H//16 只在分辨率能被 16 整除
|
||
# 时成立（1080 高度下卷积输出 68 行，按 67 取 kernel 会静默丢弃
|
||
# 末行，破坏与训练图的数值一致）
|
||
with torch.no_grad():
|
||
probe = ref_cnn(torch.zeros(
|
||
1, 3, m._image_height, m._image_width))
|
||
self.pool = nn.AvgPool2d((probe.shape[2], probe.shape[3]))
|
||
# 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())
|
||
# R2 方案一：IMU 独立 token（imu_dim>0 时追加图输入，同 MatMul+Add 手法）
|
||
self._imu_dim = m._imu_dim
|
||
if self._imu_dim > 0:
|
||
self.imu_w = nn.Parameter(m.imu_proj.weight.data.t().clone())
|
||
self.imu_b = nn.Parameter(m.imu_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, *inputs: torch.Tensor) -> torch.Tensor:
|
||
# inputs = frame0..frame{F-1}, action_history, pred_token, [imu]
|
||
# 全程 (1, 1, S, d) 布局：seq 占 W 维，避免任何 transpose
|
||
frames = list(inputs[:self._frame_len])
|
||
action_history = inputs[self._frame_len]
|
||
pred_token = inputs[self._frame_len + 1]
|
||
imu = inputs[self._frame_len + 2] if self._imu_dim > 0 else None
|
||
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)
|
||
parts = [img_tokens, act_tokens]
|
||
if self._imu_dim > 0:
|
||
# imu: (1, 1, 1, imu_dim) 归一化姿态 → (1, 1, 1, d)，位于 PRED 之前
|
||
parts.append(torch.matmul(imu, self.imu_w) + self.imu_b)
|
||
parts.append(pred_token)
|
||
tokens = torch.cat(parts, 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 check_token_ranges(wrapper: "DeployWrapper", frame_len: int, height: int,
|
||
width: int, history_len: int, action_dim: int,
|
||
pred_token: torch.Tensor, max_abs: float = MAX_TOKEN_ABS,
|
||
action_stats_path: str = None):
|
||
"""导出前检查 token 量级：幅度过大时 uint8 量化会把小信号压成 0
|
||
|
||
预训练用全零图像占位 → BN running_var 停在初始小值 → 推理时每级
|
||
BN 放大 ~50 倍，图像 token 可达 ±2e5。DFC 在 token merge 点按最大
|
||
幅度分配共享量化 scale（实测 scale=1965.92，pred_token/act_tokens
|
||
全部量化成 0），HEF 输出对任何输入恒定。必须在编译前拦截。
|
||
|
||
归一化模型（normalize_actions）的 HEF 输入是 z-score 后的历史，
|
||
检查输入须与推理端一致。
|
||
"""
|
||
rng = np.random.default_rng(42)
|
||
raw = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8)
|
||
norm = (raw.astype(np.float32) / 255.0 - IMAGENET_MEAN) / IMAGENET_STD
|
||
frames = [torch.tensor(norm[None]) for _ in range(frame_len)]
|
||
history = rng.uniform(0.0, 500.0, size=(1, 1, history_len, action_dim))
|
||
if action_stats_path is not None:
|
||
stats = np.load(action_stats_path)
|
||
if stats.shape != (2, action_dim):
|
||
raise ValueError(
|
||
f"action stats 形状应为 (2, {action_dim})，"
|
||
f"实际为 {stats.shape}: {action_stats_path}")
|
||
# 归一化模型直接在 z 空间采样：raw 均匀采样会在零方差列上
|
||
# 产生天文数字 z 值（实际推理中这些列恒为 0，z 恒为 0）
|
||
history = rng.normal(0.0, 1.5, size=(1, 1, history_len, action_dim))
|
||
history = torch.tensor(history, dtype=torch.float32)
|
||
with torch.no_grad():
|
||
feats = [wrapper._frame_feat(f, i).permute(0, 2, 3, 1)
|
||
for i, f in enumerate(frames)]
|
||
img_tokens = torch.cat(feats, dim=2) + wrapper.proj_bias
|
||
act_tokens = torch.matmul(history, wrapper.act_w) + wrapper.act_b
|
||
img_max = img_tokens.abs().max().item()
|
||
act_max = act_tokens.abs().max().item()
|
||
pred_max = pred_token.abs().max().item()
|
||
print(f"token 量级检查: |img|={img_max:.4g} |act|={act_max:.4g} "
|
||
f"|pred|={pred_max:.4g}（图像上限 {max_abs}）")
|
||
if img_max > max_abs:
|
||
raise SystemExit(
|
||
f"错误: 图像 token 幅度 {img_max:.3g} 超过可量化上限 {max_abs}。"
|
||
f"checkpoint 的 BN 统计退化（预训练用全零图像占位，running_var "
|
||
f"停在初始小值），uint8 量化会把小信号全部压成 0、HEF 输出恒定。"
|
||
f"请先用真实/随机图像数据微调后再导出。")
|
||
if act_max > MERGE_RATIO_WARN * max(img_max, 1.0):
|
||
print(f"WARNING: act_tokens 幅度 {act_max:.3g} 远大于 img_tokens "
|
||
f"{img_max:.3g}，merge 点共享 scale 会粗量化图像路径；"
|
||
f"建议用 --normalize-actions 训练（动作归一化）压小 act_tokens")
|
||
|
||
|
||
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")
|
||
parser.add_argument("--pred-token-out", default=None,
|
||
help="保存学习到的 pred_token 到 .npy（HEF 推理端喂入用，"
|
||
"配合抖动校准避免常量量化退化）")
|
||
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
|
||
stats_path = None
|
||
if cfg.data.normalize_actions:
|
||
# 与 predict.py 语义一致：相对路径按 checkpoint 目录解析
|
||
stats_path = cfg.data.action_stats_path
|
||
if not stats_path:
|
||
stats_path = os.path.join(
|
||
os.path.dirname(os.path.abspath(args.model)),
|
||
"action_stats.npy")
|
||
elif not os.path.isabs(stats_path):
|
||
stats_path = os.path.join(
|
||
os.path.dirname(os.path.abspath(args.model)), stats_path)
|
||
if not os.path.exists(stats_path):
|
||
raise FileNotFoundError(
|
||
f"模型训练时启用了动作归一化，但未找到统计文件: {stats_path}")
|
||
check_token_ranges(wrapper, f, h, w, cfg.data.history_len,
|
||
cfg.data.action_dim, model.pred_token.data.unsqueeze(1),
|
||
action_stats_path=stats_path)
|
||
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)
|
||
imu_dim = cfg.data.imu_dim
|
||
# imu_dim>0 时追加归一化姿态输入（均匀采样的归一化值域）
|
||
dummy_imu = torch.randn(1, 1, 1, imu_dim) if imu_dim > 0 else None
|
||
|
||
extra_inputs = (dummy_imu,) if imu_dim > 0 else ()
|
||
extra_names = ["imu"] if imu_dim > 0 else []
|
||
|
||
with torch.no_grad():
|
||
ref = wrapper(*dummy_frames, dummy_history, dummy_pred, *extra_inputs)
|
||
|
||
torch.onnx.export(
|
||
wrapper, (*dummy_frames, dummy_history, dummy_pred, *extra_inputs),
|
||
args.output,
|
||
input_names=[f"frame{i}" for i in range(f)] + ["action_history", "pred_token"] + extra_names,
|
||
output_names=["predicted_action"],
|
||
opset_version=18,
|
||
)
|
||
print(f"部署版 ONNX 已导出: {args.output}（{f} 个帧输入 (1,{h},{w},3) NHWC + action_history + pred_token"
|
||
f"{' + imu' if imu_dim > 0 else ''}）")
|
||
|
||
if args.pred_token_out:
|
||
import numpy as np
|
||
np.save(args.pred_token_out, model.pred_token.data.unsqueeze(1).numpy())
|
||
print(f"pred_token 已保存: {args.pred_token_out} (1,1,1,{model.pred_token.shape[-1]})")
|
||
|
||
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),
|
||
dummy_imu.squeeze(1).squeeze(1) if imu_dim > 0 else None)
|
||
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()
|
||
if imu_dim > 0:
|
||
feeds["imu"] = dummy_imu.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()
|