文章目錄

• • 一、什么是 ONNX
  • 二、PyTorch 轉(zhuǎn) ONNX
  • 三、什么是 TensorRT
  • 四、ONNX 轉(zhuǎn) TensorRT

在深度學(xué)習(xí)模型落地的過程中，會(huì)面臨將模型部署到邊端設(shè)備的問題，模型訓(xùn)練使用不同的框架，則推理的時(shí)候也需要使用相同的框架，但不同類型的平臺(tái)，調(diào)優(yōu)和實(shí)現(xiàn)起來非常困難，因?yàn)槊總€(gè)平臺(tái)都有不同的功能和特性。如果需要在該平臺(tái)上運(yùn)行多種框架，則會(huì)增加復(fù)雜性，所以 ONNX 便派上了用場?？梢酝ㄟ^將不同框架訓(xùn)練的模型轉(zhuǎn)換成通用的 ONNX 模型，再進(jìn)而轉(zhuǎn)換成各個(gè)平臺(tái)支持的格式，就可以實(shí)現(xiàn)簡化部署。

一、什么是 ONNX

ONNX 是 Open Neural Network Exchange 的簡稱，也叫開放神經(jīng)網(wǎng)絡(luò)交換，是一個(gè)用于表示深度學(xué)習(xí)模型的標(biāo)準(zhǔn)，可使模型在不同框架直接轉(zhuǎn)換。

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

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

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

典型的結(jié)構(gòu)轉(zhuǎn)換路線：

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

二、PyTorch 轉(zhuǎn) 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 是一個(gè)高性能的深度學(xué)習(xí)推理（inference）優(yōu)化器，是可以在 NVIDIA 各種 GPU 硬件平臺(tái)下運(yùn)行的一個(gè) C++ 推理框架，可以為深度學(xué)習(xí)提高低延時(shí)、高吞吐率的部署推理，可用于對(duì)嵌入式平臺(tái)、自動(dòng)駕駛平臺(tái)的推理加速。

將 TensorRT 和 NVIDIA 的 GPU 結(jié)合起來，能在幾乎所有框架中進(jìn)行快速和高效的部署推理。從而提升這個(gè)模型在英偉達(dá)GPU上運(yùn)行的速度。速度提升的比例是比較可觀的。

我們已知模型包括訓(xùn)練和推理兩個(gè)階段，訓(xùn)練的時(shí)候包含了前向傳播和反向傳播，推理只包含前向傳播，所以預(yù)測時(shí)候的速度更重要。

在訓(xùn)練的時(shí)候，一般會(huì)使用多 GPU 分布式訓(xùn)練，在部署推理的時(shí)候，往往使用單個(gè) GPU 甚至嵌入式平臺(tái)。模型訓(xùn)練時(shí)采樣的框架會(huì)不同，不同機(jī)器的性能會(huì)存在差異，導(dǎo)致推理速度變慢，無法滿足高實(shí)時(shí)性。而 TensorRT 就是推理優(yōu)化器，把 ONNX 模型轉(zhuǎn)換為 TensorRT 之后，就可以在相關(guān)邊端部署了。

TensorRT 的優(yōu)化方法：

TensorRT 有多種優(yōu)化方法，最主要的是前兩種：

• 層間融合或張量融合：

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

• 數(shù)據(jù)精度校準(zhǔn)：

  大部分深度學(xué)習(xí)框架在訓(xùn)練神經(jīng)網(wǎng)絡(luò)時(shí)網(wǎng)絡(luò)中的張量都是32位浮點(diǎn)數(shù)的精度（Full 32-bit precision，FP32），一旦網(wǎng)絡(luò)訓(xùn)練完成，在部署推理的過程中由于不需要反向傳播，完全可以適當(dāng)降低數(shù)據(jù)精度，比如降為FP16或INT8的精度。更低的數(shù)據(jù)精度將會(huì)使得內(nèi)存占用和延遲更低，模型體積更小。

• Kernel Auto-Tuning：

  網(wǎng)絡(luò)模型在推理計(jì)算時(shí)，是調(diào)用 GPU 的 CUDA 進(jìn)行計(jì)算的，TensorRT 可以真的不同的算法、不同的模型結(jié)構(gòu)、不同的 GPU 平臺(tái)等，進(jìn)行 CUDA 調(diào)整，以保證當(dāng)前模型在特定平臺(tái)上以最優(yōu)的性能計(jì)算。

  假設(shè)在 3090 和 T4 上要分別部署，則需要分別在這兩個(gè)平臺(tái)上進(jìn)行 TensorRT 的轉(zhuǎn)換，然后在對(duì)應(yīng)的平臺(tái)上使用，而不能在相同同的平臺(tái)上轉(zhuǎn)換，在不同的平臺(tái)上使用。

• Dynamic Tensor Memory：

  在每個(gè) tensor 使用期間，TensorRT 會(huì)為其指定顯存，避免顯存重復(fù)申請(qǐng)，減少內(nèi)存占用和提高重復(fù)使用效率

四、ONNX 轉(zhuǎn) 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")

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。