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）。
179 lines
8.0 KiB
Python
179 lines
8.0 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 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()
|