0
点赞
收藏
分享

微信扫一扫

2024年最新版 在AlmaLinux上安装MongoDB

认真的老去 2024-02-02 阅读 16
图像分类模型部署: PyTorch -> ONNX

1. 模型部署介绍

1.1 人工智能开发部署全流程

1.2 模型部署平台和芯片介绍

  • 设备:PC、浏览器、APP、小程序、服务器、嵌入式开发板、无人车、无人机、Jetson Nano、树莓派、机械臂、物联网设备
  • 厂商
    • 英特尔(Intel):主要生产 CPU(中央处理器)和一些 FPGA(现场可编程门阵列)芯片。代表作品包括 Intel Core 系列 CPU 和 Xeon 系列服务器 CPU,以及 FPGA 产品如 Intel Stratix 系列。
    • 英伟达(NVIDIA):以 GPU(图形处理器)为主打产品,广泛应用于图形渲染、深度学习等领域。代表作品包括 NVIDIA GeForce 系列用于游戏图形处理,NVIDIA Tesla 和 NVIDIA A100 用于深度学习加速。
    • AMD:主要生产 CPU 和 GPU。代表作品包括 AMD Ryzen 系列 CPU 和 AMD EPYC 系列服务器 CPU,以及 AMD Radeon 系列 GPU 用于游戏和专业图形处理。
    • 苹果(Apple):生产自家设计的芯片,主要包括苹果 M 系列芯片。代表作品有 M1 芯片,广泛应用于苹果的 Mac 电脑、iPad 和一些其他设备。
    • 高通(Qualcomm):主要生产移动平台芯片,包括移动处理器和调制解调器。代表作品包括 Snapdragon 系列芯片,用于智能手机和移动设备。
    • 昇腾(Ascend):由华为生产,主要生产 NPU(神经网络处理器),用于深度学习任务。代表作品包括昇腾 910 和昇腾 310。
    • 麒麟(Kirin):同样由华为生产,主要生产手机芯片,包括 CPU 和 GPU。代表作品包括麒麟 9000 系列,用于华为旗舰手机。
    • 瑞芯微(Rockchip):主要生产 VPU(视觉处理器)和一些移动平台芯片。代表作品包括 RK3288 和 RK3399,广泛应用于智能显示、机器人等领域。
芯片名英文名中文名厂商主要任务是否训练是否推理算力速度
CPUCentral Processing Unit(CPU)中央处理器各大厂商通用计算中等
GPUGraphics Processing Unit(GPU)图形处理器NVIDIA、AMD等图形渲染、深度学习加速
TPUTensor Processing Unit(TPU)张量处理器谷歌机器学习中的张量运算
NPUNeural Processing Unit(NPU)神经网络处理器华为、联发科等深度学习模型的性能提升中等
VPUVision Processing Unit(VPU)视觉处理器英特尔、博通等图像和视频处理中等中等
DSPDigital Signal Processor(DSP)数字信号处理器德州仪器、高通等数字信号处理、音频信号处理中等中等
FPGAField-Programmable Gate Array(FPGA)现场可编程门阵列英特尔、赛灵思等可编程硬件加速器中等

1.3 模型部署的通用流程

2. 使用 ONNX 的意义

从这两张图可以很明显的看到,当有了中间表示 ONNX 后,从原来的 M × N M \times N M×N 变为了 M + N M + N M+N,让模型部署的流程变得简单。

3. ONNX 的介绍

开源机器学习通用中间格式,由微软、Facebook(Meta)、亚马逊、IBM 共同发起的。它可以兼容各种深度学习框架,也可以兼容各种推理引擎和终端硬件、操作系统

4. ONNX 环境安装

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

5. PyTorch → ONNX

5.1 将一个分类模型转换为 ONNX

import torch
from torchvision import models


device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f"正在使用的设备: {device}")

# 创建一个训练好的模型
model = models.resnet18(pretrained=True)  # ImageNet 预训练权重
model = model.eval().to(device)

# 构建一个输入
dummy_input = torch.randn(size=[1, 3, 256, 256]).to(device)  # [N, B, H, W]

# 让模型推理
output = model(dummy_input)
print(f"output.shape: {output.shape}")

# 使用 PyTorch 自带的函数将模型转换为 ONNX 格式
onnx_save_path = 'ONNX/saves/resnet18_imagenet.onnx'  # 导出的ONNX模型路径 
with torch.no_grad():
    torch.onnx.export(
        model=model,                            # 要转换的模型
        args=dummy_input,                       # 模型的输入
        f=onnx_save_path,                       # 导出的ONNX模型路径 
        input_names=['input'],                  # ONNX模型输入的名字(自定义)
        output_names=['output'],                # ONNX模型输出的名字(自定义)
        opset_version=11,                       # Opset算子集合的版本(默认为17)
    )
    
print(f"ONNX 模型导出成功,路径为:{onnx_save_path}")
正在使用的设备: cpu
/home/leovin/anaconda3/envs/wsl/lib/python3.8/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.
  warnings.warn(
/home/leovin/anaconda3/envs/wsl/lib/python3.8/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=ResNet18_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet18_Weights.DEFAULT` to get the most up-to-date weights.
  warnings.warn(msg)
Downloading: "https://download.pytorch.org/models/resnet18-f37072fd.pth" to /home/leovin/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 44.7M/44.7M [00:03<00:00, 13.9MB/s]
output.shape: torch.Size([1, 1000])
ONNX 模型导出成功,路径为:ONNX/saves/resnet18_imagenet.onnx

💡 Tips:

  1. opset 算子集不同版本区别: Operators.md
  2. 虽然说 PyTorch 在提醒 pretrained=True 将会被弃用,可以使用 weights=weights=ResNet18_Weights.DEFAULTweights=ResNet18_Weights.IMAGENET1K_V1 来代替。但很明显前者比较方便,后者还需要查看对应的版本号,比较麻烦 😂

接下来我们使用 Netron 查看一下这个模型:

5.2 检查一个模型导出是否正确

import onnx


# 读取导出的模型
onnx_path = 'ONNX/saves/resnet18_imagenet.onnx'  # 导出的ONNX模型路径
onnx_model = onnx.load(onnx_path)

# 检查模型是否正常
onnx.checker.check_model(onnx_model)

print(f"模型导出正常!")
模型导出正常!

5.3 修改动态维度

前面我们导出的 ONNX 模型中,输入的维度是固定的:[1, 3, 256, 256],那么此时这个 ONNX 的输入就被限制了:

  • 如果我们想要多 Batch 的输入 → 不行
  • 如果我们输入的图片是灰度图 → 不行
  • 如果我们输入的图片尺寸不是 256×256 → 不行

torch.onnx.export() 这个函数也帮我解决了这个问题,它有一个名为 dynamic_axis 的参数,我们看一下官网对该参数的描述:

下面我们用一下这个参数:

import torch
from torchvision import models
import onnx


device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f"正在使用的设备: {device}")

# 创建一个训练好的模型
model = models.resnet18(pretrained=True)  # ImageNet 预训练权重
model = model.eval().to(device)

# 构建一个输入
dummy_input = torch.randn(size=[1, 3, 256, 256]).to(device)  # [N, B, H, W]

# 让模型推理
output = model(dummy_input)
print(f"output.shape: {output.shape}\n")

# ------ 使用 PyTorch 自带的函数将模型转换为 ONNX 格式
onnx_save_path = 'ONNX/saves/resnet18_imagenet-with_dynamic_axis.onnx'  # 导出的ONNX模型路径 
with torch.no_grad():
    torch.onnx.export(
        model=model,                            # 要转换的模型
        args=dummy_input,                       # 模型的输入
        f=onnx_save_path,                       # 导出的ONNX模型路径 
        input_names=['input'],                  # ONNX模型输入的名字(自定义)
        output_names=['output'],                # ONNX模型输出的名字(自定义)
        opset_version=11,                       # Opset算子集合的版本(默认为17)
        dynamic_axes={                          # 修改某一个维度为动态
            'input': {0: 'B', 2: 'H', 3: 'W'}   # 将原本的 [1, 3, 256, 256] 修改为 [B, 3, H, W]
        }
    )
    
print(f"ONNX 模型导出成功,路径为:{onnx_save_path}\n")

# ------ 验证导出的模型是否正确
# 读取导出的模型
onnx_model = onnx.load(onnx_save_path)

# 检查模型是否正常
onnx.checker.check_model(onnx_model)

print(f"模型导出正常!")
正在使用的设备: cpu
output.shape: torch.Size([1, 1000])

ONNX 模型导出成功,路径为:ONNX/saves/resnet18_imagenet-with_dynamic_axis.onnx

模型导出正常!

此时我们再用 Netron 看一下这个模型:

可以看到,输入的 Batch、Height、Width 均变为了动态维度,即只有当模型运行的时候才知道输入的这三个维度具体的值

6. ONNX Runtime 部署:推理单张图片

import os
import random
import numpy as np
from PIL import Image
import onnxruntime
from torchvision import transforms
import torch
import torch.nn.functional as F
import pandas as pd


# ==================================== 加载 ONNX 模型,创建推理会话 ==================================== 
ort_session = onnxruntime.InferenceSession(path_or_bytes='ONNX/saves/resnet18_imagenet-fix_axis.onnx')  # ort -> onnxruntime

# ==================================== 模型冷启动 ==================================== 
dummy_input = np.random.randn(1, 3, 256, 256).astype(np.float32)
ort_inputs = {'input': dummy_input}
ort_output = ort_session.run(output_names=['output'], input_feed=ort_inputs)[0]  # 输出被[]包围了,所以需要取出来
print(f"模型冷启动完毕! 其推理结果的shape为: {ort_output.shape}")

# ==================================== 加载真正的图像 ==================================== 
images_folder = 'Datasets/Web/images'
images_list = [os.path.join(images_folder, img) for img in os.listdir(images_folder) if img.lower().endswith(('.jpg', '.png', '.webp'))]

img_path = images_list[random.randint(0, len(images_list)-1)]
img = Image.open(fp=img_path)

# ==================================== 图像预处理 ==================================== 
# 定义预处理函数
img_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(256),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],  # imagenet专用
        std=[0.229, 0.224, 0.225]),  # imagenet专用
])

# 对图片进行预处理
input_img = img_transform(img)
print(f"input_img.type: {type(input_img)}")
print(f"input_img.shape: {input_img.shape}")

# 为图片添加batch维度
input_img = torch.unsqueeze(input_img, dim=0)

# ==================================== ONNX模型推理 ==================================== 
# 因为ONNXRuntime需要的是numpy而非torch的tensor, 所以将其转换为numpy
input_img = input_img.numpy()
print(f"input_img.type: {type(input_img)}")
print(f"input_img.shape: {input_img.shape}")

# 模型推理图片
ort_inputs = {'input': input_img, }
ort_results = ort_session.run(output_names=['output'], input_feed=ort_inputs)[0]  # 得到 1000 个类别的分数
print(f"模型推理完毕! 此时结果的shape为:{ort_results.shape}")

# ==================================== 后处理 ==================================== 
# 使用 softmax 函数将分数转换为概率
ort_results_softmax = F.softmax(input=torch.from_numpy(ort_results), dim=1)
print(f"经过softmax后的输出的shape为:{ort_results_softmax.shape}")

# 取概率最大的前 n 个结果
n = 3
top_n = torch.topk(input=ort_results_softmax, k=n)

probs = top_n.values.numpy()[0]
indices = top_n.indices.numpy()[0]

print(f"置信度最高的前{n}个结果为:\t{probs}\n"
      f"对应的类别索引为:\t\t{indices}")

# ==================================== 显示类别 ==================================== 
df = pd.read_csv('Datasets/imagenet_classes_indices.csv')

idx2labels = {}
for idx, row in df.iterrows():
    # idx2labels[row['ID']] = row['class']  # 英文标签
    idx2labels[row['ID']] = row['Chinese']  # 中文标签

print(f"=============== 推理结果 ===============\n"
      f"图片路径: {img_path}")
for i, (class_prob, idx) in enumerate(zip(probs, indices)):
    class_name = idx2labels[idx]
    text = f"\tNo.{i}: {class_name:<30} --> {class_prob:>.4f}"
    print(text)
模型冷启动完毕! 其推理结果的shape为: (1, 1000)
input_img.type: <class 'torch.Tensor'>
input_img.shape: torch.Size([3, 256, 256])
input_img.type: <class 'numpy.ndarray'>
input_img.shape: (1, 3, 256, 256)
模型推理完毕! 此时结果的shape为:(1, 1000)
经过softmax后的输出的shape为:torch.Size([1, 1000])
置信度最高的前3个结果为:       [9.9472505e-01 7.4335985e-04 5.2123831e-04]
对应的类别索引为:              [673 662 487]
=============== 推理结果 ===============
图片路径: Datasets/Web/images/mouse.jpg
        No.0: 鼠标,电脑鼠标                        --> 0.9947
        No.1: 调制解调器                          --> 0.0007
        No.2: 移动电话,手机                        --> 0.0005

7. ONNX Runtime 和 PyTorch 速度对比

  1. 不同尺度下单张图片推理 --> 对比代码链接
  2. 不同尺度下多张图片推理 --> 对比代码链接

实验环境

  • CPU:Intel i5-10400F @ 2.90 GHz
  • Memory: 8 x 2 = 16GB
  • Disk: SSD
  • GPU: RTX 3070 O8G
  • OS: Windows 10 (WSL)
  • Device: CPU
  • 模型推理次数: 50

7.1 ResNet-18

实验结果

Input ShapeONNX(fix)ONNX(fix+sim)ONNX(dyn)ONNX(dyn+sim)PyTorch(CPU)PyTorch(GPU)
[1, 3, 32, 32]0.0577s0.0597s0.0592s0.0585s0.0688s0.0787s
[1, 3, 64, 64]0.0605s0.0593s0.0588s0.0621s0.0700s0.0723s
[1, 3, 128, 128]0.0705s0.0686s0.0699s0.0694s0.0762s0.0760s
[1, 3, 256, 256]0.0784s0.0811s0.0797s0.0789s0.0949s0.0813s
[1, 3, 512, 512]0.1249s0.1241s0.1251s0.1256s0.1686s0.0996s
[1, 3, 640, 640]0.1569s0.1525s0.1572s0.1579s0.2242s0.0863s
[1, 3, 768, 768]0.1986s0.1946s0.1985s0.2038s0.2933s0.0956s
[1, 3, 1024, 1024]0.2954s0.2957s0.3094s0.3045s0.4871s0.1047s
[16, 3, 32, 32]0.2540s0.2545s0.2558s0.2498s0.2570s0.2473s
[16, 3, 64, 64]0.2811s0.2745s0.2696s0.2655s0.2824s0.2553s
[16, 3, 128, 128]0.3595s0.3181s0.3143s0.3544s0.3817s0.3518s
[16, 3, 256, 256]0.7315s0.7112s0.6767s0.6122s0.7169s0.3469s
[16, 3, 512, 512]1.3042s1.2586s1.1813s1.1949s1.6609s0.4270s
[16, 3, 640, 640]1.6340s1.6429s1.6659s1.6693s2.3923s0.5292s
[16, 3, 768, 768]2.2843s2.2830s2.3325s2.3303s3.9278s1.7851s
[16, 3, 1024, 1024]3.9132s3.9742s3.9668s3.9104s6.7532s3.6507s

画图结果

结论

  • 单 Batch
    • 静态维度和动态维度相差不大
    • 当图片尺寸在 [32, 32] ~ [256, 256] 之间时,ONNX 速度比 PyTorch-GPU 速度要快;当图片尺寸大于 [256, 256] 时,PyTorch-GPU 拥有绝对的优势
    • 当图片尺寸小于 [64, 64] 时,PyTorch-CPU 速度快于 PyTorch-GPU;当图片尺寸大于 [64, 64] 时,PyTorch-GPU 速度快于 PyTorch-CPU
    • 无论在什么时候,ONNX 速度均快于 PyTorch-CPU
  • 多 Batch
    • 静态维度和动态维度相差不大
    • 当图片尺寸小于 [128, 128] 时,ONNX、PyTorch-CPU、PyTorch-GPU 三者很难有区别(实际上 PyTorch-GPU 速度要慢一些,因为要将模型和输入放到 GPU 中,这部分会划分几秒钟的时间)
    • 当图片尺寸大于 [128, 128] 时,GPU 逐渐扩大优势(由于 OOM 的原因,[18, 3, 1024, 1024] 下 PyTorch-GPU 只推理了一次,因此速度被拉平了很多。在显存足够充裕的情况下,PyTorch-GPU 的速度是碾压其他方法的)
    • 当图片尺寸大于 [256, 256] 时,PyTorch-CPU 的速度远远慢于 ONNX
  • Sim 前后
    • 可以发现,在使用 python -m onnxsim 前后差距不大
  • 总结
    • 在使用 CPU 进行推理时,建议使用 ONNX 进行,因为不光速度有优势,而且对内存的占用也比 PyTorch-CPU 要小的多
    • 在进行多 Batch 推理时,如果有 GPU 还是使用 PyTorch-GPU,这样会缩减大量的时间(⚠️ GPU 在加载模型和输入时可能会比较耗时)
    • ⚠️ 在使用 python -m onnxsim 前后差距不大

7.2 MobileNetV3-Small

接下来我们在 MobileNetV3-Small 上也进行相同的实验。

Input ShapeONNX(fix)ONNX(dyn)PyTorch(CPU)PyTorch(GPU)
[1, 3, 32, 32]0.0575s0.0619s0.0636s0.0731s
[1, 3, 64, 64]0.0585s0.0591s0.0643s0.0701s
[1, 3, 128, 128]0.0611s0.0597s0.0629s0.0700s
[1, 3, 256, 256]0.0627s0.0622s0.0690s0.0731s
[1, 3, 512, 512]0.0714s0.0703s0.0841s0.0765s
[1, 3, 640, 640]0.0776s0.0785s0.0975s0.0823s
[1, 3, 768, 768]0.0867s0.0861s0.1138s0.0851s
[1, 3, 1024, 1024]0.1103s0.1126s0.1630s0.0958s
[16, 3, 32, 32]0.2410s0.2295s0.2538s0.2446s
[16, 3, 64, 64]0.2443s0.2421s0.2576s0.2481s
[16, 3, 128, 128]0.2618s0.2576s0.2804s0.2727s
[16, 3, 256, 256]0.3097s0.3131s0.3502s0.3043s
[16, 3, 512, 512]0.5556s0.5873s0.7655s0.3970s
[16, 3, 640, 640]0.7191s0.7130s0.8988s0.4877s
[16, 3, 768, 768]0.9293s0.9285s1.5091s0.5754s
[16, 3, 1024, 1024]1.4768s1.4945s3.3530s1.1316s

画图结果

其实可以发现,与 ResNet18 的结论是一致的。

7.3 为什么 python -m onnxsim 没有效果

我们看一下这个过程:

-------------- ResNet-18 --------------

python -m onnxsim ONNX/saves/resnet18-dynamic_dims.onnx ONNX/saves/resnet18-dynamic_dims-sim.onnx
Simplifying...
Finish! Here is the difference:
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┓
┃                   ┃ Original Model ┃ Simplified Model ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━┩
│ Add               │ 8              │ 8                │
│ Constant          │ 42             │ 42               │
│ Conv              │ 20             │ 20               │
│ Flatten           │ 1              │ 1                │
│ Gemm              │ 1              │ 1                │
│ GlobalAveragePool │ 1              │ 1                │
│ MaxPool           │ 1              │ 1                │
│ Relu              │ 17             │ 17               │
│ Model Size        │ 44.6MiB        │ 44.6MiB          │
└───────────────────┴────────────────┴──────────────────┘

-------------- MobileNetV3-Small --------------

python -m onnxsim ONNX/saves/mobilenetv3small-dynamic_dims.onnx ONNX/saves/mobilenetv3small-dynamic_dims-sim.onnx
Simplifying...
Finish! Here is the difference:
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┓
┃                   ┃ Original Model ┃ Simplified Model ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━┩
│ Add               │ 6              │ 6                │
│ Constant          │ 108            │ 108              │
│ Conv              │ 52             │ 52               │
│ Flatten           │ 1              │ 1                │
│ Gemm              │ 2              │ 2                │
│ GlobalAveragePool │ 10             │ 10               │
│ HardSigmoid       │ 9              │ 9                │
│ HardSwish         │ 19             │ 19               │
│ Mul               │ 9              │ 9                │
│ Relu              │ 14             │ 14               │
│ Model Size        │ 9.7MiB         │ 9.7MiB           │
└───────────────────┴────────────────┴──────────────────┘

可以看到,其实根本没有变化,所以速度也没有提升。

8. ONNX 与 PyTorch 精度对比

我们现在有如下的模型:

  • weights.pth: PyTorch 权重
  • weights.onnx: ONNX 权重
  • weights-sim.onnx: ONNX 精简后的权重

模型的关系如下:

现在我们想要搞清楚,这样转换后的模型精度是怎么样的?

import os
import argparse
import numpy as np
import pandas as pd
from PIL import Image
import onnxruntime
import torch
import torch.nn.functional as F
from torchvision import transforms, models
from rich.progress import track


# ==================================== 参数 ==================================== 
parser = argparse.ArgumentParser()
parser.add_argument('--image_folder_path', type=str, default='Datasets/Web/images', help='图片路径')
parser.add_argument('--input-shape', type=int, nargs=2, default=[256, 256])
parser.add_argument('--verbose', action='store_true', help='')
args = parser.parse_args()  # 解析命令行参数

onnx_weights = 'ONNX/saves/model-dynamic_dims.onnx'
onnx_weights_sim = 'ONNX/saves/model-dynamic_dims-sim.onnx'
# ==============================================================================

# 定义模型
onnx_model = onnxruntime.InferenceSession(path_or_bytes=onnx_weights)
onnx_model_sim = onnxruntime.InferenceSession(path_or_bytes=onnx_weights_sim)
pytorch_model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1).eval()  # ⚠️ 一定要 .eval
# pytorch_model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)

# 定义预处理函数
img_transform = transforms.Compose([
    transforms.Resize(args.input_shape[-1]),
    transforms.CenterCrop(args.input_shape[-1]),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],  # imagenet专用
        std=[0.229, 0.224, 0.225]),  # imagenet专用
])

image_list = [os.path.join(args.image_folder_path, img) for img in os.listdir(args.image_folder_path) \
               if img.lower().endswith(('.jpg', '.jpeg', '.png', '.webp'))]

for img_idx, image_path in track(enumerate(image_list), description='Precision Comparison'):
    # 读取图片
    img = Image.open(fp=image_path)  # 读取图片
    input_img = img_transform(img)
    input_img = input_img.unsqueeze(0)
    print(f"inputs.type: {type(input_img)}") if args.verbose else ...
    print(f"inputs.shape: {input_img.shape}") if args.verbose else ...

    model_ls = ['pt', 'onnx', 'onnx-sim']
    for model_name in model_ls:
        if model_name != 'pt':
            if not isinstance(input_img, np.ndarray):
                input_img = input_img.numpy()
            model_input = {'input': input_img, }
            model_result = onnx_model.run(output_names=['output'], input_feed=model_input)[0]
        else:
            model_result = pytorch_model(input_img)
        
        if not isinstance(model_result, torch.Tensor):
            model_result = torch.from_numpy(model_result)
        
        model_result_softmax = F.softmax(input=model_result, dim=1)  # [1, 1000]

        # 取概率最大的前 n 个结果
        n = 3
        top_n = torch.topk(input=model_result_softmax, k=n, dim=1)

        probs = top_n.values.detach().numpy()[0]  # torch.Size([18, 3])
        indices = top_n.indices.detach().numpy()[0]  # torch.Size([18, 3])
        print(f"probs: {probs}") if args.verbose else ...
        print(f"indices: {indices}") if args.verbose else ...

        df = pd.read_csv('Datasets/imagenet_classes_indices.csv')

        idx2labels = {}
        for _, row in df.iterrows():
            idx2labels[row['ID']] = row['Chinese']  # 中文标签

        print(f"============================== 推理结果-{model_name} ==============================")  if args.verbose else ...
        
        _results = []
        for i, (prob, idx) in enumerate(zip(probs, indices)):
            class_name = idx2labels[idx]
            text = f"No.{i}: {class_name:<30} --> {prob:>.5f}"  if args.verbose else ...
            _results.append(prob)
            print(text)
        print(f"=====================================================================")  if args.verbose else ...

        with open("ONNX/saves/Precision-comparison.txt", 'a') as f:
            if model_name == 'pt':
                f.write(f"|[{img_idx+1}] {os.path.basename(image_path)}"
                        f"|{_results[0]:>.5f}</br>{_results[1]:>.5f}</br>{_results[2]:>.5f}")
            elif model_name == 'onnx':
                f.write(f"|{_results[0]:>.5f}</br>{_results[1]:>.5f}</br>{_results[2]:>.5f}")
            else:
                f.write(f"|{_results[0]:>.5f}</br>{_results[1]:>.5f}</br>{_results[2]:>.5f}|\n")

实验结果

图片名称PyTorchONNXONNX-sim
[1] book.jpg0.739730.050490.023580.739730.050490.023580.739730.050490.02358
[2] butterfly.jpg0.897040.047720.015420.897040.047720.015420.897040.047720.01542
[3] camera.jpg0.276580.177090.109250.276580.177090.109250.276580.177090.10925
[4] cat.jpg0.277730.183930.172540.277730.183930.172540.277730.183930.17254
[5] dog.jpg0.517870.253840.059290.517870.253840.059290.517870.253840.05929
[6] dogs_orange.jpg0.352890.301140.077910.352890.301140.077910.352890.301140.07791
[7] female.jpg0.156000.080310.048080.156000.080310.048080.156000.080310.04808
[8] free-images.jpg0.455950.176260.084140.455950.176260.084140.455950.176260.08414
[9] gull.jpg0.647110.233240.044300.647110.233240.044300.647110.233240.04430
[10] laptop-phone.jpg0.493790.354050.060630.493790.354050.060630.493790.354050.06063
[11] monitor.jpg0.516780.441930.022320.516780.441930.022320.516780.441930.02232
[12] motorcycle.jpg0.317120.224350.156310.317120.224350.156310.317120.224350.15631
[13] mouse.jpg0.994730.000740.000520.994730.000740.000520.994730.000740.00052
[14] panda.jpg0.945590.031990.005610.945590.031990.005610.945590.031990.00561
[15] share_flower_fullsize.jpg0.788060.056910.024830.788060.056910.024830.788060.056910.02483
[16] tiger.jpeg0.617490.380010.000520.617490.380010.000520.617490.380010.00052

可以看到,转换前后模型并没有精度的丢失。

9. 〔拓展知识〕为什么 .pt 模型在推理时一定要 .eval()

在PyTorch中,.eval() 是一个用于将模型切换到评估(inference)模式的方法。在评估模式下,模型的行为会有所变化,主要体现在两个方面:DropoutBatch Normalization

  1. Dropout:

    • 在训练阶段,为了防止过拟合,通常会使用 dropout 策略,即在每个训练步骤中,以一定的概率随机丢弃某些神经元的输出。
    • 在推理阶段,我们希望获得模型的确定性输出,而不是在每次推理时都丢弃不同的神经元。因此,在推理时应该关闭 dropout。通过调用 .eval(),PyTorch 会将所有 dropout 层设置为评估模式,即不进行随机丢弃。
  2. Batch Normalization:

    • Batch Normalization(批标准化)在训练时通过对每个 mini-batch 进行标准化来加速训练,但在推理时,我们通常不是基于 mini-batch 进行预测,因此需要使用整个数据集的统计信息进行标准化。
    • .eval() 模式下,Batch Normalization 会使用训练时计算的移动平均和方差,而不是使用当前 mini-batch 的统计信息。

因此,为了确保在推理时得到一致和可靠的结果,需要在推理之前调用 .eval() 方法,以确保模型处于评估模式,关闭了 dropout,并使用适当的 Batch Normalization 统计信息。


举个例子,对于一张猫咪图片而言,如果我们的 .pt 模型没有开启 .eval() 就进行推理,那么得到的结果如下:

============================== 推理结果-pt ==========================
No.0: 桶                                --> 0.00780
No.1: 手压皮碗泵                        --> 0.00680
No.2: 钩爪                              --> 0.00601
====================================================================
probs: [0.27773306 0.18392678 0.17254312]
indices: [281 285 287]
============================== 推理结果-onnx ========================
No.0: 虎斑猫                            --> 0.27773
No.1: 埃及猫                            --> 0.18393
No.2: 猞猁,山猫                         --> 0.17254
====================================================================
probs: [0.27773306 0.18392678 0.17254312]
indices: [281 285 287]
============================== 推理结果-onnx-sim ====================
No.0: 虎斑猫                            --> 0.27773
No.1: 埃及猫                            --> 0.18393
No.2: 猞猁,山猫                         --> 0.17254
====================================================================

可以看到,对于 ONNX 模型而言,推理相对来说是比较正确的。但对于 PyTorch 模型,推理与猫无关了,所以 ⚠️ 在推理时开启 .eval() 是非常重要的!

参考

  1. 图像分类模型部署-Pytorch转ONNX
  2. Pytorch图像分类模型部署-ONNX Runtime本地终端推理
举报

相关推荐

0 条评论