feat(pybot×actpred): R2 姿态反馈方案一——MPU6050 独立 IMU token 全链路
- pybot:recorder 逐帧落盘 imu.npy(T×8,roll,pitch,ax,ay,az,gx,gy,gz,与帧/动作同调用对齐);run_loop 对声明 imu_dim 的 predictor 以 predict(..., imu=) 调用(鸭子类型,旧两参调用向后兼容) - actpred:新增 --imu-dim;model 序列 [img×F, act×H, (IMU), PRED](imu_dim>0 时 F+H+2,IMU token 位于 PRED 前);dataset/train 读 imu.npy(缺失按零向量兜底);predict/ONNX/Hailo 三后端与 export/gen_calibration/dfc_compile/_npu_test 部署链适配 imu 输入 - 姿态按物理量程固定缩放归一化(imu_norm.py,训练/推理共用,零状态) - 测试:actpred 128 passed / pybot loop+recorder 28 passed
This commit is contained in:
1 parent
3e773debc9
commit
a2ba797410
22 files changed
+1236
-133
No files matched your search
+18
-11
@@ -23,31 +23,38 @@ def main():
|
||||
dummy_images = torch.randn(
|
||||
1, data_cfg.frame_len, 3, data_cfg.image_height, data_cfg.image_width)
|
||||
dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
|
||||
# R2 方案一:imu_dim>0 时追加归一化姿态输入 (1, imu_dim)
|
||||
imu_dim = data_cfg.imu_dim
|
||||
extra_inputs = (torch.randn(1, imu_dim),) if imu_dim > 0 else ()
|
||||
extra_names = ["imu"] if imu_dim > 0 else []
|
||||
|
||||
batch = Dim("batch")
|
||||
dynamic = {"images": {0: batch}, "action_history": {0: batch}}
|
||||
if imu_dim > 0:
|
||||
dynamic["imu"] = {0: batch}
|
||||
torch.onnx.export(
|
||||
model, (dummy_images, dummy_history),
|
||||
model, (dummy_images, dummy_history, *extra_inputs),
|
||||
args.output,
|
||||
input_names=["images", "action_history"],
|
||||
input_names=["images", "action_history"] + extra_names,
|
||||
output_names=["predicted_action"],
|
||||
dynamic_shapes={
|
||||
"images": {0: batch},
|
||||
"action_history": {0: batch},
|
||||
},
|
||||
dynamic_shapes=dynamic,
|
||||
opset_version=18,
|
||||
)
|
||||
print(f"ONNX model exported to: {args.output}")
|
||||
print(f"ONNX model exported to: {args.output}"
|
||||
f"{'(含 imu 输入)' if imu_dim > 0 else ''}")
|
||||
|
||||
# 可选:验证导出的 ONNX 模型一致性
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
import numpy as np
|
||||
sess = ort.InferenceSession(args.output)
|
||||
img_np = dummy_images.numpy()
|
||||
hist_np = dummy_history.numpy()
|
||||
ort_out = sess.run(None, {"images": img_np, "action_history": hist_np})[0]
|
||||
feeds = {"images": dummy_images.numpy(),
|
||||
"action_history": dummy_history.numpy()}
|
||||
for name, t in zip(extra_names, extra_inputs):
|
||||
feeds[name] = t.numpy()
|
||||
ort_out = sess.run(None, feeds)[0]
|
||||
with torch.no_grad():
|
||||
torch_out = model(dummy_images, dummy_history).numpy()
|
||||
torch_out = model(dummy_images, dummy_history, *extra_inputs).numpy()
|
||||
max_diff = np.abs(ort_out - torch_out).max()
|
||||
print(f"ONNX 一致性验证通过(最大差异: {max_diff:.6e})")
|
||||
except ImportError:
|
||||
|
||||
Reference in new issue
Block a user