HCHUNGW的部落格

12.2 TensorRT 工作流及cuda-python

前言

本節將在python中採用TensorRT進行resnet50模型推理，通過一個案例瞭解TensorRT的工作步驟流程，為後續各模組深入研究打下基礎。

本節核心內容包括TensorRT中各模組概念，python中進行推理步驟，python中的cuda庫使用。

本節用到的resnet50_bs_1.engine檔，需要提前通過trtexec來生成（resnet50_bs_1.onnx通過11章內容生成，或者從網盤下載-提取碼：24uq

trtexec --onnx=resnet50_bs_1.onnx --saveEngine=resnet50_bs_1.engine

workflow基礎概念

在python中使用TensorRT進行推理，需要瞭解cuda程式設計的概念，在這裡會出現一些新名詞，這裡進行了總結，幫助大家理解python中使用TRT的代碼。

借鑒nvidia官方教程中的一幅圖來說明TRT從構建模型到使用模型推理，需要涉及的概念。

1. Logger：用於在TensorRT日誌中記錄消息,可以來實現自訂日誌記錄器，以便更好地瞭解TensorRT運行時發生的事情。
2. Builder：用於構建一個優化的TensorRT引擎，可以使用Builder來定義和優化網路
3. BuilderConfig：用於配置Builder的行為，例如最大批處理大小、優化級別等
4. Network：用於描述一個計算圖，它由一組層組成，用來定義和構建深度學習模型，通常有三種方式獲取network，包括TensorRT API、parser、訓練框架中trt工具。
5. SerializeNetwork：用於將網路序列化為二進位格式，可以將網路模型保存到磁片上的.plan檔中，以便稍後載入和使用。
6. .plan：.plan是TensorRT引擎的序列化檔案格式，有的地方用.engine，.plan文件和.engine都是序列化的TensorRT引擎檔，但是它們有一些區別。
   • .plan檔是通用序列化檔案格式，它包含了優化後的網路和權重參數。而.engine檔是Nvidia TensorRT引擎的專用二進位格式，它包含了優化後的網路，權重參數和硬體相關資訊。 可
   • 移植性方面，由於.plan檔是通用格式，所以可以在不同的硬體平臺上使用。而.engine檔是特定於硬體的，需要在具有相同GPU架構的系統上使用。
   • 載入速度上：.plan檔的載入速度通常比.engine檔快，因為它不包含硬體相關資訊，而.engine檔必須在運行時進行硬體特定的編譯和優化。
7. Engine：Engine是TensorRT中的主要對象，它包含了優化後的網路，可以用於進行推理。可以使用Engine類載入.plan檔，創建Engine物件，並使用它來進行推理。
8. Context：Context是Engine的一個實例，它提供了對引擎計算的訪問。可以使用Context來執行推理，並在執行過程中管理輸入和輸出緩衝區。
9. Buffer：用於管理記憶體緩衝區。可以使用Buffer類來分配和釋放記憶體，並將其用作輸入和輸出緩衝區。
10. Execute：Execute是Context類的一個方法，用於執行推理。您可以使用Execute方法來執行推理，並通過傳遞輸入和輸出緩衝區來管理資料流程。

在這裡面需要重點瞭解的是Network的構建有三種方式，包括TensorRT API、parser、訓練框架中trt工具。

•  
  1. TensorRT API是手寫網路結構，一層一層的搭建
  2. parser 是採用解析器對常見的模型進行解析、轉換，常用的有ONNX Parser。本小節那裡中採用parser進行onnx模型解析，實現trt模型創建。
  3. 直接採用pytorch/tensorflow框架中的trt工具匯出模型

此處先採用parser形式獲取trt模型，後續章節再講解API形式構建trt模型。

TensorRT resnet50推理

到這裡環境有了，基礎概念瞭解了，下面通過python代碼實現resnet50的推理，通過本案例代碼梳理在代碼層面的workflow。

pycuda庫與cuda庫

正式看代碼前，有必要簡單介紹pycuda庫與cuda庫的差別。在早期版本中，代碼多以pycuda進行buffer的管理，在python3.7後，官方推薦採用cuda庫進行buffer的管理。

• cuda庫是NVIDIA提供的用於CUDA GPU程式設計的python介面,包含在cuda toolkit中。主要作用是: 直接調用cuda runtime API,如記憶體管理、執行kernel等。
• pycuda庫是一個協力廠商庫,提供了另一個python綁定到CUDA runtime API。主要作用是: 封裝cuda runtime API到python調用，管理GPU Context、Stream等。

cuda庫更基礎，pycuda庫更全面。前者集成在CUDA toolkit中,後者更靈活。

但在最新的trt版本（v8.6.1）中，官方推薦採用cuda庫，兩者使用上略有不同，在配套代碼中會實現兩種buffer管理的代碼。

python workflow

正式跑代碼前，瞭解一下代碼層面的workflow：

1. 初始化模型，獲得context
2.  
   1. 創建logger：logger = trt.Logger(trt.Logger.WARNING)
   2. 創建engine：engine = runtime.deserialize_cuda_engine(ff.read())
   3. 創建context： context = engine.create_execution_context()
3. 記憶體申請：
4.  
   1. 申請host（cpu）記憶體：進行變數據量的賦值，即完成變數記憶體分配。
   2. 申請device（gpu）記憶體： 採用cudart函數，獲得記憶體位址。 d_input = cudart.cudaMalloc(h_input.nbytes)[1]
   3. 告知context gpu地址：context.set_tensor_address(l_tensor_name[0], d_input)
5. 推理
6.  
   1. 資料拷貝 host 2 device： cudart.cudaMemcpy(d_input, h_input.ctypes.data, h_input.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)
   2. 推理： context.execute_async_v3(0)
   3. 資料拷貝 device 2 host： cudart.cudaMemcpy(h_output.ctypes.data, d_output, h_output.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)
7. 記憶體釋放
8.  
   1. cudart.cudaFree(d_input)
9. 取結果
10.  
    1. 在host的變數上即可拿到模型的輸出結果。

這裡採用配套代碼實現ResNet圖像分類，可以得到與上一章ONNX中一樣的分類效果，並且輸送量與trtexec中差別不大，大約在470 it/s。

100%|██████████| 3000/3000 [00:06<00:00, 467.81it/s]。

Copy

cuda庫的buffer管理

採用cuda庫進行buffer管理，可分3個部分，記憶體和顯存的申請、資料拷貝、顯存釋放。

這裡推薦官方教程以及官方教程代碼

在教程配套代碼的 model_infer()函數是對上述兩份資料進行了結合，下面詳細介紹cuda部分的代碼含義。

def model_infer(context, engine, img_chw_array):

    n_io = engine.num_io_tensors  # since TensorRT 8.5, the concept of Binding is replaced by I/O Tensor, all the APIs with "binding" in their name are deprecated

    l_tensor_name = [engine.get_tensor_name(ii) for ii in range(n_io)]  # get a list of I/O tensor names of the engine, because all I/O tensor in Engine and Excution Context are indexed by name, not binding number like TensorRT 8.4 or before

    # 記憶體、顯存的申請

    h_input = np.ascontiguousarray(img_chw_array)

    h_output = np.empty(context.get_tensor_shape(l_tensor_name[1]), dtype=trt.nptype(engine.get_tensor_dtype(l_tensor_name[1])))

    d_input = cudart.cudaMalloc(h_input.nbytes)[1]

    d_output = cudart.cudaMalloc(h_output.nbytes)[1]

    # 分配地址

    context.set_tensor_address(l_tensor_name[0], d_input)  # 'input'

    context.set_tensor_address(l_tensor_name[1], d_output)  # 'output'

    # 資料拷貝

    cudart.cudaMemcpy(d_input, h_input.ctypes.data, h_input.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)

    # 推理

    context.execute_async_v3(0)  # do inference computation

    # 資料拷貝

    cudart.cudaMemcpy(h_output.ctypes.data, d_output, h_output.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)

    # 釋放顯存

    cudart.cudaFree(d_input)

    cudart.cudaFree(d_output)

    return h_output

Copy

• 第3行：獲取模型輸入、輸出變數的數量，在本例中是2。
• 第4行：獲取模型輸入、輸出變數的名字，存儲在list中。本案例中是['input', 'output']，這兩個名字是在onnx導入時候設定的。
• 第7/8行：申請host端的記憶體，可看出只需要進行兩個numpy的賦值，即可開闢記憶體空間存儲變數。
• 第9/10行：調用cudart進行顯存空間申請，變數獲取的是記憶體位址。例如”47348061184 “
• 第13/14行：將顯存地址告知context，context在推理的時候才能找到它們。
• 第17行：將記憶體中資料拷貝到顯存中
• 第19行：context執行推理，此時運算結果已經到了顯存
• 第21行：將顯存中資料拷貝到記憶體中
• 第24/25行：釋放顯存中的變數。

pycuda庫的buffer管理

注意：代碼在v8.6.1上運行通過，在v10.0.0.6上未能正確使用context.execute_async_v3，因此請注意版本。

更多介面參考：https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/ExecutionContext.html

在配套代碼中pycuda進行buffer的申請及管理代碼如下：

h_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)

h_output = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)), dtype=np.float32)

d_input = cuda.mem_alloc(h_input.nbytes)

d_output = cuda.mem_alloc(h_output.nbytes)

def model_infer(context, h_input, h_output, d_input, d_output, stream, img_chw_array):

    # 圖像資料遷到 input buffer

    np.copyto(h_input, img_chw_array.ravel())

    # 資料移轉, H2D

    cuda.memcpy_htod_async(d_input, h_input, stream)

    # 推理

    context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)

    # 資料移轉，D2H

    cuda.memcpy_dtoh_async(h_output, d_output, stream)

    stream.synchronize()

    return h_output

Copy

• 第1,2行：通過cuda.pagelocked_empty創建一個陣列，該陣列位於GPU中的鎖頁記憶體，鎖頁記憶體（pinned Memory/ page locked memory）的概念可以到作業系統中瞭解，它是為了提高資料讀取、傳輸速率，用空間換時間，在記憶體中開闢一塊固定的區域進行獨享。申請的大小及資料類型，通過trt.volume(context.get_binding_shape(0)和np.float32進行設置。其中trt.volume(context.get_binding_shape(0)是獲取輸入資料的元素個數，此案例，輸入是[1, 3, 224, 224]，因此得到的數是1x3x224x224 = 150528
• 第3,4行：通過cuda.mem_alloc函數在GPU上分配了一段記憶體，返回一個指向這段記憶體的指標
• 第8行：將圖像資料遷移到input buffer中
• 第10行，將輸入資料從CPU記憶體非同步複製到GPU記憶體。其中，cuda.memcpy_htod_async函數非同步將資料從CPU記憶體複製到GPU記憶體，d_input是GPU上的記憶體指標，h_input是CPU上的numpy陣列，stream是CUDA流
• 第14行，同理，從GPU中把資料複製回到CPU端的h_output，模型最終使用h_output來表示模型預測結果

小結

本節介紹了TensorRT的工作流程，其中涉及10個主要模組，各模組的概念剛開始不好理解，可以先跳過，在後續實踐中去理解。

隨後基於onnx parser實現TensorRT模型的創建，engine檔的生成通過trtexec工具生成，並進行圖片推理，獲得了與trtexec中類似的輸送量。

最後介紹了pycuda庫和cuda庫進行buffer管理的詳細步驟，兩者在代碼效率上是一樣的，輸送量幾乎一致。