梅河口信息网上海百度seo牛巨微

如何在Windows中利用python版本的onnxruntime部署pytorch模型？
部署的目的是为了脱离原有的训练环境，让模型在生产环境能够运行，并且加速。
我这里以deeplabv3+为例。

第一步：

在pytorch训练环境中安装onnx

pip install onnx

第二步：

利用pytorch导出onnx格式：

# 导出训练好的模型，以供后续部署
import os
import argparse
import torch
from modeling.deeplab import DeepLab#换成自己的模型
if __name__ == '__main__':    # 读取模型pth文件model_path = r'F:\dataset\CoverSegData_train_result\deeplab-resnet-20250606\model_best.pth.tar'if not os.path.exists(model_path):raise FileNotFoundError(f"Model file {model_path} does not exist.")model = DeepLab(num_classes=2, backbone='resnet', output_stride=16, sync_bn=False, freeze_bn=False)checkpoint = torch.load(model_path, weights_only=False)model.load_state_dict(checkpoint['state_dict'])model.eval()  # 设置模型为评估模式# 导出为torchscript格式dumpy = torch.randn(1, 3, 1024, 1024)  # 假设输入是1张1024x1024的RGB图像# 使用torch.jit.trace进行模型跟踪trace_model = torch.jit.trace(model, dumpy)   export_path = r'F:\dataset\CoverSegData_train_result\deeplab-resnet-20250606\deeplab_resnet_exported.pt'trace_model.save(export_path)print(f"Model exported to {export_path}")# 导出为ONNX格式onnx_export_path = r'F:\dataset\CoverSegData_train_result\deeplab-resnet-20250606\deeplab_resnet_exported.onnx'torch.onnx.export(model,dumpy,onnx_export_path,export_params=True,opset_version=11,do_constant_folding=True,input_names=['input'],output_names=['output'])print(f"Model exported to {onnx_export_path}")

第三步：

conda新建一个推理环境，并安装onnxruntime

conda create -n ort_env python=3.10

第四步：

在ort_env环境中安装onnxruntime
如果安装cpu版本：

pip install onnxruntime

如果安装gpu版本：

pip install onnxruntime-gpu

第五步：

如果安装的cpu版本跳过这一步。
安装cuda和cudnn。

win+R->cmd窗口，查询Windows系统中的cuda和cudnn
1、查询cuda路径：

echo %CUDA_PATH%

这一条命令会找到你的cuda安装的位置。我的电脑显示E:\nvidia\cuda12.2。我下面就以我的路径为例，自己电脑换成自己的路径即可。
表示我安装的是cuda12.2，也可以通过nvidia-smi查看。
如果找不到目录，去官网下载安装cuda
重启电脑，这时候重新运行上面的命令，如果还是找不到cuda目录，在系统环境变量->path中添加：CUDA_PATH，E:\nvidia\cuda12.2，重启电脑生效。
2、查询cudnn位置

where cudart64_*.dll

这一条命令会找到你的cudnn安装的位置。cudnn的目录一定在cuda\bin里面。我这里输出的是E:\nvidia\cuda12.2\bin\cudnn64_9.dll。
如果找不到目录，官网下载安装cudnn.安装完cudnn之后，把cudnn目录复制到cuda目录下面：

复制 bin\*     到 E:\nvidia\cuda12.2\bin
复制 include\* 到 E:\nvidia\cuda12.2\include
复制 lib\*     到 E:\nvidia\cuda12.2\lib\x64

第六步

在推理阶段一定要保证和导出onnx输入的照片(张量变量：dumpy)通道和尺寸一样！
核心推理代码：

      if self.use_gpu and 'CUDAExecutionProvider' in ort.get_available_providers():providers = ['CUDAExecutionProvider']self.provider = 'GPU'else:providers = ['CPUExecutionProvider']self.provider = 'CPU'print(f"使用提供者: {self.provider}")# 创建会话选项session_options = ort.SessionOptions()session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALLself.session = ort.InferenceSessionself.onnx_path, providers=providers,sess_options=session_optionsif self.provider == 'GPU':# 使用 IO 绑定优化 GPU 推理io_binding = self.session.io_binding()# 绑定输入到 GPUinput_tensor = ort.OrtValue.ortvalue_from_numpy(input_img, 'cuda', 0)io_binding.bind_input(self.input_name, 'cuda', 0, np.float32, input_img.shape, input_tensor.data_ptr())output_shape = (1, 2, self.input_shape[2], self.input_shape[3])output_array = np.empty(output_shape, dtype=np.float32)output_ortvalue = ort.OrtValue.ortvalue_from_numpy(output_array, 'cuda', 0)# 绑定输出io_binding.bind_ortvalue_output(self.output_name, output_ortvalue)  # 更简洁的新API# 运行推理self.session.run_with_iobinding(io_binding)# 获取输出# outputs = output_tensor.numpy()outputs = output_ortvalue.numpy()else:# 普通 CPU 推理outputs = self.session.run([self.output_name], {self.input_name: input_img})[0]

因为推理的第一张照片需要编译模型并且拷贝显存，所以不能反应真实的推理时间，所以在我的代码中支持输入文件夹批量推理，并对比cpu和gpu推理的时间（去掉第一张照片），代码如下：

import onnxruntime as ort
import numpy as np
import time
import cv2
import os
import glob
from tqdm import tqdmdef preprocess_image(img, target_size=(1024, 1024)):"""预处理图像以匹配模型输入要求"""# 如果输入是文件路径，则读取图像if isinstance(img, str):img = cv2.imread(img)if img is None:raise ValueError(f"无法读取图像: {img}")# 确保图像是 NumPy 数组if not isinstance(img, np.ndarray):raise TypeError("输入必须是文件路径或 NumPy 数组")# 转换颜色空间 (OpenCV 默认是 BGR)if len(img.shape) == 3 and img.shape[2] == 3:img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)# 调整大小img = cv2.resize(img, target_size)# 归一化 (根据模型训练时的预处理)img = img.astype(np.float32) / 255.0# 转换为 CHW 格式 [Height, Width, Channel] -> [Channel, Height, Width]if len(img.shape) == 3:img = np.transpose(img, (2, 0, 1))# 添加批次维度 [Channel, Height, Width] -> [Batch, Channel, Height, Width]img = np.expand_dims(img, axis=0)return imgdef visualize_segmentation(output, original_img, class_colors=None):"""可视化分割结果"""if class_colors is None:# 默认颜色映射 (背景, 前景)class_colors = [(0, 0, 0), (255, 0, 0)]  # 黑色背景，红色前景# 获取预测类别pred_class = np.argmax(output[0], axis=0).astype(np.uint8)# 创建彩色分割图seg_image = np.zeros((pred_class.shape[0], pred_class.shape[1], 3), dtype=np.uint8)for class_idx, color in enumerate(class_colors):seg_image[pred_class == class_idx] = color# 如果原始图像是预处理过的，还原它if original_img.shape[0] == 3 or original_img.shape[0] == 1:  # [C, H, W]orig_image = (np.transpose(original_img, (1, 2, 0)) * 255).astype(np.uint8)else:  # [H, W, C]orig_image = (original_img * 255).astype(np.uint8)# 叠加原图与分割结果overlay = cv2.addWeighted(orig_image, 0.7, seg_image, 0.3, 0)# 创建纯分割掩码mask = seg_imagereturn overlay, maskclass ONNXModelInferencer:def __init__(self, onnx_path, use_gpu=True):"""初始化 ONNX 模型推理器"""self.onnx_path = onnx_pathself.use_gpu = use_gpuself.session = Noneself.input_name = Noneself.output_name = Noneself.input_shape = Noneself.provider = Noneself._initialize_session()def _initialize_session(self):"""初始化 ONNX Runtime 会话"""# 选择提供者if self.use_gpu and 'CUDAExecutionProvider' in ort.get_available_providers():providers = ['CUDAExecutionProvider']self.provider = 'GPU'else:providers = ['CPUExecutionProvider']self.provider = 'CPU'print(f"使用提供者: {self.provider}")# 创建会话选项session_options = ort.SessionOptions()session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALLtry:self.session = ort.InferenceSession(self.onnx_path, providers=providers,sess_options=session_options)# 获取输入输出信息self.input_name = self.session.get_inputs()[0].nameinput_info = self.session.get_inputs()[0]self.input_shape = input_info.shape  # [batch, channels, height, width]self.output_name = self.session.get_outputs()[0].nameprint(f" ONNX 模型加载成功 (提供者: {self.provider})")print(f"输入名称: {self.input_name}, 形状: {self.input_shape}")print(f"输出名称: {self.output_name}")except Exception as e:print(f" 无法创建 ONNX 会话: {e}")raisedef infer_image(self, image_input, return_raw=False):"""对单个图像进行推理"""# 预处理图像if isinstance(image_input, str):  # 文件路径img = cv2.imread(image_input)if img is None:raise ValueError(f"无法读取图像: {image_input}")original_img = img.copy()else:  # NumPy 数组original_img = image_input.copy()# 预处理图像input_img = preprocess_image(original_img, target_size=(self.input_shape[3], self.input_shape[2])  # (width, height))# 运行推理start_time = time.time()if self.provider == 'GPU':# 使用 IO 绑定优化 GPU 推理io_binding = self.session.io_binding()# 绑定输入到 GPUinput_tensor = ort.OrtValue.ortvalue_from_numpy(input_img, 'cuda', 0)io_binding.bind_input(self.input_name, 'cuda', 0, np.float32, input_img.shape, input_tensor.data_ptr())# 绑定输出到 GPU# output_shape = (1, 2, self.input_shape[2], self.input_shape[3])  # [1, num_classes, H, W]# output_tensor = ort.OrtValue.ortvalue_from_shape(output_shape, np.float32, 'cuda', 0)# io_binding.bind_output(#     self.output_name, #     'cuda', #     0, #     np.float32, #     output_shape, #     output_tensor.data_ptr()# )output_shape = (1, 2, self.input_shape[2], self.input_shape[3])output_array = np.empty(output_shape, dtype=np.float32)output_ortvalue = ort.OrtValue.ortvalue_from_numpy(output_array, 'cuda', 0)# 绑定输出io_binding.bind_ortvalue_output(self.output_name, output_ortvalue)  # 更简洁的新API# 运行推理self.session.run_with_iobinding(io_binding)# 获取输出# outputs = output_tensor.numpy()outputs = output_ortvalue.numpy()else:# 普通 CPU 推理outputs = self.session.run([self.output_name], {self.input_name: input_img})[0]inference_time = time.time() - start_timeif return_raw:return outputs, inference_time# 可视化结果overlay, mask = visualize_segmentation(outputs, original_img)return overlay, mask, inference_timedef batch_infer_folder(self, input_folder, output_folder, save_overlay=True, save_mask=True):"""批量处理整个文件夹中的图像"""# 创建输出文件夹os.makedirs(output_folder, exist_ok=True)# 获取所有支持的图像文件image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff']image_paths = []for ext in image_extensions:image_paths.extend(glob.glob(os.path.join(input_folder, ext)))if not image_paths:print(f" 在文件夹 {input_folder} 中未找到图像文件")return 0  # 返回0表示没有处理任何图像print(f" 找到 {len(image_paths)} 张图像进行处理")# 处理统计total_time = 0min_time = float('inf')max_time = 0processed_count = 0  # 记录成功处理的图像数量# 处理每张图像for img_path in tqdm(image_paths, desc=f"处理图像 ({self.provider})"):try:# 推理图像overlay, mask, infer_time = self.infer_image(img_path)# 更新计时统计# 不统计第一张照片，因为第一张照片涉及编译和内存分配等开销if processed_count != 0:total_time += infer_timeprocessed_count += 1min_time = min(min_time, infer_time)max_time = max(max_time, infer_time)# 准备输出路径filename = os.path.basename(img_path)name, ext = os.path.splitext(filename)# 保存结果if save_overlay:overlay_bgr = cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR)overlay_path = os.path.join(output_folder, f"{name}_overlay{ext}")cv2.imwrite(overlay_path, overlay_bgr)if save_mask:mask_bgr = cv2.cvtColor(mask, cv2.COLOR_RGB2BGR)mask_path = os.path.join(output_folder, f"{name}_mask{ext}")cv2.imwrite(mask_path, mask_bgr)except Exception as e:print(f" 处理 {img_path} 时出错: {str(e)}")# 打印统计信息print("\n" + "="*50)print(f"处理完成 ({self.provider})")print(f"总图像数: {len(image_paths)}")print(f"成功处理: {processed_count} 张图像")if processed_count > 0:avg_time = total_time / (processed_count-1)print(f"总推理时间: {total_time:.4f} 秒")print(f"平均推理时间: {avg_time:.4f} 秒/图像")print(f"最短推理时间: {min_time:.4f} 秒")print(f"最长推理时间: {max_time:.4f} 秒")print(f"速度: {(processed_count-1)/total_time:.2f} 图像/秒")else:print(" 没有图像被成功处理")print("="*50)return total_timedef compare_gpu_cpu_performance(onnx_path, input_folder, output_base_folder):"""比较 GPU 和 CPU 性能"""# 创建 GPU 推理器print("\n" + "="*50)print("初始化 GPU 推理器")print("="*50)gpu_inferencer = ONNXModelInferencer(onnx_path, use_gpu=True)# 创建 CPU 推理器print("\n" + "="*50)print("初始化 CPU 推理器")print("="*50)cpu_inferencer = ONNXModelInferencer(onnx_path, use_gpu=False)# 创建输出文件夹gpu_output_folder = os.path.join(output_base_folder, "gpu_results")cpu_output_folder = os.path.join(output_base_folder, "cpu_results")os.makedirs(gpu_output_folder, exist_ok=True)os.makedirs(cpu_output_folder, exist_ok=True)# 运行 GPU 推理print("\n" + "="*50)print("开始 GPU 批量推理")print("="*50)gpu_time = gpu_inferencer.batch_infer_folder(input_folder, gpu_output_folder)# 运行 CPU 推理print("\n" + "="*50)print("开始 CPU 批量推理")print("="*50)cpu_time = cpu_inferencer.batch_infer_folder(input_folder, cpu_output_folder)# 性能比较 - 只有当两者都成功处理了图像时才进行比较if gpu_time is not None and cpu_time is not None and gpu_time > 0 and cpu_time > 0:print("\n" + "="*50)print("性能比较结果")print("="*50)print(f"GPU 总时间: {gpu_time:.4f} 秒")print(f"CPU 总时间: {cpu_time:.4f} 秒")print(f"加速比: {cpu_time/gpu_time:.2f}x")print("="*50)else:print("\n 无法比较性能，因为一个或多个设备没有处理任何图像")def batch_infer_all_images(onnx_path, input_folder, output_folder, use_gpu=True):"""批量处理整个文件夹中的图像（单一设备）"""# 创建推理器print("\n" + "="*50)device_name = "GPU" if use_gpu else "CPU"print(f"初始化 {device_name} 推理器")print("="*50)inferencer = ONNXModelInferencer(onnx_path, use_gpu=use_gpu)# 运行推理print("\n" + "="*50)print(f"开始 {device_name} 批量推理")print("="*50)inferencer.batch_infer_folder(input_folder, output_folder)if __name__ == "__main__":# 配置参数onnx_model_path = r'F:\dataset\CoverSegData_train_result\deeplab-resnet-20250606\deeplab_resnet_exported.onnx'input_image_folder = r'F:\dataset\test'  # 替换为你的图像文件夹路径output_base_folder = r'F:\dataset\test_results'  # 结果输出文件夹# 检查 ONNX 文件是否存在if not os.path.exists(onnx_model_path):raise FileNotFoundError(f"ONNX 模型文件 {onnx_model_path} 不存在")# 检查输入文件夹是否存在if not os.path.exists(input_image_folder):raise FileNotFoundError(f"输入文件夹 {input_image_folder} 不存在")# 创建输出基础文件夹os.makedirs(output_base_folder, exist_ok=True)# 选择运行模式mode = input("请选择运行模式:\n1. 比较 GPU 和 CPU 性能\n2. 仅使用 GPU 推理\n3. 仅使用 CPU 推理\n> ")if mode == '1':# 比较 GPU 和 CPU 性能compare_gpu_cpu_performance(onnx_model_path, input_image_folder, output_base_folder)elif mode == '2':# 仅使用 GPU 推理gpu_output_folder = os.path.join(output_base_folder, "gpu_only")batch_infer_all_images(onnx_model_path, input_image_folder, gpu_output_folder, use_gpu=True)elif mode == '3':# 仅使用 CPU 推理cpu_output_folder = os.path.join(output_base_folder, "cpu_only")batch_infer_all_images(onnx_model_path, input_image_folder, cpu_output_folder, use_gpu=False)else:print(" 无效的选择，请重新运行程序并输入 1, 2 或 3")

结果：

PS F:\python_projects\deeplab_deploy> & E:/anaconda3/python.exe f:/python_projects/deeplab_deploy/test2.py
请选择运行模式:
1. 比较 GPU 和 CPU 性能
2. 仅使用 GPU 推理
3. 仅使用 CPU 推理
> 1==================================================
初始化 GPU 推理器
==================================================
使用提供者: GPU
ONNX 模型加载成功 (提供者: GPU)
输入名称: input, 形状: [1, 3, 1024, 1024]
输出名称: output==================================================
初始化 CPU 推理器
==================================================
使用提供者: CPU
ONNX 模型加载成功 (提供者: CPU)
输入名称: input, 形状: [1, 3, 1024, 1024]
输出名称: output性能比较结果
==================================================
GPU 总时间: 0.4545 秒
CPU 总时间: 6.2390 秒
加速比: 13.73x
==================================================

可以看出来速度提高了13倍。
参考：
https://onnxruntime.ai/docs/get-started/with-python.html
https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements