1. 背景与平台介绍海思 Hi3519DV500 是面向智能 IPC、人脸识别、车辆识别等场景的高性能视觉 SoC集成了双核 NNIENeural Network Inference Engine以及 DSP 加速单元支持多种深度学习网络的高效推理。将训练好的模型部署到这类终端设备上核心环节之一就是模型转换本文主要介绍将 PyTorch 训练得到的 .pth 权重文件转换为海思平台专用 .om 权重文件的方法。2. 准备工作2.1 以Ubuntu系统为例需要在conda环境中新建一个python3.7.5虚拟环境conda create-npy3_7_5python3.7.5 conda activate py3_7_52.2 下载Hi3519DV500官方SDK找到文件夹下的MindStudio安装Ascend-cann-toolkitcdMindStudiochmodx Ascend-cann-toolkit_6.10.t03spc011b090_linux.x86_64.runsudo./Ascend-cann-toolkit_6.10.t03spc011b090_linux.x86_64.run--install安装成功后安装路径为: /usr/local/Ascend设置环境变量source/usr/local/Ascend/ascend-toolkit/svp_latest/x86_64-linux/script/setenv.sh3. 模型转换3.1 将训练好的.pth文件转换为.onnx文件以ResNet50为例。importosimporttorch from resnet50importresnet50 def export_to_onnx(weights_path, onnx_path,num_classes2,opset_version13): 将PyTorch模型(.pth/.pt)转换为ONNX格式 参数: weights_path: PyTorch权重文件路径(.pth 或 .pt)onnx_path: 输出ONNX文件路径 num_classes: 分类类别数量 opset_version: ONNX opset版本(推荐13) devicetorch.device(cpu)# 创建模型print(fCreating ResNet50 model with {num_classes} classes...)modelresnet50(num_classesnum_classes)# 加载权重assert os.path.exists(weights_path), fWeights file not found: {weights_path}print(fLoading weights from: {weights_path})model.load_state_dict(torch.load(weights_path,map_locationdevice))# 设置为推理模式model.eval()# 创建示例输入 (1, 3, 224, 224)dummy_inputtorch.randn(1,3,224,224,devicedevice)# 导出为ONNXprint(fExporting to ONNX: {onnx_path})torch.onnx.export(model, dummy_input, onnx_path,export_paramsTrue,opset_versionopset_version,do_constant_foldingTrue,input_names[input],output_names[output],dynamic_axes{input:{0:batch_size},output:{0:batch_size}},dynamoFalse)print(fONNX model saved to: {onnx_path})# 验证ONNX模型try:importonnx print(Verifying ONNX model...)onnx_modelonnx.load(onnx_path)onnx.checker.check_model(onnx_model)print(ONNX model verification passed!)# 打印模型信息print(f\nONNX Model Info:)print(f - Opset version: {onnx_model.opset_import[0].version})print(f - Inputs: {[i.name for i in onnx_model.graph.input]})print(f - Outputs: {[o.name for o in onnx_model.graph.output]})except ImportError: print(onnx package not installed, skipping verification.)print(Install with: pip install onnx)except Exception as e: print(fONNX verification failed: {e})if__name____main__:# 配置参数WEIGHTS_PATH./resNet50.pthONNX_PATH./resnet50.onnxNUM_CLASSES2OPSET_VERSION13export_to_onnx(WEIGHTS_PATH, ONNX_PATH, NUM_CLASSES, OPSET_VERSION)3.2 在官方源码中找到对应的模型官方sample的readme中有使用ATC工具进行模型转换的相关指令以onnx格式为例atc--model./model/resnet50.onnx--framework5--output./model/resnet50--input_shapeinput:1,3,224,224--input_formatNCHW--soc_versionHi3519DV500--insert_op_conf./insert_op.cfg--image_list./data/image_paths.txt--input_typeinput:UINT8--batch_num1--log_level0--dump_data0--compile_mode0--save_original_modeltrue--online_model_type0--model指定ONNX模型文件路径。--framework5声明输入为ONNX格式5固定标识ONNX。 --output设置.om离线模型输出路径及文件名前缀。 --input_shape定义输入张量形状格式节点名:批大小,通道,高,宽必须与实际数据一致。--input_formatNCHW指定输入数据内存排布顺序N批大小,C通道,H高,W宽。--soc_versionHi3519DV500必须匹配目标芯片型号用于适配硬件特性。 --insert_op_conf指定AIPP预处理配置文件实现硬件级图像预处理如归一化。 --image_list提供校准图像路径列表每行1张INT8量化必需直接影响精度。--input_typeUINT8声明输入为8位整型0~255与AIPP配置联动。--batch_num1量化校准时单次处理图像数量嵌入式设备通常为1。--save_original_modeltrue保留原始模型中间表示用于问题分析。--log_level0日志级别0仅错误调试建议设3。--dump_data0是否转存中间数据0不转存问题定位时设1。--compile_mode0编译模式0标准优化1快速模式。--online_model_type0模型类型0离线模型ATC仅支持0。其中需要注意input_shape需要通过Netron工具查看实际输入点击最顶部节点通常是Conv层在右侧的属性面板中查看 Inputs 对应的名称和维度。image_list是ATC模型转换时用于校准量化的图像路径列表文件其内容格式和路径要求极为严格。必须满足每行是一个可访问的图像绝对路径且所有图像尺寸必须严格匹配input_shape参数指定的尺寸。