PyTorch模型转TensorRT是怎么实现的?
ubuntu-18.04PyTorch-1.8.1onnx-1.9.0onnxruntime-1.7.2cuda-11.1cudnn-8.2.0TensorRT-7.2.3.4PyTorch转ONNX
Step1：安装ONNX和ONNXRUNTIME
网上找到的安装方式是通过pip
pip install onnxpip install onnxruntime
如果使用的是Anaconda环境，conda安装也是可以的。
conda install -c conda-forge onnxconda install -c conda-forge onnxruntime
Step2：安装netron
netron是用于可视化网络结构的，便于debug。
pip install netron
Step3 PyTorch转ONNx
安装完成后，可以根据下面code进行转换。
#--*-- coding:utf-8 --*--import onnx # 注意这里导入onnx时必须在torch导入之前，否则会出现segmentation faultimport torchimport torchvision from model import Netmodel= 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 inputexport_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 netronimport onnxruntimeimport numpy as npfrom PIL import Imageimport cv2netron.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_pathdpkg -i nv-tensorrt-repo-ubuntu1804-cuda11.1-trt7.2.3.4-ga-20210226_1-1_amd64.debsudo apt-get updatesudo 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 pippip install nvidia-pyindexpip install nvidia-tensorrt
验证一下这是Anconda环境的python是否可以import tensorrt。
import tensorrtprint(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 cudaimport tensorrt as trtimport torch import time from PIL import Imageimport cv2,osimport torchvision import numpy as npfrom 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 imageclass HostDeviceMem(object): def __init__(self, host_mem, device_mem):'''host_mom指代cpu内存，device_mem指代GPU内存'''self.host = host_memself.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_sizedtype = trt.nptype(engine.get_binding_dtype(binding))# Allocate host and device buffershost_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, streamdef 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创建一个计算图 INetworkDefinitionexplicit_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 updatedraise 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 enginedef 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 hardwaresfp16_mode = Falseint8_mode = Falsetrt_engine_path = ’./model_fp16_{}_int8_{}.trt’.format(fp16_mode, int8_mode)# Build an engineengine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode, int8_mode , save_engine=True)# Create the context for this enginecontext = engine.create_execution_context()# Allocate buffers for input and outputinputs, outputs, bindings, stream = allocate_buffers(engine) # input, output: host # bindings# Do inferenceimg_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 inputt1 = time.time()trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) # numpy datat2 = 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内容请搜索好吧啦网以前的文章或继续浏览下面的相关文章希望大家以后多多支持好吧啦网！
相关文章：
网公网安备