使用onnx部署自己的神经网络

使用onnx部署自己的神经网络

一、安装环境

安装Pytorch：详见我的博客基于vscode安装pytorch

安装onnx：

pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple

安装onnxruntime:

pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple

安装可能会用到的包：

pip install numpy pandas matplotlib tqdm opencv-python pillow -i https://pypi.tuna.tsinghua.edu.cn/simple

测试是否安装成功

import torch
print('PyTorch 版本', torch.__version__)
import onnx
print('ONNX 版本', onnx.__version__)
import onnxruntime as ort
print('ONNX Runtime 版本', ort.__version__)

有下面类似的输出就没问题了

二、使用onnx来将pth转换为onnx

我这里直接使用torchvision来加载模型来演示

import torch
from torchvision import models
import onnx

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device', device)
#这里是选择cpu和gpu
model = models.resnet18(pretrained=True)
model = model.eval().to(device)
#将模型的验证加载到device上，注意这里的model是带参数的

#设置一个随机输入
x = torch.randn(1, 3, 256, 256).to(device)
#开始转换
with torch.no_grad():
    torch.onnx.export(
        model,                       # 要转换的模型
        x,                           # 模型的任意一组输入
        'resnet18_imagenet.onnx',    # 导出的 ONNX 文件名
        opset_version=11,            # ONNX 算子集版本
        input_names=['input'],       # 输入 Tensor 的名称（自己起名字）
        output_names=['output']      # 输出 Tensor 的名称（自己起名字）
    )

# 读取 ONNX 模型
onnx_model = onnx.load('resnet18_imagenet.onnx')
# 检查模型格式是否正确
if not onnx.checker.check_model(onnx_model)
	print('onnx模型成功加载！')
#打印计算图，可视化
print(onnx.helper.printable_graph(onnx_model.graph))

三、加载onnx文件并进行PC本地部署

这里演示对一张图片进行加载

import onnxruntime
import numpy as np
import torch
import torch.nn.functional as F
import pandas as pd
#加载onnx模型并获取onnxruntime推理器，这里我使用cpu来进行推理，需要在providers中指定
ort_session = onnxruntime.InferenceSession('resnet18_imagenet.onnx'，providers=['CPUExecutionProvider'])

#跟imagenet设置一样的图片处理函数
test_transform = transforms.Compose([transforms.Resize(256),
                                     transforms.CenterCrop(256),
                                     transforms.ToTensor(),
                                     transforms.Normalize(
                                         mean=[0.485, 0.456, 0.406], 
                                         std=[0.229, 0.224, 0.225])])
读取图片
img_path = '你自己的图片.jpg'
from PIL import Image
img_pil = Image.open(img_path)    
input_img = test_transform(img_pil)
#拓展维度，因为原模型输入的第一维度是batch，我们这里只有一张图片，所以需要增加一个维度
input_tensor = input_img.unsqueeze(0).numpy()
print("输入的tensor维度是：",input_tensor.shape)
# ONNX Runtime 输入
ort_inputs = {'input': input_tensor}
# ONNX Runtime 输出
pred_logits = ort_session.run(['output'], ort_inputs)[0]
#这里的['output']是希望网络输出上面命名为['output']的输出,是上面你自己起的名字
#预测结果
pred_logits = torch.tensor(pred_logits)
#查看topk结果
k = 5
top_n = torch.topk(pred_softmax, k)
# 预测类别
pred_ids = top_n.indices.numpy()[0]
# 预测置信度
confs = top_n.values.numpy()[0]