PyTorch模型轉(zhuǎn)TensorRT是怎么實現(xiàn)的?

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

pip install onnx
pip install onnxruntime

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

pip install netron

#--*-- coding:utf-8 --*--
import onnx 
# 注意這里導(dǎo)入onnx時必須在torch導(dǎo)入之前，否則會出現(xiàn)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'])#載入訓(xùn)練好的權(quán)重文件
print ("Model and weights LOADED successfully")

export_onnx_file = './net.onnx'
x = torch.onnx.export(net,
					torch.randn(1,1,224,224,device='cuda'), #根據(jù)輸入要求初始化一個dummy input
					export_onnx_file,
					verbose=False, #是否以字符串形式顯示計算圖
					input_names = ["inputs"]+["params_%d"%i for i in range(120)],#輸入節(jié)點的名稱，這里也可以給一個list，list中名稱分別對應(yīng)每一層可學(xué)習(xí)的參數(shù)，便于后續(xù)查詢
					output_names = ["outputs"],# 輸出節(jié)點的名稱
					opset_version  = 10,#onnx 支持采用的operator set, 應(yīng)該和pytorch版本相關(guān)
					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 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)
#根據(jù)需要處理一下outputs[0],并可視化一下結(jié)果，看看結(jié)果是否正常

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

pip install 'pycuda2021.1'

sudo apt-get install python3-libnvinfer-dev

sudo apt-get install onnx-graphsurgeon

dpkg -l | grep TensorRT

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

import tensorrt
print(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 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函數(shù)根據(jù)需要進(jìn)行修改

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內(nèi)存，device_mem指代GPU內(nèi)存
        """
        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:      預(yù)先指定大小好分配顯存
    params onnx_file_path:      onnx文件路徑
    params engine_file_path:    待保存的序列化的引擎文件路徑
    params fp16_mode:           是否采用FP16
    params int8_mode:           是否采用INT8
    params save_engine:         是否保存引擎
    returns:                    ICudaEngine
    """
    # 如果已經(jīng)存在序列化之后的引擎，則直接反序列化得到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創(chuàng)建cudaEngine
        
        # 使用logger創(chuàng)建一個builder 
        # builder創(chuàng)建一個計算圖 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的解析器綁定計算圖，后續(xù)將通過解析填充計算圖
            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 = 130  # 預(yù)先分配的工作空間大小,即ICudaEngine執(zhí)行時GPU最大需要的空間
            builder.max_batch_size = max_batch_size # 執(zhí)行時最大可以使用的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: # 二值化的網(wǎng)絡(luò)結(jié)果和參數(shù)
            #     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從計算圖中創(chuàng)建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))

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的解析器綁定計算圖，后續(xù)將通過解析填充計算圖
            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 = 130  # 預(yù)先分配的工作空間大小,即ICudaEngine執(zhí)行時GPU最大需要的空間
            engine = builder.build_engine(network,config)