Pytorch轉ONNX轉TensorRT加速推理過程

我們訓練好深度學習模型後，這時其仍然需在特定的深度學習框架下運行，往往不能進行高性能推理。

NVIDIA提供了一套高效推理的框架——TensorRT，可將已訓練好的模型轉爲TensorRT引擎格式，然後進行高效推理。

對於Pytorch用戶而言，該技術路線爲：pytorch model-->onnx file-->TensorRT engine。

因此，我們需要做的只有三步：

將Pytorch模型轉爲ONNX作爲中間格式；
將ONNX文件轉爲TensorRT引擎（格式包括：FP32、FP16、INT8）；
使用TensorRT引擎文件進行推理計算。

關於TensorRT的介紹網上資料較多，這裏就不再贅述。下面將結合這三個步驟對整個過程進行簡單介紹。

詳細的代碼文件我已整理到GitHub（https://github.com/qq995431104/Pytorch2TensorRT.git），歡迎大家參考並給個star~~

1、Pytorch to ONNX

2、ONNX to TensorRT

3、推理

1、Pytorch to ONNX

這一步比較簡單，只要你的模型中所有OP均被ONNX支持，即可利用Pytorch中的ONN庫進行轉換。參考如下代碼：

import torch

def get_model():
    """ Define your own model and return it
    :return: Your own model
    """
    pass

def get_onnx(model, onnx_save_path, example_tensor):

    example_tensor = example_tensor.cuda()

    _ = torch.onnx.export(model,  # model being run
                                  example_tensor,  # model input (or a tuple for multiple inputs)
                                  onnx_save_path,
                                  verbose=False,  # store the trained parameter weights inside the model file
                                  training=False,
                                  do_constant_folding=True,
                                  input_names=['input'],
                                  output_names=['output']
                                  )

if __name__ == '__main__':

    model = get_model()
    onnx_save_path = "onnx/resnet50_2.onnx"
    example_tensor = torch.randn(1, 3, 288, 512, device='cuda')

    # 導出模型
    get_onnx(model, onnx_save_path, example_tensor)

需要提供的有：加載好的Pytorch模型、一個輸入樣例。其中，模型需要按照自己的方式導入並加載模型，輸入樣例的格式爲BCHW，B爲batch_size，CHW爲通道、高、寬，CHW的值需要與你自己的模型相匹配，否則後面轉換成功後輸出結果也不對。

如果出現“RuntimeError: ONNX export failed: Couldn't export Python operator XXXX”錯誤提示，說明你的模型中有ONNX不支持的OP，可以嘗試升級Pytorch版本，或者編寫自定義op，這部分尚未進行研究，後續有了進展會更新上來。

2、ONNX to TensorRT

現在已經有了ONNX文件了，需要利用TensorRT提供的OnnxParser解析該文件，同理：Caffe對應的有CaffPaser、TensorFlow的UFF格式對應的有UffParser。

先使用TensorRT創建一個builder，然後創建一個network，然後利用對應的Parser將ONNX文件加載進去；

接着，對builder指定一些參數設置，如：max_batch_size、max_workspace_size；

如需轉爲特定格式，如fp16或int8，需指定相應參數：fp16_mode或int8_mode設爲True；

對於Int8格式，需要：

準備一個校準集，用於在轉換過程中尋找使得轉換後的激活值分佈與原來的FP32類型的激活值分佈差異最小的閾值;
並寫一個校準器類，該類需繼承trt.IInt8EntropyCalibrator2父類，並重寫get_batch_size, get_batch, read_calibration_cache, write_calibration_cache這幾個方法。具體做法參考腳本myCalibrator.py.
使用時，需額外指定cache_file，該參數是校準集cache文件的路徑，會在校準過程中生成，方便下一次校準時快速提取。

示例代碼如下：

def ONNX2TRT(args, calib=None):
    ''' convert onnx to tensorrt engine, use mode of ['fp32', 'fp16', 'int8']
    :return: trt engine
    '''

    assert args.mode.lower() in ['fp32', 'fp16', 'int8'], "mode should be in ['fp32', 'fp16', 'int8']"

    G_LOGGER = trt.Logger(trt.Logger.WARNING)
    with trt.Builder(G_LOGGER) as builder, builder.create_network() as network, \
            trt.OnnxParser(network, G_LOGGER) as parser:

        builder.max_batch_size = args.batch_size
        builder.max_workspace_size = 1 << 30
        if args.mode.lower() == 'int8':
            assert (builder.platform_has_fast_int8 == True), "not support int8"
            builder.int8_mode = True
            builder.int8_calibrator = calib
        elif args.mode.lower() == 'fp16':
            assert (builder.platform_has_fast_fp16 == True), "not support fp16"
            builder.fp16_mode = True

        print('Loading ONNX file from path {}...'.format(args.onnx_file_path))
        with open(args.onnx_file_path, 'rb') as model:
            print('Beginning ONNX file parsing')
            parser.parse(model.read())
        print('Completed parsing of ONNX file')

        print('Building an engine from file {}; this may take a while...'.format(args.onnx_file_path))
        engine = builder.build_cuda_engine(network)
        print("Created engine success! ")

        # 保存計劃文件
        print('Saving TRT engine file to path {}...'.format(args.engine_file_path))
        with open(args.engine_file_path, "wb") as f:
            f.write(engine.serialize())
        print('Engine file has already saved to {}!'.format(args.engine_file_path))
        return engine

3、推理

推理過程就完全獨立於我們原先模型所依賴的框架了。

基本過程如下：

按照原模型的輸入輸出格式，準備數據，如：輸入的shape、均值、方差，輸出的shape等；
根據第二步得到的引擎文件，利用TensorRT Runtime反序列化爲引擎engine；
創建上下文環境engine.create_execution_context()；
使用Pycuda的mem_alloc對輸入輸出分配cuda內存；
創建Stream；
使用memcpy_htod_async將IO數據放入device（一般爲GPU）；
使用context.execute_async執行推理；
使用memcpy_dtoh_async取出結果；

根據引擎文件反序列化爲TensorRT引擎的示例代碼如下：

def loadEngine2TensorRT(filepath):
    G_LOGGER = trt.Logger(trt.Logger.WARNING)
    # 反序列化引擎
    with open(filepath, "rb") as f, trt.Runtime(G_LOGGER) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())
        return engine

推理過程示例如下：

# 通過engine文件創建引擎
engine = loadEngine2TensorRT('path_to_engine_file')

# 準備輸入輸出數據
img = Image.open('XXX.jpg')
img = D.transform(img).unsqueeze(0)
img = img.numpy()
output = np.empty((1, 2), dtype=np.float32)

#創建上下文
context = engine.create_execution_context()

# 分配內存
d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
bindings = [int(d_input), int(d_output)]

# pycuda操作緩衝區
stream = cuda.Stream()

# 將輸入數據放入device
cuda.memcpy_htod_async(d_input, img, stream)

# 執行模型
context.execute_async(100, bindings, stream.handle, None)

# 將預測結果從從緩衝區取出
cuda.memcpy_dtoh_async(output, d_output, stream)
# 線程同步
stream.synchronize()

print(output)

*更多詳細內容，請參閱GitHub倉庫：https://github.com/qq995431104/Pytorch2TensorRT.git

Pytorch轉ONNX轉TensorRT加速推理過程

1、Pytorch to ONNX

2、ONNX to TensorRT

3、推理

Python實現大麥網搶票的四大關鍵技術點解析

salesforce零基礎學習（一百三十八）零碎知識點小總結（十）

關於接口協議，你必須要知道這些！

TensorRT INT8量化原理以及如何編寫校準器類進行校準

數據結構與算法（二）：Python 單向鏈表及其排序

數據結構與算法（八）：順序隊列、棧式隊列

數據結構與算法（三）：Python實現雙向鏈表並遍歷

求整數的近似平方根（牛頓迭代法）

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結