用途:在华为昇腾 NPU 上将 PyTorch 模型经 ONNX 转为 OM,并完成端到端推理适配,支持任意标准 PyTorch 或 ONNX 模型。
主要能力:六个工作流。Workflow 1 通过静态代码分析(搜索 cv2.resize、transforms.Resize、argparse、config 等)或只读探针脚本发现 input_shape、input_names、opset 等参数,遵循反硬编码规则,禁止使用猜测默认值;识别分类/检测/分割/姿态/生成等后处理类型。Workflow 2 用 export_onnx.py 导出 ONNX,支持动态 axes 与 torchvision 预训练模型。Workflow 3 用 get_onnx_info.py 检查模型并执行 ATC 转换。Workflow 4 复用原仓库预处理/后处理逻辑生成 e2e_infer_om.py,替换推理为 OM。Workflow 5 通过 infer_om.py 和 compare_precision.py 做冒烟测试与 ONNX/OM 精度对比。Workflow 6 生成可复现 README。
必要用法:先收集模型权重路径与代码仓库地址。导出示例:python3 scripts/export_onnx.py --pt_model model.pt --output model.onnx --input_shape 1,3,224,224。环境配置:./scripts/setup_env.sh。ATC 示例:atc --model=model.onnx --framework=5 --output=model_om --soc_version=Ascend910B3 --input_shape="input:1,3,224,224"。SoC 版本须与设备一致,用 npu-smi info 查询。
适用场景:昇腾 NPU 上 PyTorch 模型的部署、模型迁移、精度与性能验证。
限制:Python 必须 ≤3.10(推荐 3.10),NumPy 必须 <2.0,ONNX opset 推荐 11;支持 CANN 8.1.RC1、8.3.RC1、8.5.0+。纯 state dict 不支持导出(需模型架构定义)。end2end 模型张量级对比可能不适用,需以业务结果为准。