当前位置：首页 > news >正文

12.6深度学习_模型优化和迁移_模型移植

news 2025/7/27 11:49:28

八、模型移植

1. 认识ONNX

https://onnx.ai/

Open Neural Network Exchange（ONNX，开放神经网络交换）格式，是一个用于表示深度学习模型的标准，可使模型在不同框架之间进行转移。

ONNX的规范及代码主要由微软，亚马逊，Face book 和 IBM等公司共同开发，以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有： Caffe2, PyTorch, PaddlePaddle， TensorFlow等。

2. 导出ONNX

2.1 安装依赖包

pip install onnx
pip install onnxruntime

2.2 导出ONNX模型

import os
import torch
import torch.nn as nn
from torchvision.models import resnet18if __name__ == "__main__":dir = os.path.dirname(__file__)weightpath = os.path.join(os.path.dirname(__file__), "pth", "resnet18_default_weight.pth")onnxpath = os.path.join(os.path.dirname(__file__), "pth", "resnet18_default_weight.onnx")device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = resnet18(pretrained=False)model.conv1 = nn.Conv2d(#in_channels=3,out_channels=64,kernel_size=3,stride=1,padding=0,bias=False,)# 删除池化层model.maxpool = nn.MaxPool2d(kernel_size=1, stride=1, padding=0)# 修改全连接层in_feature = model.fc.in_featuresmodel.fc = nn.Linear(in_feature, 10)model.load_state_dict(torch.load(weightpath, map_location=device))model.to(device)# 创建一个实例输入x = torch.randn(1, 3, 224, 224, device=device)# 导出onnxtorch.onnx.export(model,x,onnxpath,#verbose=True, # 输出转换过程input_names=["input"],output_names=["output"],)print("onnx导出成功")

2.3 ONNX结构可视化

可以直接在线查看：https://netron.app/

也可以下载桌面版：https://github.com/lutzroeder/netron

3. ONNX推理

ONNX在做推理时不再需要导入网络，且适用于Python、JAVA、PyQT等各种语言，不再依赖于PyTorch框架；

3.1 简单推理

import onnxruntime as ort
import torchvision.transforms as transforms
import cv2 as cv
import os
import numpy as npimg_size = 224
transformtest = transforms.Compose([transforms.ToPILImage(),  # 将numpy数组转换为PIL图像transforms.Resize((img_size, img_size)),transforms.ToTensor(),transforms.Normalize(# 均值和标准差mean=[0.4914, 0.4822, 0.4465],std=[0.2471, 0.2435, 0.2616],),]
)def softmax(x):e_x = np.exp(x - np.max(x))return e_x / e_x.sum(axis=1, keepdims=True)def cv_imread(file_path):cv_img = cv.imdecode(np.fromfile(file_path, dtype=np.uint8), cv.IMREAD_COLOR)return cv_imglablename = "飞机、汽车、鸟类、猫、鹿、狗、青蛙、马、船和卡车".split("、")if __name__ == "__main__":dir = os.path.dirname(__file__)weightpath = os.path.join(os.path.dirname(__file__), "pth", "resnet18_default_weight.pth")onnxpath = os.path.join(os.path.dirname(__file__), "pth", "resnet18_default_weight.onnx")# 读取图片img_path = os.path.join(dir, "test", "5.jpg")img = cv_imread(img_path)img = cv.cvtColor(img, cv.COLOR_BGR2RGB)img_tensor = transformtest(img)# 将图片转换为ONNX运行时所需的格式img_numpy = img_tensor.numpy()img_numpy = np.expand_dims(img_numpy, axis=0)  # 增加batch_size维度# 加载onnx模型sess = ort.InferenceSession(onnxpath)# 运行onnx模型outputs = sess.run(None, {"input": img_numpy})output = outputs[0]# 应用softmaxprobabilities = softmax(output)print(probabilities)# 获得预测结果pred_index = np.argmax(probabilities, axis=1)pred_value = probabilities[0][pred_index[0]]print(pred_index)print("预测目标：",lablename[pred_index[0]],"预测概率：",str(pred_value * 100)[:5] + "%",)

输出结果：

[[6.7321511e-05 9.7113671e-11 7.6417709e-05 2.8661249e-02 7.0206769e-043.9052707e-04 9.7010124e-01 6.8206714e-07 4.1351362e-07 5.7089373e-09]]
[6]
预测目标： 青蛙 预测概率： 97.01%

3.2 使用GPU推理

需要安装依赖包：

pip install onnxruntime-gpu

代码：

# 导入FileSystemStorage
import time
import random
import os# 人工智能推理用到的模块
import onnxruntime as ort
import torchvision.transforms as transforms
import numpy as np
import PIL.Image as Imageimg_size = 32
transformtest = transforms.Compose([transforms.Resize((img_size, img_size)),transforms.ToTensor(),transforms.Normalize(# 均值和标准差mean=[0.4914, 0.4822, 0.4465],std=[0.2471, 0.2435, 0.2616],),]
)def softmax(x):e_x = np.exp(x - np.max(x))return e_x / e_x.sum(axis=1, keepdims=True)def imgclass():# AI推理# 读取图片imgpath = os.path.join(os.path.dirname(__file__), "..", "static/ai", filename)# 加载并预处理图像image = Image.open(imgpath)input_tensor = transformtest(image)input_tensor = input_tensor.unsqueeze(0)  # 添加批量维度# 将图片转换为ONNX运行时所需的格式img_numpy = input_tensor.numpy()# 加载模型onnxPath = os.path.join(#os.path.dirname(__file__),"..","onnx","resnet18_default_weight_1.onnx",)# 设置 ONNX Runtime 使用 GPUproviders = ["CUDAExecutionProvider"]sess = ort.InferenceSession(onnxPath, providers=providers)# 使用模型对图片进行推理运算output = sess.run(None, {"input": img_numpy})output = softmax(output[0])print(output)ind = np.argmax(output, axis=1)print(ind)lablename = "飞机、汽车、鸟类、猫、鹿、狗、青蛙、马、船、卡车".split("、")res = {"code": 200, "msg": "处理成功", "url": img, "class": lablename[ind[0]]}