摘要

最近，學習了一些模型轉化和加速推理的知識，本文是對學習成果的總結，

對模型的轉化，本文實作了pytorch模型轉onnx模型和onnx轉TensorRT，在轉為TensorRT模型的程序中，實作了模型單精度的壓縮，

對于加速推理，本文實作GPU環境下的onnxruntime推理、TensorRT動態推理和TensorRT靜態推理，

希望本文能幫助大家，

環境配置

CUDA版本：11.3.1

cuDNN版本：8.2.1

TensorRT版本：8.0.3.4

顯卡：1650

pytorch：1.10.2

模型的轉化和推理對版本有要求，如果版本對應不上很可能出現千奇百怪的問題，所以我把我的版本資訊列出來給大家做參考，

ONNX

ONNX，全稱：Open Neural Network Exchange（ONNX，開放神經網路交換），是一個用于表示深度學習模型的標準，可使模型在不同框架之間進行轉移，

ONNX是一種針對機器學習所設計的開放式的檔案格式，用于存盤訓練好的模型，它使得不同的人工智能框架（如Pytorch, MXNet）可以采用相同格式存盤模型資料并互動， ONNX的規范及代碼主要由微軟，亞馬遜 ，Facebook 和 IBM 等公司共同開發，以開放源代碼的方式托管在Github上，目前官方支持加載ONNX模型并進行推理的深度學習框架有： Caffe2, PyTorch, MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow 也非官方的支持ONNX，—維基百科

onnx模型可以看作是模型轉化之間的中間模型，同時也是支持做推理的，一般來說，onnx的推理速度要比pytorch快上一倍，

模型轉化

onnx模型轉換和推理需要安裝Python包，如下：

pip install onnx
pip install onnxruntime-gpu

新建模型轉換腳本pytorch2onnx.py，

import torch
from torch.autograd import Variable
import onnx
import netron
print(torch.__version__)
input_name = ['input']
output_name = ['output']
input = Variable(torch.randn(1, 3, 224, 224)).cuda()
model = torch.load('model.pth', map_location="cuda:0")
torch.onnx.export(model, input, 'model_onnx.onnx',opset_version=13, input_names=input_name, output_names=output_name, verbose=True)
# 模型可視化
netron.start('model_onnx.onnx')

匯入需要的包，

列印pytorch版本，

定義input_name和output_name變數，

定義輸入格式，

加載pytorch模型，

匯出onnx模型，這里注意一下引數opset_version在8.X版本中設定為13，在7.X版本中設定為12，

yolov5中這么寫的，

 		if trt.__version__[0] == '7':  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
            grid = model.model[-1].anchor_grid
            model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
            export_onnx(model, im, file, 12, train, False, simplify)  # opset 12
            model.model[-1].anchor_grid = grid
        else:  # TensorRT >= 8
            check_version(trt.__version__, '8.0.0', hard=True)  # require tensorrt>=8.0.0
            export_onnx(model, im, file, 13, train, False, simplify)  # opset 13

查看轉化后的模型，如下圖：

推理的寫法有兩種，一種直接寫，另一種將其封裝為通用的推理類，

第一種推理寫法

先看第一種寫法，新建test_onnx.py，插入下面的代碼：

import os, sys
import time
sys.path.append(os.getcwd())
import onnxruntime
import numpy as np
import torchvision.transforms as transforms
from PIL import Image

匯入包

def get_test_transform():
    return transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
image = Image.open('11.jpg') # 289
img = get_test_transform()(image)
img = img.unsqueeze_(0) # -> NCHW, 1,3,224,224
print("input img mean {} and std {}".format(img.mean(), img.std()))
img =  np.array(img)

定義get_test_transform函式，實作影像的歸一化和resize，

讀取影像，

對影像做resize和歸一化，

增加一維batchsize，

將圖片轉為陣列，

onnx_model_path = "model_onnx.onnx"
##onnx測驗
session = onnxruntime.InferenceSession(onnx_model_path,providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'])
#compute ONNX Runtime output prediction
inputs = {session.get_inputs()[0].name: img}
time3=time.time()
outs = session.run(None, inputs)[0]
y_pred_binary = np.argmax(outs, axis=1)
print("onnx prediction", y_pred_binary[0])
time4=time.time()
print(time4-time3)

定義onnx_model_path模型的路徑，

加載onnx模型，

定義輸入，

執行推理，

獲取預測結果，

到這里第一種寫法就完成了，是不是很簡單，接下來看第二種寫法，

第二種推理寫法

新建onnx.py腳本，加入以下代碼：

import onnxruntime
class ONNXModel():
    def __init__(self, onnx_path):
        """
        :param onnx_path:
        """
        self.onnx_session = onnxruntime.InferenceSession(onnx_path,providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'])
        self.input_name = self.get_input_name(self.onnx_session)
        self.output_name = self.get_output_name(self.onnx_session)
        print("input_name:{}".format(self.input_name))
        print("output_name:{}".format(self.output_name))

    def get_output_name(self, onnx_session):
        """
        output_name = onnx_session.get_outputs()[0].name
        :param onnx_session:
        :return:
        """
        output_name = []
        for node in onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name

    def get_input_name(self, onnx_session):
        """
        input_name = onnx_session.get_inputs()[0].name
        :param onnx_session:
        :return:
        """
        input_name = []
        for node in onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name

    def get_input_feed(self, input_name, image_numpy):
        """
        input_feed={self.input_name: image_numpy}
        :param input_name:
        :param image_numpy:
        :return:
        """
        input_feed = {}
        for name in input_name:
            input_feed[name] = image_numpy
        return input_feed

    def forward(self, image_numpy):
        # 輸入資料的型別必須與模型一致,以下三種寫法都是可以的
        # scores, boxes = self.onnx_session.run(None, {self.input_name: image_numpy})
        # scores, boxes = self.onnx_session.run(self.output_name, input_feed={self.input_name: iimage_numpy})
        input_feed = self.get_input_feed(self.input_name, image_numpy)
        scores = self.onnx_session.run(self.output_name, input_feed=input_feed)
        return scores

呼叫onnx.py實作推理，新建test_onnx1.py插入代碼：

import os, sys
sys.path.append(os.getcwd())
import numpy as np
import torchvision.transforms as transforms
from PIL import Image
from models.onnx import ONNXModel

def get_test_transform():
    return transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

image = Image.open('11.jpg') # 289

img = get_test_transform()(image)
img = img.unsqueeze_(0) # -> NCHW, 1,3,224,224
print("input img mean {} and std {}".format(img.mean(), img.std()))
img =  np.array(img)
onnx_model_path = "model_onnx.onnx"
model1 = ONNXModel(onnx_model_path)
out = model1.forward(img)
y_pred_binary = np.argmax(out[0], axis=1)
print("onnx prediction1", y_pred_binary[0])

輸出結果如下：

TensorRT

TensorRT是英偉達推出的一個高性能的深度學習推理（Inference）優化器，可以為深度學習應用提供低延遲、高吞吐率的部署推理，TensorRT可用于對超大規模資料中心、嵌入式平臺或自動駕駛平臺進行推理加速，TensorRT現已能支持TensorFlow、Caffe、Mxnet、Pytorch等幾乎所有的深度學習框架，將TensorRT和NVIDIA的GPU結合起來，能在幾乎所有的框架中進行快速和高效的部署推理，

TensorRT 是一個C++庫，從 TensorRT 3 開始提供C++ API和Python API，主要用來針對 NVIDIA GPU進行 高性能推理（Inference）加速，

TensorRT的安裝可以參考我以前的文章：

https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/120360288，

本次用的8.X版本的，安裝方式一樣，

本文實作Python版本的 TensorRT 推理加速,需要安裝tensorrt包檔案，這個檔案不能直接通過pip下載，我在下載的TensorRT安裝包里，不過我下載的8.0.3.4版本中并沒有，在8.2.1.8的版本中存在這個包檔案，

所以我安裝了8.2.1.8中的whl檔案，

安裝方式，進入模型所在的目錄，執行：

pip install tensorrt-8.2.1.8-cp39-none-win_amd64.whl

模型推理用到了pycuda，執行安裝命令：

pip install pycuda

模型轉化

將onnx模型轉為TensorRT 模型，新建onnx2trt.py，插入代碼：

import tensorrt as trt

def build_engine(onnx_file_path,engine_file_path,half=False):
    """Takes an ONNX file and creates a TensorRT engine to run inference with"""
    logger = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(logger)
    config = builder.create_builder_config()
    config.max_workspace_size = 4 * 1 << 30
    flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    network = builder.create_network(flag)
    parser = trt.OnnxParser(network, logger)
    if not parser.parse_from_file(str(onnx_file_path)):
        raise RuntimeError(f'failed to load ONNX file: {onnx_file_path}')
    half &= builder.platform_has_fast_fp16
    if half:
        config.set_flag(trt.BuilderFlag.FP16)
    with builder.build_engine(network, config) as engine, open(engine_file_path, 'wb') as t:
        t.write(engine.serialize())
    return engine_file_path
if __name__ =="__main__":
    onnx_path1 = 'model_onnx.onnx'
    engine_path = 'model_trt.engine'
    build_engine(onnx_path1,engine_path,True)

build_engine函式共有三個引數：

onnx_file_path：onnx模型的路徑，

engine_file_path：TensorRT模型的路徑，

half：是否使用單精度，

單精度的模型速度更快，所以我選擇使用單精度，

通過上面的代碼就可以完成模型的轉化，下面開始實作推理部分，推理分為動態推理和靜態推理，

動態推理

新建test_trt,py檔案，插入代碼：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import torchvision.transforms as transforms
from PIL import Image

匯入需要的包，

def load_engine(engine_path):
    # TRT_LOGGER = trt.Logger(trt.Logger.WARNING)  # INFO
    TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

# 2. 讀取資料，資料處理為可以和網路結構輸入對應起來的的shape，資料可增加預處理
def get_test_transform():
    return transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

定義load_engine函式和get_test_transform函式，

load_engine用于加載TensorRT模型，

get_test_transform實作影像的resize和歸一化，

image = Image.open('11.jpg') # 289
image = get_test_transform()(image)
image = image.unsqueeze_(0) # -> NCHW, 1,3,224,224
print("input img mean {} and std {}".format(image.mean(), image.std()))
image =  np.array(image)

圖片的預處理，和onnx一樣，最后轉為陣列，

path = 'model_trt.engine'
# 1. 建立模型，構建背景關系管理器
engine = load_engine(path)
context = engine.create_execution_context()
context.active_optimization_profile = 0

# 3.分配記憶體空間，并進行資料cpu到gpu的拷貝
# 動態尺寸，每次都要set一下模型輸入的shape，0代表的就是輸入，輸出根據具體的網路結構而定，可以是0,1,2,3...其中的某個頭，
context.set_binding_shape(0, image.shape)
d_input = cuda.mem_alloc(image.nbytes)  # 分配輸入的記憶體，
output_shape = context.get_binding_shape(1)
buffer = np.empty(output_shape, dtype=np.float32)
d_output = cuda.mem_alloc(buffer.nbytes)  # 分配輸出記憶體，
cuda.memcpy_htod(d_input, image)
bindings = [d_input, d_output]

# 4.進行推理，并將結果從gpu拷貝到cpu，
context.execute_v2(bindings)  # 可異步和同步
cuda.memcpy_dtoh(buffer, d_output)
output = buffer.reshape(output_shape)
y_pred_binary = np.argmax(output, axis=1)
print(y_pred_binary[0])

輸出結果：

靜態推理

靜態推理和動態推理的代碼差不多，唯一不同的是不需要

import time

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
import time
import torchvision.transforms as transforms
from PIL import Image
def load_engine(engine_path):
    # TRT_LOGGER = trt.Logger(trt.Logger.WARNING)  # INFO
    TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

path = 'model_trt.engine'
engine = load_engine(path)
context = engine.create_execution_context()
outshape = context.get_binding_shape(1)

# 2. 讀取資料，資料處理為可以和網路結構輸入對應起來的的shape，資料可增加預處理
def get_test_transform():
    return transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

image = Image.open('11.jpg') 
image = get_test_transform()(image)
image = image.unsqueeze_(0) # -> NCHW, 1,3,224,224
print("input img mean {} and std {}".format(image.mean(), image.std()))

image =  np.array(image)
# image = np.expand_dims(image, axis=1)
image = image.astype(np.float32)

image = image.ravel()  # 資料平鋪

output = np.empty((outshape), dtype=np.float32)
d_input = cuda.mem_alloc(1 * image.size * image.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
bindings = [int(d_input), int(d_output)]
stream = cuda.Stream()


cuda.memcpy_htod(d_input, image)
context.execute_v2(bindings)
cuda.memcpy_dtoh(output, d_output)
output = output.reshape(outshape)
y_pred_binary = np.argmax(output, axis=1)
print(y_pred_binary[0])

運行結果: