qing
63cfdd2066
- 第 1 轮安全审计修复: F-1~F-10（路径校验、预训练架构修正等） - 第 2 轮复审修复 11 项: 3 项上轮修复不完整（R01 checkpoint 键校验/ R02 fused_hidden_dim 值域/R03 pretrained=False）+ 8 项新问题 （actions.npy 形状校验、帧数警告、CLI 值域校验、stats 形状与路径、 onnxscript 依赖、归一化语义、解析去重） - 单元测试由 17 扩展至 51 用例（50 通过 + 1 环境性跳过） - 审计文档两轮合并至 actpred_CODE_AUDIT_20260808.md，更新 INDEX
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
import argparse
|
||
import torch
|
||
from config import parse_checkpoint_config
|
||
from model import ActionPredictor
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Export actpred model to ONNX")
|
||
parser.add_argument("--model", required=True, help="checkpoint path (best_inference.pt)")
|
||
parser.add_argument("--output", default="model.onnx", help="output ONNX file path")
|
||
args = parser.parse_args()
|
||
|
||
checkpoint = torch.load(args.model, map_location="cpu", weights_only=True)
|
||
cfg = parse_checkpoint_config(checkpoint["config"])
|
||
data_cfg = cfg.data
|
||
model_cfg = cfg.model
|
||
|
||
model = ActionPredictor(model_cfg, data_cfg, pretrained=False)
|
||
model.load_state_dict(checkpoint["model_state_dict"])
|
||
model.eval()
|
||
|
||
dummy_image = torch.randn(1, 3, data_cfg.image_size, data_cfg.image_size)
|
||
dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
|
||
|
||
torch.onnx.export(
|
||
model, (dummy_image, dummy_history),
|
||
args.output,
|
||
input_names=["image", "action_history"],
|
||
output_names=["predicted_action"],
|
||
dynamic_axes={
|
||
"image": {0: "batch"},
|
||
"action_history": {0: "batch"},
|
||
"predicted_action": {0: "batch"},
|
||
},
|
||
opset_version=17,
|
||
)
|
||
print(f"ONNX model exported to: {args.output}")
|
||
|
||
# 可选：验证导出的 ONNX 模型一致性
|
||
try:
|
||
import onnxruntime as ort
|
||
import numpy as np
|
||
sess = ort.InferenceSession(args.output)
|
||
img_np = dummy_image.numpy()
|
||
hist_np = dummy_history.numpy()
|
||
ort_out = sess.run(None, {"image": img_np, "action_history": hist_np})[0]
|
||
with torch.no_grad():
|
||
torch_out = model(dummy_image, dummy_history).numpy()
|
||
max_diff = np.abs(ort_out - torch_out).max()
|
||
print(f"ONNX 一致性验证通过（最大差异: {max_diff:.6e}）")
|
||
except ImportError:
|
||
print("提示: 安装 onnxruntime 可验证导出模型一致性 (pip install onnxruntime)")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|