DamoFD模型ONNX格式转换：多框架兼容方案

DamoFD模型ONNX格式转换多框架兼容方案1. 引言在实际的人脸检测项目部署中我们经常会遇到这样的困境训练好的PyTorch模型需要在不同的推理框架中运行比如TensorRT、OpenVINO或者移动端的NCNN。每次适配都要重新写一堆代码特别麻烦。DamoFD作为一款轻量级的人脸检测模型在精度和速度方面都有不错的表现。但如果只能局限在PyTorch环境中使用就大大限制了它的应用场景。今天我们就来解决这个问题——通过ONNX格式转换让DamoFD模型能够在各种推理框架中无缝运行。无论你是想在TensorRT上获得极致推理速度还是在移动端部署人脸检测功能这篇文章都会手把手教你完成整个转换过程并解决可能遇到的各种问题。2. 环境准备与模型下载2.1 基础环境配置首先确保你的环境中已经安装了必要的依赖库# 创建虚拟环境可选但推荐 conda create -n damofd_onnx python3.8 conda activate damofd_onnx # 安装核心依赖 pip install torch1.8.1 torchvision0.9.1 pip install onnx1.14.0 onnxruntime1.15.1 pip install modelscope1.9.5如果你打算在GPU上运行ONNX模型还需要安装对应的CUDA版本和cuDNN。这里我们以CUDA 11.1为例pip install onnxruntime-gpu1.15.12.2 下载DamoFD模型通过ModelScope库可以很方便地获取预训练好的DamoFD模型from modelscope.hub.snapshot_download import snapshot_download model_dir snapshot_download(damo/cv_ddsar_face-detection_iclr23-damofd) print(f模型下载到: {model_dir})下载完成后你会在指定目录看到模型文件和相关配置文件。主要关注的是.pth格式的模型权重文件。3. ONNX转换核心步骤3.1 加载原始PyTorch模型在进行转换之前我们需要先加载原始的PyTorch模型import torch from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks # 创建人脸检测pipeline face_detection pipeline( taskTasks.face_detection, modeldamo/cv_ddsar_face-detection_iclr23-damofd ) # 获取模型实例 model face_detection.model model.eval() # 设置为评估模式3.2 准备示例输入ONNX转换需要提供一个示例输入用于确定模型的输入形状和类型# 创建一个示例输入张量 # DamoFD的典型输入尺寸为640x6403通道 dummy_input torch.randn(1, 3, 640, 640).to(cuda if torch.cuda.is_available() else cpu) # 如果你有具体的图片也可以使用真实输入 # import cv2 # image cv2.imread(your_image.jpg) # image cv2.resize(image, (640, 640)) # image torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()3.3 执行ONNX转换现在开始最重要的转换步骤# 导出ONNX模型 onnx_path damofd_0.5g.onnx torch.onnx.export( model, # 要转换的模型 dummy_input, # 示例输入 onnx_path, # 输出文件路径 export_paramsTrue, # 导出模型参数 opset_version13, # ONNX算子集版本 do_constant_foldingTrue, # 优化常量折叠 input_names[input], # 输入节点名称 output_names[boxes, scores, landmarks], # 输出节点名称 dynamic_axes{ # 动态轴配置支持动态输入尺寸 input: {0: batch_size, 2: height, 3: width}, boxes: {0: batch_size, 1: num_detections}, scores: {0: batch_size, 1: num_detections}, landmarks: {0: batch_size, 1: num_detections} } ) print(fONNX模型已保存到: {onnx_path})这个转换过程通常只需要几秒钟。完成后你会得到一个.onnx文件这就是可以在多个框架中使用的通用模型格式。4. 转换验证与测试4.1 验证ONNX模型有效性转换完成后我们需要验证生成的ONNX模型是否有效import onnx # 加载并验证ONNX模型 onnx_model onnx.load(onnx_path) onnx.checker.check_model(onnx_model) print(ONNX模型验证通过) # 查看模型输入输出信息 print(输入信息:) for input in onnx_model.graph.input: print(f {input.name}: {[dim.dim_value for dim in input.type.tensor_type.shape.dim]}) print(输出信息:) for output in onnx_model.graph.output: print(f {output.name}: {[dim.dim_value for dim in output.type.tensor_type.shape.dim]})4.2 对比原始模型与ONNX模型推理结果为了确保转换没有影响模型精度我们需要对比两者的推理结果import numpy as np import onnxruntime as ort # 原始PyTorch模型推理 with torch.no_grad(): original_output model(dummy_input) # ONNX模型推理 ort_session ort.InferenceSession(onnx_path) onnx_output ort_session.run( None, {input: dummy_input.cpu().numpy()} ) # 对比输出结果 print(原始模型输出形状:, [o.shape for o in original_output]) print(ONNX模型输出形状:, [o.shape for o in onnx_output]) # 计算数值差异应该非常小 if isinstance(original_output, tuple): for i, (orig, onnx) in enumerate(zip(original_output, onnx_output)): diff np.max(np.abs(orig.cpu().numpy() - onnx)) print(f输出{i}最大差异: {diff}) else: diff np.max(np.abs(original_output.cpu().numpy() - onnx_output)) print(f最大差异: {diff})如果最大差异在可接受范围内通常小于1e-5说明转换成功。5. 常见问题与解决方案5.1 算子不支持错误在转换过程中你可能会遇到某些PyTorch算子不被ONNX支持的错误# 如果遇到不支持的算子可以尝试以下方法 # 1. 使用不同的opset版本 # opset_version11 或 opset_version12 # 2. 添加自定义算子支持 # 需要实现对应的ONNX算子 # 3. 修改模型结构避免使用不支持的算子对于DamoFD模型使用opset 13通常能够顺利转换。5.2 动态形状支持如果你需要处理不同尺寸的输入图片确保在转换时正确配置动态轴# 更详细的动态轴配置 dynamic_axes{ input: { 0: batch_size, # 批量大小动态 2: height, # 高度动态 3: width # 宽度动态 }, boxes: {0: batch_size, 1: num_detections}, scores: {0: batch_size, 1: num_detections}, landmarks: {0: batch_size, 1: num_detections, 2: points} }5.3 性能优化建议转换后的ONNX模型还可以进一步优化# 使用ONNX Runtime的优化工具 python -m onnxruntime.tools.convert_onnx_models_to_ort damofd_0.5g.onnx # 或者使用ONNX Optimizer pip install onnxoptimizer python -c import onnx from onnxoptimizer import optimize model onnx.load(damofd_0.5g.onnx) optimized_model optimize(model) onnx.save(optimized_model, damofd_0.5g_optimized.onnx) 6. 多框架部署示例6.1 TensorRT部署转换后的ONNX模型可以轻松部署到TensorRT# 使用trtexec工具转换ONNX到TensorRT # trtexec --onnxdamofd_0.5g.onnx --saveEnginedamofd_0.5g.trt --fp16 # 或者在Python中使用TensorRT API import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, logger) with open(damofd_0.5g.onnx, rb) as model: parser.parse(model.read()) # 构建优化引擎 config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) # 1GB serialized_engine builder.build_serialized_network(network, config) with open(damofd_0.5g.trt, wb) as f: f.write(serialized_engine)6.2 OpenVINO部署对于Intel硬件平台可以使用OpenVINO# 使用OpenVINO的模型优化器 mo --input_model damofd_0.5g.onnx --output_dir openvino_model6.3 ONNX Runtime推理最简单的方式是使用ONNX Runtime进行推理import onnxruntime as ort import numpy as np import cv2 # 创建推理会话 ort_session ort.InferenceSession(damofd_0.5g.onnx) # 预处理输入图像 image cv2.imread(test_image.jpg) image cv2.resize(image, (640, 640)) image image.transpose(2, 0, 1) # HWC to CHW image np.expand_dims(image, axis0).astype(np.float32) # 推理 outputs ort_session.run(None, {input: image}) # 后处理 boxes, scores, landmarks outputs print(f检测到 {len(boxes[0])} 张人脸)7. 总结通过ONNX格式转换我们成功地将DamoFD人脸检测模型从PyTorch框架中解放出来使其能够在多种推理框架中运行。这个过程其实并不复杂关键是要注意输入输出的匹配和动态形状的配置。实际使用中ONNX格式确实带来了很大的便利性。我们团队在多个项目中都采用了这种方案大大减少了模型部署的工作量。特别是在需要同时支持多种硬件平台的场景下ONNX的统一接口显得格外有价值。如果你在转换过程中遇到问题建议先从简单的静态形状开始成功后再尝试动态形状。另外ONNX社区很活跃遇到特定算子的支持问题通常都能找到解决方案。转换后的模型可以在保持精度的同时获得在不同硬件上的加速效果这对于实际应用来说是非常有价值的。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。