PyTorch模型转TensorRT是怎么实现的?

转换步骤概览

  • 准备好模型定义文件(.py文件)
  • 准备好训练完成的权重文件(.pth或.pth.tar)
  • 安装onnx和onnxruntime
  • 将训练好的模型转换为.onnx格式
  • 安装tensorRT

环境参数

ubuntu-18.04
PyTorch-1.8.1
onnx-1.9.0
onnxruntime-1.7.2
cuda-11.1
cudnn-8.2.0
TensorRT-7.2.3.4

PyTorch转ONNX

Step1:安装ONNX和ONNXRUNTIME

网上找到的安装方式是通过pip

pip install onnx
pip install onnxruntime

如果使用的是Anaconda环境,conda安装也是可以的。

conda install -c conda-forge onnx
conda install -c conda-forge onnxruntime

Step2:安装netron

netron是用于可视化网络结构的,便于debug。

pip install netron

Step3 PyTorch转ONNx

安装完成后,可以根据下面code进行转换。

#--*-- coding:utf-8 --*--
import onnx
# 注意这里导入onnx时必须在torch导入之前,否则会出现segmentation fault
import torch
import torchvision 

from model import Net

model= Net(args).cuda()#初始化模型
checkpoint = torch.load(checkpoint_path)
net.load_state_dict(checkpoint['state_dict'])#载入训练好的权重文件
print ("Model and weights LOADED successfully")

export_onnx_file = './net.onnx'
x = torch.onnx.export(net,
					torch.randn(1,1,224,224,device='cuda'), #根据输入要求初始化一个dummy input
					export_onnx_file,
					verbose=False, #是否以字符串形式显示计算图
					input_names = ["inputs"]+["params_%d"%i for i in range(120)],#输入节点的名称,这里也可以给一个list,list中名称分别对应每一层可学习的参数,便于后续查询
					output_names = ["outputs"],# 输出节点的名称
					opset_version  = 10,#onnx 支持采用的operator set, 应该和pytorch版本相关
					do_constant_folding = True,
					dynamic_axes = {"inputs":{0:"batch_size"}, 2:"h", 3:"w"}, "outputs":{0: "batch_size"},})

net = onnx.load('./erfnet.onnx') #加载onnx 计算图
onnx.checker.check_model(net) # 检查文件模型是否正确
onnx.helper.printable_graph(net.graph) #输出onnx的计算图

dynamic_axes用于指定输入、输出中的可变维度。输入输出的batch_size在这里都设为了可变,输入的第2和第3维也设置为了可变。

Step 4:验证ONNX模型

下面可视化onnx模型,同时测试模型是否正确运行

import netron
import onnxruntime
import numpy as np
from PIL import Image
import cv2

netron.start('./net.onnx')
test_image = np.asarray(Image.open(test_image_path).convert('L'),dtype='float32') /255.
test_image = cv2.resize(np.array(test_image),(224,224),interpolation = cv2.INTER_CUBIC)
test_image = test_image[np.newaxis,np.newaxis,:,:]
session = onnxruntime.InferenceSession('./net.onnx')
outputs = session.run(None, {"inputs": test_image})
print(len(outputs))
print(outputs[0].shape)
#根据需要处理一下outputs[0],并可视化一下结果,看看结果是否正常

ONNX转TensorRT

Step1:从NVIDIA下载TensorRT下载安装包 https://developer.nvidia.com/tensorrt

根据自己的cuda版本选择,我选择的是TensorRT 7.2.3,下载到本地。

cd download_path
dpkg -i nv-tensorrt-repo-ubuntu1804-cuda11.1-trt7.2.3.4-ga-20210226_1-1_amd64.deb
sudo apt-get update
sudo apt-get install tensorrt

查了一下NVIDIA的官方安装教程https://docs.nvidia.com/deeplearning/tensorrt/quick-start-guide/index.html#install,由于可能需要调用TensorRT Python API,我们还需要先安装PyCUDA。这边先插入一下PyCUDA的安装。

pip install 'pycuda<2021.1'

遇到任何问题,请参考官方说明 https://wiki.tiker.net/PyCuda/Installation/Linux/#step-1-download-and-unpack-pycuda
如果使用的是Python 3.X,再执行一下以下安装。

sudo apt-get install python3-libnvinfer-dev

如果需要ONNX graphsurgeon或使用Python模块,还需要执行以下命令。

sudo apt-get install onnx-graphsurgeon

验证是否安装成功。

dpkg -l | grep TensorRT

得到类似上图的结果就是安装成功了。

问题:此时在python中import tensorrt,得到ModuleNotFoundError: No module named 'tensorrt'的报错信息。

网上查了一下,通过dpkg安装的tensorrt是默认安装在系统python中,而不是Anaconda环境的python里的。由于系统默认的python是3.6,而Anaconda里使用的是3.8.8,通过export PYTHONPATH的方式,又会出现python版本不匹配的问题。

重新搜索了一下如何在anaconda环境里安装tensorRT。

pip3 install --upgrade setuptools pip
pip install nvidia-pyindex
pip install nvidia-tensorrt

验证一下这是Anconda环境的python是否可以import tensorrt。

import tensorrt
print(tensorrt.__version__)
#输出8.0.0.3

Step 2:ONNX转TensorRT

先说一下,在这一步里遇到了*** AttributeError: ‘tensorrt.tensorrt.Builder' object has no attribute 'max_workspace_size'的报错信息。网上查了一下,是8.0.0.3版本的bug,要退回到7.2.3.4。
emmm…

pip unintall nvidia-tensorrt #先把8.0.0.3版本卸载掉
pip install nvidia-tensorrt==7.2.* --index-url https://pypi.ngc.nvidia.com # 安装7.2.3.4banben 

转换代码

import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import torch
import time
from PIL import Image
import cv2,os
import torchvision
import numpy as np
from scipy.special import softmax

### get_img_np_nchw h和postprocess_the_output函数根据需要进行修改

TRT_LOGGER = trt.Logger()

def get_img_np_nchw(img_path):
	img = Image.open(img_path).convert('L')
	img = np.asarray(img, dtype='float32')
	img = cv2.resize(np.array(img),(224, 224), interpolation = cv2.INTER_CUBIC)
	img = img / 255.
	img = img[np.newaxis, np.newaxis]
	return image
class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        """host_mom指代cpu内存,device_mem指代GPU内存
        """
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()

def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream

def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="",fp16_mode=False, int8_mode=False,save_engine=False):
    """
    params max_batch_size:      预先指定大小好分配显存
    params onnx_file_path:      onnx文件路径
    params engine_file_path:    待保存的序列化的引擎文件路径
    params fp16_mode:           是否采用FP16
    params int8_mode:           是否采用INT8
    params save_engine:         是否保存引擎
    returns:                    ICudaEngine
    """
    # 如果已经存在序列化之后的引擎,则直接反序列化得到cudaEngine
    if os.path.exists(engine_file_path):
        print("Reading engine from file: {}".format(engine_file_path))
        with open(engine_file_path, 'rb') as f, \
            trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())  # 反序列化
    else:  # 由onnx创建cudaEngine

        # 使用logger创建一个builder
        # builder创建一个计算图 INetworkDefinition
        explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        # In TensorRT 7.0, the ONNX parser only supports full-dimensions mode, meaning that your network definition must be created with the explicitBatch flag set. For more information, see Working With Dynamic Shapes.

        with trt.Builder(TRT_LOGGER) as builder, \
            builder.create_network(explicit_batch) as network,  \
            trt.OnnxParser(network, TRT_LOGGER) as parser, \
            builder.create_builder_config() as config: # 使用onnx的解析器绑定计算图,后续将通过解析填充计算图
            profile = builder.create_optimization_profile()
            profile.set_shape("inputs", (1, 1, 224, 224),(1,1,224,224),(1,1,224,224))
            config.add_optimization_profile(profile)

            config.max_workspace_size = 1<<30  # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间
            builder.max_batch_size = max_batch_size # 执行时最大可以使用的batchsize
            builder.fp16_mode = fp16_mode
            builder.int8_mode = int8_mode

            if int8_mode:
                # To be updated
                raise NotImplementedError

            # 解析onnx文件,填充计算图
            if not os.path.exists(onnx_file_path):
                quit("ONNX file {} not found!".format(onnx_file_path))
            print('loading onnx file from path {} ...'.format(onnx_file_path))
            # with open(onnx_file_path, 'rb') as model: # 二值化的网络结果和参数
            #     print("Begining onnx file parsing")
            #     parser.parse(model.read())  # 解析onnx文件
            parser.parse_from_file(onnx_file_path) # parser还有一个从文件解析onnx的方法

            print("Completed parsing of onnx file")
            # 填充计算图完成后,则使用builder从计算图中创建CudaEngine
            print("Building an engine from file{}' this may take a while...".format(onnx_file_path))

            #################
            # import pdb;pdb.set_trace()
            print(network.get_layer(network.num_layers-1).get_output(0).shape)
            # network.mark_output(network.get_layer(network.num_layers -1).get_output(0))
            engine = builder.build_engine(network,config)  # 注意,这里的network是INetworkDefinition类型,即填充后的计算图
            print("Completed creating Engine")
            if save_engine:  #保存engine供以后直接反序列化使用
                with open(engine_file_path, 'wb') as f:
                    f.write(engine.serialize())  # 序列化
            return engine

def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
    # Transfer data from CPU to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference.
    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]

def postprocess_the_outputs(outputs, shape_of_output):
    outputs = outputs.reshape(*shape_of_output)
    out = np.argmax(softmax(outputs,axis=1)[0,...],axis=0)
    # import pdb;pdb.set_trace()
    return out
# 验证TensorRT模型是否正确
onnx_model_path = './Net.onnx'
max_batch_size = 1
# These two modes are dependent on hardwares
fp16_mode = False
int8_mode = False
trt_engine_path = './model_fp16_{}_int8_{}.trt'.format(fp16_mode, int8_mode)
# Build an engine
engine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode, int8_mode , save_engine=True)
# Create the context for this engine
context = engine.create_execution_context()
# Allocate buffers for input and output
inputs, outputs, bindings, stream = allocate_buffers(engine)  # input, output: host # bindings

# Do inference
img_np_nchw = get_img_np_nchw(img_path)
inputs[0].host = img_np_nchw.reshape(-1)
shape_of_output = (max_batch_size, 2, 224, 224)

# inputs[1].host = ... for multiple input
t1 = time.time()
trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) # numpy data
t2 = time.time()
feat = postprocess_the_outputs(trt_outputs[0], shape_of_output)

print('TensorRT ok')
print("Inference time with the TensorRT engine: {}".format(t2-t1))

根据https://www.jb51.net/article/187266.htm文章里的方法,转换的时候会报下面的错误:

原来我是根据链接里的代买进行转换的,后来进行了修改,按我文中的转换代码不会有问题,

修改的地方在

with trt.Builder(TRT_LOGGER) as builder, \
            builder.create_network(explicit_batch) as network,  \
            trt.OnnxParser(network, TRT_LOGGER) as parser, \
            builder.create_builder_config() as config: # 使用onnx的解析器绑定计算图,后续将通过解析填充计算图
            profile = builder.create_optimization_profile()
            profile.set_shape("inputs", (1, 1, 224, 224),(1,1,224,224),(1,1,224,224))
            config.add_optimization_profile(profile)

            config.max_workspace_size = 1<<30  # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间
            engine = builder.build_engine(network,config)

将链接中相应的代码进行修改或添加,就没有这个问题了。

到此这篇关于PyTorch模型转TensorRT是怎么实现的?的文章就介绍到这了,更多相关PyTorch模型转TensorRT内容请搜索我们以前的文章或继续浏览下面的相关文章希望大家以后多多支持我们!

(0)

相关推荐

  • Pytorch通过保存为ONNX模型转TensorRT5的实现

    1 Pytorch以ONNX方式保存模型 def saveONNX(model, filepath): ''' 保存ONNX模型 :param model: 神经网络模型 :param filepath: 文件保存路径 ''' # 神经网络输入数据类型 dummy_input = torch.randn(self.config.BATCH_SIZE, 1, 28, 28, device='cuda') torch.onnx.export(model, dummy_input, filepath,

  • PyTorch 如何检查模型梯度是否可导

    一.PyTorch 检查模型梯度是否可导 当我们构建复杂网络模型或在模型中加入复杂操作时,可能会需要验证该模型或操作是否可导,即模型是否能够优化,在PyTorch框架下,我们可以使用torch.autograd.gradcheck函数来实现这一功能. 首先看一下官方文档中关于该函数的介绍: 可以看到官方文档中介绍了该函数基于何种方法,以及其参数列表,下面给出几个例子介绍其使用方法,注意: Tensor需要是双精度浮点型且设置requires_grad = True 第一个例子:检查某一操作是否可

  • 解决Pytorch修改预训练模型时遇到key不匹配的情况

    一.Pytorch修改预训练模型时遇到key不匹配 最近想着修改网络的预训练模型vgg.pth,但是发现当我加载预训练模型权重到新建的模型并保存之后. 在我使用新赋值的网络模型时出现了key不匹配的问题 #加载后保存(未修改网络) base_weights = torch.load(args.save_folder + args.basenet) ssd_net.vgg.load_state_dict(base_weights) torch.save(ssd_net.state_dict(),

  • pytorch 预训练模型读取修改相关参数的填坑问题

    pytorch 预训练模型读取修改相关参数的填坑 修改部分层,仍然调用之前的模型参数. resnet = resnet50(pretrained=False) resnet.load_state_dict(torch.load(args.predir)) res_conv31 = Bottleneck_dilated(1024, 256,dilated_rate = 2) print("---------------------",res_conv31) print("---

  • pytorch_pretrained_bert如何将tensorflow模型转化为pytorch模型

    pytorch_pretrained_bert将tensorflow模型转化为pytorch模型 BERT仓库里的模型是TensorFlow版本的,需要进行相应的转换才能在pytorch中使用 在Google BERT仓库里下载需要的模型,这里使用的是中文预训练模型(chinese_L-12_H-768_A_12) 下载chinese_L-12_H-768_A-12.zip后解压,里面有5个文件 chinese_L-12_H-768_A-12.zip后解压,里面有5个文件 bert_config

  • pytorch模型的保存和加载、checkpoint操作

    其实之前笔者写代码的时候用到模型的保存和加载,需要用的时候就去度娘搜一下大致代码,现在有时间就来整理下整个pytorch模型的保存和加载,开始学习把~ pytorch的模型和参数是分开的,可以分别保存或加载模型和参数.所以pytorch的保存和加载对应存在两种方式: 1. 直接保存加载模型 (1)保存和加载整个模型 # 保存模型 torch.save(model, 'model.pth\pkl\pt') #一般形式torch.save(net, PATH) # 加载模型 model = torc

  • PyTorch模型转TensorRT是怎么实现的?

    转换步骤概览 准备好模型定义文件(.py文件) 准备好训练完成的权重文件(.pth或.pth.tar) 安装onnx和onnxruntime 将训练好的模型转换为.onnx格式 安装tensorRT 环境参数 ubuntu-18.04 PyTorch-1.8.1 onnx-1.9.0 onnxruntime-1.7.2 cuda-11.1 cudnn-8.2.0 TensorRT-7.2.3.4 PyTorch转ONNX Step1:安装ONNX和ONNXRUNTIME 网上找到的安装方式是通过

  • pytorch 模型可视化的例子

    如下所示: 一. visualize.py from graphviz import Digraph import torch from torch.autograd import Variable def make_dot(var, params=None): """ Produces Graphviz representation of PyTorch autograd graph Blue nodes are the Variables that require gra

  • 画pytorch模型图,以及参数计算的方法

    刚入pytorch的坑,代码还没看太懂.之前用keras用习惯了,第一次使用pytorch还有些不适应,希望广大老司机多多指教. 首先说说,我们如何可视化模型.在keras中就一句话,keras.summary(),或者plot_model(),就可以把模型展现的淋漓尽致. 但是pytorch中好像没有这样一个api让我们直观的看到模型的样子.但是有网友提供了一段代码,可以把模型画出来,对我来说简直就是如有神助啊. 话不多说,上代码吧. import torch from torch.autog

  • 将Pytorch模型从CPU转换成GPU的实现方法

    最近将Pytorch程序迁移到GPU上去的一些工作和思考 环境:Ubuntu 16.04.3 Python版本:3.5.2 Pytorch版本:0.4.0 0. 序言 大家知道,在深度学习中使用GPU来对模型进行训练是可以通过并行化其计算来提高运行效率,这里就不多谈了. 最近申请到了实验室的服务器来跑程序,成功将我简陋的程序改成了"高大上"GPU版本. 看到网上总体来说少了很多介绍,这里决定将我的一些思考和工作记录下来. 1. 如何进行迁移 由于我使用的是Pytorch写的模型,网上给

  • 把vgg-face.mat权重迁移到pytorch模型示例

    最近使用pytorch时,需要用到一个预训练好的人脸识别模型提取人脸ID特征,想到很多人都在用用vgg-face,但是vgg-face没有pytorch的模型,于是写个vgg-face.mat转到pytorch模型的代码 #!/usr/bin/env python2 # -*- coding: utf-8 -*- """ Created on Thu May 10 10:41:40 2018 @author: hy """ import torc

  • Pytorch模型转onnx模型实例

    如下所示: import io import torch import torch.onnx from models.C3AEModel import PlainC3AENetCBAM device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") def test(): model = PlainC3AENetCBAM() pthfile = r'/home/joy/Projects/

  • pytorch模型预测结果与ndarray互转方式

    预测结果转为numpy: logits=model(feature) #如果模型是跑在GPU上 result=logits.data.cpu().numpy() / logits.cpu().numpy() #如果模型跑在cpu上 result=logits.data.numpy() / logits.numpy() 将矩阵转为tensor: np_arr = np.array([1,2,3,4]) tensor=torch.from_numpy(np_arr) 以上这篇pytorch模型预测结

  • pytorch模型存储的2种实现方法

    1.保存整个网络结构信息和模型参数信息: torch.save(model_object, './model.pth') 直接加载即可使用: model = torch.load('./model.pth') 2.只保存网络的模型参数-推荐使用 torch.save(model_object.state_dict(), './params.pth') 加载则要先从本地网络模块导入网络,然后再加载参数: from models import AgeModel model = AgeModel()

  • pytorch 模型的train模式与eval模式实例

    原因 对于一些含有batch normalization或者是Dropout层的模型来说,训练时的froward和验证时的forward有计算上是不同的,因此在前向传递过程中需要指定模型是在训练还是在验证. 源代码 [docs] def train(self, mode=True): r"""Sets the module in training mode. This has any effect only on certain modules. See documentat

  • 浅谈pytorch 模型 .pt, .pth, .pkl的区别及模型保存方式

    我们经常会看到后缀名为.pt, .pth, .pkl的pytorch模型文件,这几种模型文件在格式上有什么区别吗? 其实它们并不是在格式上有区别,只是后缀不同而已(仅此而已),在用torch.save()函数保存模型文件时,各人有不同的喜好,有些人喜欢用.pt后缀,有些人喜欢用.pth或.pkl.用相同的torch.save()语句保存出来的模型文件没有什么不同. 在pytorch官方的文档/代码里,有用.pt的,也有用.pth的.一般惯例是使用.pth,但是官方文档里貌似.pt更多,而且官方也

随机推荐