pytorch模型转ONNX、并进行比较推理

简介: pytorch模型转ONNX、并进行比较推理

13.png

pytorch模型转ONNX


概述


ONNX(Open Neural Network Exchange)是一种开放式的深度学习模型交换格式,旨在促进不同深度学习框架之间的互操作性。通过将深度学习模型转换为ONNX格式,可以将其从一个深度学习框架移植到另一个框架中,而无需重新训练模型或手动重新实现模型结构。


ONNX是由微软和Facebook于2017年联合发布的,目前得到了众多公司和社区的支持和贡献。ONNX格式支持多种深度学习模型类型,包括卷积神经网络(CNN)、循环神经网络(RNN)和变换器(Transformer)等。


在使用ONNX时,通常需要两个步骤:模型的导出和模型的导入。首先,将模型从一个深度学习框架中导出到ONNX格式。导出后,可以使用ONNX运行时(ONNX Runtime)或其他支持ONNX格式的深度学习框架将模型导入到所需的平台中进行推理。


在导出模型时,需要注意确保模型的兼容性。由于每个框架实现深度学习算法的方式不同,因此在转换模型时可能会出现一些限制或不兼容性。此外,由于ONNX不支持某些框架特有的操作或层,因此在导出模型时可能需要进行一些自定义操作或调整模型结构,以便使模型兼容ONNX格式。


总的来说,ONNX是一个非常有用的工具,可以简化深度学习模型的转换和部署过程,促进不同框架之间的互操作性和合作。


PT转ONNX

import torch
import onnx
import onnxruntime
from efficient_net import build_model
# Step 1: 加载 PyTorch 模型
# model = EfficientNet.from_pretrained('efficientnet-b0')
model = build_model(model_name="efficientnet-b2",weights_path="ef_classify/models/best.pt", num_classes=2)
model.eval()
# Step 2: 将模型转换为 ONNX 格式(包含动态批次)
output_path = "efficientnet-b2.onnx"  # ONNX 模型保存路径
# 导出 ONNX 模型
dummy_input = torch.randn(1, 3, 260, 260)
dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
torch.onnx.export(model, dummy_input, output_path, verbose=True, input_names=['input'], output_names=['output'],
                  dynamic_axes=dynamic_axes)

比较PT和ONNX


要比较转换后的ONNX模型和原始的PyTorch模型之间的差异,可以使用ONNX Runtime工具包。


1.加载模型:在使用ONNX Runtime进行模型比较之前,您需要加载PyTorch模型和转换后的ONNX模型。您可以使用PyTorch或ONNX Runtime加载模型。

# 加载 PyTorch 模型
import torch
from efficient_net import build_model
# model_pt = torch.load('model.pt')
model_pt = build_model(model_name="efficientnet-b2",weights_path="ef_classify/models/best.pt", num_classes=2)
# 加载转换后的 ONNX 模型
import onnxruntime as ort
model_onnx = ort.InferenceSession('efficientnet-b2.onnx',providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

2.输入数据:为了比较模型,需要提供相同的输入数据。可以生成一些随机数据并将其提供给模型进行比较。

import numpy as np
input_shape = (1, 3, 260, 260)
input_data = np.random.random(input_shape).astype(np.float32)

3.推理和比较:使用PyTorch和ONNX Runtime对相同的输入数据进行推理并比较输出结果。如果输出相同,则说明两个模型在这些输入数据上是等价的。

# 使用 PyTorch 进行推理
model_pt.eval()
with torch.no_grad():
    output_pt = model_pt(torch.tensor(input_data))
# 使用 ONNX Runtime 进行推理
output_onnx = model_onnx.run(None, {'input': input_data})[0]
# 比较输出结果
if np.allclose(output_pt.numpy(), output_onnx, rtol=1e-3, atol=1e-5):
    print("模型输出相同!")
else:
    print("模型输出不同!")

使用这种方法,可以比较转换后的ONNX模型和原始的PyTorch模型在相同输入数据上的输出结果是否相同,从而评估模型的差距。


ONNX批量推理图片


在下面的代码中,我们使用了一个批处理的方式来进行推理。我们首先获取了所有图像文件的路径,然后将其分成了多个大小为batch_size的批次进行推理。在每个批次中,我们将多张图像一次性输入模型进行推理,并将输出结果保存在outputs列表中。最后,我们将所有输出结果合并并输出每个图像的预测结果。

import os
import cv2
import numpy as np
import onnxruntime as ort
# 加载ONNX模型
model_path = "efficientnet-b2.onnx"
session =  ort.InferenceSession(model_path,providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
input_name = session.get_inputs()[0].name
# 获取所有图像文件的路径
images_path = "ef_classify/dataset/dog"
image_paths = [os.path.join(images_path, f) for f in os.listdir(images_path) if os.path.isfile(os.path.join(images_path, f))]
# 批量读取、预处理和推理图像
batch_size = 8
outputs = []
for i in range(0, len(image_paths), batch_size):
    batch_images = []
    for image_path in image_paths[i:i+batch_size]:
        # 读取图像并进行预处理
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (260, 260))
        img = image.astype(np.float32)
        img /= 255.0
        img_mean = (0.485,0.456,0.406)
        img_std = (0.229,0.224,0.225)
        img = img - img_mean
        img = (img/img_std)
        image = img.transpose((2, 0, 1))
        image = np.expand_dims(image, axis=0).astype(np.float32)
        batch_images.append(image)
    # 进行推理并保存输出结果
    batch_images = np.concatenate(batch_images, axis=0)
    batch_outputs = session.run(None, {input_name: batch_images})
    outputs.append(batch_outputs[0])
# 将所有输出结果合并并输出每个图像的预测结果
outputs = np.concatenate(outputs, axis=0)
for i, image_path in enumerate(image_paths):
    predicted_class = np.argmax(outputs[i])
    print(f"Image: {image_path} - Predicted Class: {predicted_class}")
目录
相关文章
|
1月前
|
机器学习/深度学习 数据采集 人工智能
PyTorch学习实战:AI从数学基础到模型优化全流程精解
本文系统讲解人工智能、机器学习与深度学习的层级关系,涵盖PyTorch环境配置、张量操作、数据预处理、神经网络基础及模型训练全流程,结合数学原理与代码实践,深入浅出地介绍激活函数、反向传播等核心概念,助力快速入门深度学习。
90 1
|
3月前
|
机器学习/深度学习 PyTorch 测试技术
从训练到推理:Intel Extension for PyTorch混合精度优化完整指南
PyTorch作为主流深度学习框架,凭借动态计算图和异构计算支持,广泛应用于视觉与自然语言处理。Intel Extension for PyTorch针对Intel硬件深度优化,尤其在GPU上通过自动混合精度(AMP)提升训练与推理性能。本文以ResNet-50在CIFAR-10上的实验为例,详解如何利用该扩展实现高效深度学习优化。
192 0
|
5月前
|
机器学习/深度学习 PyTorch API
PyTorch量化感知训练技术:模型压缩与高精度边缘部署实践
本文深入探讨神经网络模型量化技术,重点讲解训练后量化(PTQ)与量化感知训练(QAT)两种主流方法。PTQ通过校准数据集确定量化参数,快速实现模型压缩,但精度损失较大;QAT在训练中引入伪量化操作,使模型适应低精度环境,显著提升量化后性能。文章结合PyTorch实现细节,介绍Eager模式、FX图模式及PyTorch 2导出量化等工具,并分享大语言模型Int4/Int8混合精度实践。最后总结量化最佳策略,包括逐通道量化、混合精度设置及目标硬件适配,助力高效部署深度学习模型。
742 21
PyTorch量化感知训练技术:模型压缩与高精度边缘部署实践
|
6月前
|
机器学习/深度学习 存储 缓存
加速LLM大模型推理,KV缓存技术详解与PyTorch实现
大型语言模型(LLM)的推理效率是AI领域的重要挑战。本文聚焦KV缓存技术,通过存储复用注意力机制中的Key和Value张量,减少冗余计算,显著提升推理效率。文章从理论到实践,详细解析KV缓存原理、实现与性能优势,并提供PyTorch代码示例。实验表明,该技术在长序列生成中可将推理时间降低近60%,为大模型优化提供了有效方案。
1065 15
加速LLM大模型推理,KV缓存技术详解与PyTorch实现
|
17天前
|
边缘计算 人工智能 PyTorch
130_知识蒸馏技术:温度参数与损失函数设计 - 教师-学生模型的优化策略与PyTorch实现
随着大型语言模型(LLM)的规模不断增长,部署这些模型面临着巨大的计算和资源挑战。以DeepSeek-R1为例,其671B参数的规模即使经过INT4量化后,仍需要至少6张高端GPU才能运行,这对于大多数中小型企业和研究机构来说成本过高。知识蒸馏作为一种有效的模型压缩技术,通过将大型教师模型的知识迁移到小型学生模型中,在显著降低模型复杂度的同时保留核心性能,成为解决这一问题的关键技术之一。
|
1月前
|
机器学习/深度学习 存储 PyTorch
Neural ODE原理与PyTorch实现:深度学习模型的自适应深度调节
Neural ODE将神经网络与微分方程结合,用连续思维建模数据演化,突破传统离散层的限制,实现自适应深度与高效连续学习。
77 3
Neural ODE原理与PyTorch实现:深度学习模型的自适应深度调节
|
2月前
|
PyTorch 算法框架/工具 异构计算
PyTorch 2.0性能优化实战:4种常见代码错误严重拖慢模型
我们将深入探讨图中断(graph breaks)和多图问题对性能的负面影响,并分析PyTorch模型开发中应当避免的常见错误模式。
169 9
|
7月前
|
机器学习/深度学习 JavaScript PyTorch
9个主流GAN损失函数的数学原理和Pytorch代码实现:从经典模型到现代变体
生成对抗网络(GAN)的训练效果高度依赖于损失函数的选择。本文介绍了经典GAN损失函数理论,并用PyTorch实现多种变体,包括原始GAN、LS-GAN、WGAN及WGAN-GP等。通过分析其原理与优劣,如LS-GAN提升训练稳定性、WGAN-GP改善图像质量,展示了不同场景下损失函数的设计思路。代码实现覆盖生成器与判别器的核心逻辑,为实际应用提供了重要参考。未来可探索组合优化与自适应设计以提升性能。
516 7
9个主流GAN损失函数的数学原理和Pytorch代码实现:从经典模型到现代变体
|
4月前
|
机器学习/深度学习 存储 PyTorch
PyTorch + MLFlow 实战:从零构建可追踪的深度学习模型训练系统
本文通过使用 Kaggle 数据集训练情感分析模型的实例,详细演示了如何将 PyTorch 与 MLFlow 进行深度集成,实现完整的实验跟踪、模型记录和结果可复现性管理。文章将系统性地介绍训练代码的核心组件,展示指标和工件的记录方法,并提供 MLFlow UI 的详细界面截图。
178 2
PyTorch + MLFlow 实战:从零构建可追踪的深度学习模型训练系统

推荐镜像

更多