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

八、模型移植

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 resnet18

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"
    )

    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_features
    model.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)
    # 导出onnx
    torch.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 np

img_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_img

lablename = "飞机、汽车、鸟类、猫、鹿、狗、青蛙、马、船和卡车".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]

    # 应用softmax
    probabilities = 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-04
  3.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 Image

img_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 使用 GPU
    providers = ["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]]}