文章目錄

• • 一、什么是 ONNX
  • 二、PyTorch 轉 ONNX
  • 三、什么是 TensorRT
  • 四、ONNX 轉 TensorRT

在深度學習模型落地的過程中，會面臨將模型部署到邊端設備的問題，模型訓練使用不同的框架，則推理的時候也需要使用相同的框架，但不同類型的平臺，調優和實現起來非常困難，因為每個平臺都有不同的功能和特性。如果需要在該平臺上運行多種框架，則會增加復雜性，所以 ONNX 便派上了用場。可以通過將不同框架訓練的模型轉換成通用的 ONNX 模型，再進而轉換成各個平臺支持的格式，就可以實現簡化部署。

一、什么是 ONNX

ONNX 是 Open Neural Network Exchange 的簡稱，也叫開放神經網絡交換，是一個用于表示深度學習模型的標準，可使模型在不同框架直接轉換。

ONNX 是邁向開放式生態系統的第一步，使得開發人員不局限于某種特定的開發工具，為模型提供了開源格式。

ONNX 目前支持的框架有：Caffe2、PyTorch、TensorFlow、MXNet、TensorRT、CNTK 等

ONNX 通常來說就是一個中介，是一種手段，在把模型轉換成 ONNX 之后，再轉換成可部署的形式，如 TensorRT。

典型的結構轉換路線：

• Pytorch → ONNX → TensorRT
• Pytorch → ONNX → TVM
• TF → ONNX → NCNN

二、PyTorch 轉 ONNX

import onnxruntime import torch torch.onnx.export(model, (img_list, ), 'tmp.onnx',input_names=['input.1'],output_names=['output'],export_params=True,keep_initializers_as_inputs=False,verbose=show,opset_version=opset_version,dynamic_axes=dynamic_axes)) # onnx 模型簡化： python3 -m onnxsim tmp.onnx tmp_simplify.onnx

三、什么是 TensorRT

TensorRT 是一個高性能的深度學習推理（inference）優化器，是可以在 NVIDIA 各種 GPU 硬件平臺下運行的一個 C++ 推理框架，可以為深度學習提高低延時、高吞吐率的部署推理，可用于對嵌入式平臺、自動駕駛平臺的推理加速。

將 TensorRT 和 NVIDIA 的 GPU 結合起來，能在幾乎所有框架中進行快速和高效的部署推理。從而提升這個模型在英偉達GPU上運行的速度。速度提升的比例是比較可觀的。

我們已知模型包括訓練和推理兩個階段，訓練的時候包含了前向傳播和反向傳播，推理只包含前向傳播，所以預測時候的速度更重要。

在訓練的時候，一般會使用多 GPU 分布式訓練，在部署推理的時候，往往使用單個 GPU 甚至嵌入式平臺。模型訓練時采樣的框架會不同，不同機器的性能會存在差異，導致推理速度變慢，無法滿足高實時性。而 TensorRT 就是推理優化器，把 ONNX 模型轉換為 TensorRT 之后，就可以在相關邊端部署了。

TensorRT 的優化方法：

TensorRT 有多種優化方法，最主要的是前兩種：

• 層間融合或張量融合：

  TensorRT通過對層間的橫向或縱向合并（合并后的結構稱為CBR，意指 convolution, bias, and ReLU layers are fused to form a single layer），使得層的數量大大減少。橫向合并可以把卷積、偏置和激活層合并成一個CBR結構，只占用一個CUDA核心。縱向合并可以把結構相同，但是權值不同的層合并成一個更寬的層，也只占用一個CUDA核心。合并之后的計算圖）的層次更少了，占用的CUDA核心數也少了，因此整個模型結構會更小，更快，更高效。

• 數據精度校準：

  大部分深度學習框架在訓練神經網絡時網絡中的張量都是32位浮點數的精度（Full 32-bit precision，FP32），一旦網絡訓練完成，在部署推理的過程中由于不需要反向傳播，完全可以適當降低數據精度，比如降為FP16或INT8的精度。更低的數據精度將會使得內存占用和延遲更低，模型體積更小。

• Kernel Auto-Tuning：

  網絡模型在推理計算時，是調用 GPU 的 CUDA 進行計算的，TensorRT 可以真的不同的算法、不同的模型結構、不同的 GPU 平臺等，進行 CUDA 調整，以保證當前模型在特定平臺上以最優的性能計算。

  假設在 3090 和 T4 上要分別部署，則需要分別在這兩個平臺上進行 TensorRT 的轉換，然后在對應的平臺上使用，而不能在相同同的平臺上轉換，在不同的平臺上使用。

• Dynamic Tensor Memory：

  在每個 tensor 使用期間，TensorRT 會為其指定顯存，避免顯存重復申請，減少內存占用和提高重復使用效率

四、ONNX 轉 TensorRT

def convert_tensorrt_engine(onnx_fn, trt_fn, max_batch_size, fp16=True, int8_calibrator=None, workspace=2_000_000_000):network_creation_flag = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)with trt.Builder(TRT_LOGGER) as builder, builder.create_network(network_creation_flag) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:builder.max_workspace_size = workspacebuilder.max_batch_size = max_batch_sizebuilder.fp16_mode = fp16if int8_calibrator:builder.int8_mode = Truebuilder.int8_calibrator = int8_calibratorwith open(onnx_fn, "rb") as f:if not parser.parse(f.read()):print("got {} errors: ".format(parser.num_errors))for i in range(parser.num_errors):e = parser.get_error(i)print(e.code(), e.desc(), e.node())returnelse:print("parse successful")print("inputs: ", network.num_inputs)# inputs = [network.get_input(i) for i in range(network.num_inputs)]# opt_profiles = create_optimization_profiles(builder, inputs)# add_profiles(config, inputs, opt_profiles)for i in range(network.num_inputs):print(i, network.get_input(i).name, network.get_input(i).shape)print("outputs: ", network.num_outputs)for i in range(network.num_outputs):output = network.get_output(i)print(i, output.name, output.shape)engine = builder.build_cuda_engine(network)with open(trt_fn, "wb") as f:f.write(engine.serialize())print("done")

總結

以上是生活随笔為你收集整理的【TensorRT】将 PyTorch 转化为可部署的 TensorRT的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。