自学教程：PyTorch模型转TensorRT是怎么实现的?

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

pip install onnxpip install onnxruntime

conda install -c conda-forge onnxconda install -c conda-forge onnxruntime

pip install netron

#--*-- 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 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的计算图

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],并可视化一下结果，看看结果是否正常

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

pip install 'pycuda<2021.1'

sudo apt-get install python3-libnvinfer-dev

sudo apt-get install onnx-graphsurgeon

dpkg -l | grep TensorRT

pip3 install --upgrade setuptools pippip install nvidia-pyindexpip install nvidia-tensorrt

import tensorrtprint(tensorrt.__version__)#输出8.0.0.3

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_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, 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创建一个计算图 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 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))

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)