安裝配置環(huán)境
代碼運(yùn)行云GPU平臺(tái)：公眾號(hào) 人工智能小技巧 回復(fù) gpu
同濟(jì)子豪兄 2022-8-22 2023-4-28 2023-5-8
安裝 Pytorch

pip3 install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113

安裝 ONNX

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

安裝推理引擎 ONNX Runtime

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

驗(yàn)證安裝配置成功

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

Pytorch圖像分類模型轉(zhuǎn)ONNX-ImageNet1000類
把Pytorch預(yù)訓(xùn)練ImageNet圖像分類模型，導(dǎo)出為ONNX格式，用于后續(xù)在推理引擎上部署。
代碼運(yùn)行云GPU平臺(tái)：公眾號(hào) 人工智能小技巧 回復(fù) gpu
同濟(jì)子豪兄 2022-8-22 2023-4-28 2023-5-8

import torch
from torchvision import models
import onnx

device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:',device)

# 載入ImageNet預(yù)訓(xùn)練PyTorch圖像分類模型
model=models.resnet18(pretrained=True)
model=model.eval().to(device)
# 構(gòu)造一個(gè)輸入圖像Tensor
x=torch.rand(1,3,256,256).to(device)
print(x.shape)
# 輸入PyTorch模型推理預(yù)測(cè)，獲得1000個(gè)類別預(yù)測(cè)結(jié)果
output=model(x)
print(output.shape)
# PyTorch模型轉(zhuǎn)ONNX格式
with torch.no_grad():
    torch.onnx.export(
        model,                    # 要轉(zhuǎn)換的模型
        x,                        # 模型的任意一組輸入
        'resnet18_imagenet.onnx', # 導(dǎo)出的ONNX文件名
        input_names=['input'],    # 輸入Tensor的名稱（自己起名字）
        output_names=['output'],  # 輸出Tensor的名稱（自己起名字）
        opset_version=11,         # ONNX算子集版本
    )
# 驗(yàn)證onnx模型導(dǎo)出成功
# 讀取ONNX模型
onnx_model=onnx.load('resnet18_imagenet.onnx')
# 檢查模型格式是否正確
onnx.checker.check_model(onnx_model)
print('無報(bào)錯(cuò)，onnx模型導(dǎo)入成功!')
# 以可讀的形式打印計(jì)算圖
print(onnx.helper.printable_graph(onnx_model.graph))
#使用Netron可視化模型結(jié)構(gòu)
#Netron：https://netron.app
#視頻教程：https://www.bilibili.com/video/BV1TV4y1P7AP

PyTorch圖像分類模型轉(zhuǎn)ONNX-花分類

import torch
from torchvision import models
import onnx

device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:',device)

# 導(dǎo)入訓(xùn)練好的模型
model=torch.load('MobileNetV2.pth')
model=model.eval().to(device)
# 構(gòu)造一個(gè)輸入圖像Tensor
x=torch.rand(1,3,224,224).to(device)
print(x.shape)
# 輸入PyTorch模型推理預(yù)測(cè)，獲得N個(gè)類別預(yù)測(cè)結(jié)果
output=model(x)
print(output.shape)
# PyTorch模型轉(zhuǎn)ONNX格式
with torch.no_grad():
    torch.onnx.export(
        model,                    # 要轉(zhuǎn)換的模型
        x,                        # 模型的任意一組輸入
        'mobilenetv2_flowers5.onnx', # 導(dǎo)出的ONNX文件名
        input_names=['input'],    # 輸入Tensor的名稱（自己起名字）
        output_names=['output'],  # 輸出Tensor的名稱（自己起名字）
        opset_version=11,         # ONNX算子集版本
    )
# 驗(yàn)證onnx模型導(dǎo)出成功
# 讀取ONNX模型
onnx_model=onnx.load('mobilenetv2_flowers5.onnx')
# 檢查模型格式是否正確
onnx.checker.check_model(onnx_model)
print('無報(bào)錯(cuò)，onnx模型導(dǎo)入成功!')
# 以可讀的形式打印計(jì)算圖
print(onnx.helper.printable_graph(onnx_model.graph))
# 登錄netron.app在線模型可視化網(wǎng)頁(yè)打開可視化模型結(jié)構(gòu)

擴(kuò)展閱讀
【同濟(jì)子豪兄】?jī)商旄愣▓D像分類、目標(biāo)檢測(cè)、語(yǔ)義分割、實(shí)例分割、關(guān)鍵點(diǎn)檢測(cè)畢業(yè)設(shè)計(jì)：https://github.com/TommyZihao/Train_Custom_Dataset
ONNX主頁(yè)： https://onnx.ai
ONNX Github主頁(yè)： https://github.com/onnx/onnx
ONNX支持的深度學(xué)習(xí)框架：https://onnx.ai/supported-tools.html
ONNX支持的計(jì)算平臺(tái)：https://onnx.ai/supported-tools.html
ONNX算子：https://onnx.ai/onnx/operators
模型在線可視化網(wǎng)站 Netron：https://netron.app
Netron視頻教程：https://www.bilibili.com/video/BV1TV4y1P7AP
周奕帆博客
模型部署入門教程（一）：模型部署簡(jiǎn)介 https://zhuanlan.zhihu.com/p/477743341
模型部署入門教程（二）：解決模型部署中的難題 https://zhuanlan.zhihu.com/p/479290520
模型部署入門教程（三）：PyTorch 轉(zhuǎn) ONNX 詳解 https://zhuanlan.zhihu.com/p/498425043文章來源地址http://www.zghlxwxcb.cn/news/detail-708258.html

到了這里，關(guān)于Pytorch圖像分類模型轉(zhuǎn)ONNX（同濟(jì)子豪兄學(xué)習(xí)筆記）的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！