Docker实现Pytorch网络模型转Onnx格式,多种方法(opencv、onnxruntime)调用Onnx

0. 环境准备

  • 1. 安装Docker:

    curl -fsSL https://get.docker.com | bash -s docker --mirror Aliyun
    或
    sudo apt-get install docker.io          # 上面由于curl问题  采用的该命令
    
  • 2. 添加库

    创建.sh文件

    touch script.sh
    

    编辑.sh文件

    gedit script.sh
    

    添加如下内容:

    curl -s -L https://nvidia.github.io/nvidia-container-runtime/gpgkey | \
      sudo apt-key add -
    distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
    curl -s -L https://nvidia.github.io/nvidia-container-runtime/$distribution/nvidia-container-runtime.list | \
      sudo tee /etc/apt/sources.list.d/nvidia-container-runtime.list
    sudo apt-get update
    

    运行.sh文件 sh/bash

    sh script.sh
    
  • 将当前用户加入到docker用户组:
    1、创建docker组(安装docker时会自动创建,一般无需重新创建)

    $ sudo groupadd docker
    

    2、将当前用户加入到docker用户组

    $ sudo gpasswd -a  ${USER} docker
    

    3 、重新启动docker

     sudo service docker restart
     #或者执行以下命令,无须重新登录
     newgrp docker    # 建议用这个
    
  • 3. 创建镜像并同时创建容器:

    nldy@xiaoxie-Z10PE-D8-WS:~$ docker pull python:3.6
    nldy@xiaoxie-Z10PE-D8-WS:~$ docker images
    REPOSITORY   TAG       IMAGE ID       CREATED      SIZE
    python       3.6       1498723101b5   9 days ago   902MB
    
  • 4.测试:

    nldy@xiaoxie-Z10PE-D8-WS:~$ docker run -it --gpus=all python:3.6 /bin/bash
    root@bda3e7f60958:/# nvidia-smi
    

    在这里插入图片描述

  • 深度学习环境的配置:
    https://blog.csdn.net/qq_44747572/article/details/120831629?spm=1001.2014.3001.5501

过程中出现的问题:

1. 数据准备

classification_classes_ILSVRC2012.txt: 下载链接
dog.jpg:
在这里插入图片描述
VGG16模型: https://download.pytorch.org/models/vgg16-397923af.pth

2. Torchvison模型推理(未使用GPU,以下同)

Opencv-Onnx.py

#-*- codeing = utf-8 -*-
#@Function: pytorch->onnx  opencv.dnn 调用 onnx   Onnxruntime 调用 Onnx
#@Time : 2021/10/14 12:34
#@Author : yx
#@File : Opencv-Onnx.py
#@Software : PyCharm

import torch
import torchvision
import cv2
import onnx
import numpy as np
import timm
import os
from PIL import Image
from torchvision import transforms
import onnxruntime
from onnxsim import simplify
import torchvision.models as models

print(torch.__version__)
print(cv2.__version__)
print(np.__version__)
print(onnx.__version__)
classes = None
class_file = './classification_classes_ILSVRC2012.txt'
with open(class_file, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')

def init_model(model_name):
    if model_name == 'vgg16':
        path = "./vgg16.pth"
        vgg16 = models.vgg16(pretrained=False).eval()
        vgg16.load_state_dict(torch.load(path))
    model = vgg16
    dummy = torch.randn(1, 3, 224, 224)
    
    return model, dummy

model, dummy = init_model('vgg16')

img_file = './dog.jpg'

################################torchvison模型推理#####################################
transform = transforms.Compose([  # [1]
    transforms.Resize(256),  # [2]
    transforms.CenterCrop(224),  # [3]
    transforms.ToTensor(),  # [4]
    transforms.Normalize(  # [5]
        mean=[0.485, 0.456, 0.406],  # [6]
        std=[0.229, 0.224, 0.225]  # [7]
    )])

img = Image.open(img_file)
img_t = transform(img)            # 输入给模型的图像数据要先进行转换
batch_t = torch.unsqueeze(img_t, 0)
tc_out = model(batch_t).detach().cpu().numpy()
# Get a class with a highest score.
tc_out = tc_out.flatten()
classId = np.argmax(tc_out)
confidence = tc_out[classId]
print(confidence)

label = '%s: %.4f' % (classes[classId] if classes else 'Class #%d' % classId, confidence)
print(label)
################################ torchvison模型推理 ######################################

运行程序:

xiaoyuanzi@xiaoyuanzi-virtual-machine:~$ docker run  -v $PWD/python/onnx:/usr/src/python/onnx  -w /usr/src/python/onnx new_python:3.6 python Opencv_Onnx.py
1.9.1+cu102
4.5.3
1.19.5
1.10.1
19.722366
French bulldog: 19.7224

3. Pytorch -> Onnx

################################ pytorch转Onnx ######################################

onnx_name = 'vgg16.onnx'
torch.onnx.export(model, dummy, './vgg16.onnx')

# 模型测试(可跳过)
print("----- 模型测试 -----")
# 可以跳过该步骤,一般不会有问题
# 检查输出
def check_onnx_output(filename, input_data, torch_output):
    session = onnxruntime.InferenceSession(filename)
    input_name = session.get_inputs()[0].name
    result = session.run([], {input_name: input_data.detach().cpu().numpy()})
    for test_result, gold_result in zip(result, torch_output.values()):
        np.testing.assert_almost_equal(
            gold_result.cpu().numpy(), test_result, decimal=3,
        )
    return result
# 检查模型
def check_onnx_model(model, onnx_filename, input_image):
    with torch.no_grad():
        torch_out = {"output": model(input_image)}
    check_onnx_output(onnx_filename, input_image, torch_out)
    onnx_model = onnx.load(onnx_filename)
    onnx.checker.check_model(onnx_model)
    print("模型测试成功")
    return onnx_model
# 检测导出的onnx模型是否完整
# 一般出现问题程序直接报错,不过很少出现问题
onnx_model = check_onnx_model(model, onnx_name, dummy)

# -----  模型简化
print("-----  模型简化 -----")
# 基于onnx-simplifier简化模型,https://github.com/daquexian/onnx-simplifier
# 也可以命令行输入python3 -m onnxsim input_onnx_model output_onnx_model
# 或者使用在线网站直接转换https://convertmodel.com/

# 输出模型名
filename = onnx_name + "sim.onnx"
# 简化模型
# 设置skip_fuse_bn=True表示跳过融合bn层,pytorch高版本融合bn层会出错
simplified_model, check = simplify(onnx_model, skip_fuse_bn=True)
onnx.save_model(simplified_model, filename)
onnx.checker.check_model(simplified_model)
# 如果出错
assert check, "简化模型失败"
print("模型简化成功")
################################ pytorch转Onnx ######################################
xiaoyuanzi@xiaoyuanzi-virtual-machine:~$ sudo docker run  -v $PWD/python/onnx:/usr/src/python/onnx  -w /usr/src/python/onnx new_python:3.6 python Opencv_Onnx.py
1.10.0+cu102
4.5.4-dev
1.19.5
1.10.1
----- 模型测试 -----
模型测试成功
-----  模型简化 -----
模型简化成功

4. Opencv 调用 Onnx

############################### Opencv dnn 调用 Onnx #####################################
# 载入onnx模块
onnx_name = './vgg16.onnx'
model_ = onnx.load(onnx_name)
# 检查IR是否良好
onnx.checker.check_model(model_)
# opencv dnn加载
net = cv2.dnn.readNetFromONNX(onnx_name)

frame = cv2.imread(img_file)
# Create a 4D blob from a frame.
inpWidth = dummy.shape[-2]
inpHeight = dummy.shape[-2]
# blob = cv2.dnn.blobFromImage(frame, size=(inpWidth, inpHeight), crop=False)
blob = cv2.dnn.blobFromImage(frame,
                             scalefactor=1.0 / 255,
                             size=(inpWidth, inpHeight),
                             mean=[0.485, 0.456, 0.406],
                             swapRB=True,
                             crop=False)
# Run a model
net.setInput(blob)
out = net.forward()
print(out.shape)        # 得到推理的结果

# Get a class with a highest score.
out = out.flatten()
classId = np.argmax(out)
confidence = out[classId]

# Put efficiency information.
t, _ = net.getPerfProfile()
label = 'Inference time: %.2f ms' % (t * 1000.0 / cv2.getTickFrequency())
print(label)

# Print predicted class.
label = '%s: %.4f' % (classes[classId] if classes else 'Class #%d' % classId, confidence)
print(label)
################################ Opencv dnn 调用 Onnx #####################################

结果:

xiaoyuanzi@xiaoyuanzi-virtual-machine:~$ sudo docker run  -v $PWD/python/onnx:/usr/src/python/onnx  -w /usr/src/python/onnx new_python:3.6 python Opencv_Onnx.py
1.10.0+cu102
4.5.4-dev
1.19.5
1.10.1
(1, 1000)
Inference time: 206.04 ms
diaper, nappy, napkin: 11.8783

注:
Linux不支持使用cv2.imshow()
建议使用plt.imshow()

5. Onnxruntime 调用 Onnx

################################ Onnxruntime 调用 Onnx #####################################
transform = transforms.Compose([  # [1]
    transforms.Resize(256),  # [2]
    transforms.CenterCrop(224),  # [3]
    transforms.ToTensor(),  # [4]
    transforms.Normalize(  # [5]
        mean=[0.485, 0.456, 0.406],  # [6]
        std=[0.229, 0.224, 0.225]  # [7]
    )])
img = Image.open(img_file)
img_t = transform(img)              # 输入给模型的图像数据要先进行转换
batch_t = torch.unsqueeze(img_t, 0)
img_np = batch_t.cpu().numpy()


OnnxR_model = './vgg16.onnx'
session = onnxruntime.InferenceSession(OnnxR_model)
#compute ONNX Runtime output prediction
inputs = {session.get_inputs()[0].name: img_np}
outs = session.run(None, inputs)[0]

# Get a class with a highest score.
outs = outs.flatten()
classId = np.argmax(outs)
print(classId)
confidence = outs[classId]

# Print predicted class.
label = '%s: %.4f' % (classes[classId] if classes else 'Class #%d' % classId, confidence)
print(label)

################################ Onnxruntime 调用 Onnx #####################################

结果:

French bulldog: 19.8963
  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 3
    评论
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

只搬烫手的砖

你的鼓励将是我最大的动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值