TensorRT 的核心是一個 c++ 的 library，透過 TensorRT 將 training framework 最佳化成一個 inference engine，這個 engine 能夠高效率的於 Nvidia GPU 進行 inference。
如今 TensorRT 已經支援了很多深度學習的框架，但是有些框架需先轉換成 ONNX 的通用深度學習模型，才可以透過 TensorRT 進行最佳化，像是 Caffe2/CNTK/Chainer/PyTorch/MxNet 等等，而有些框架則已經把 TensorRT 集成到框架中了，像是TensorFlow/MATLAB 等等。
為了最佳化模型需先經過 TensorRT 的 network definition 再 build 成 engine，這個 build 的過程會需要一段時間，因此可以透過 serialize 將 engine 存成 trt 檔，下次若要使用可以透過 deserialize 將 trt 檔還原成 engine 以供預測。
TensorRT 主要做了下面幾件事，来提升模型的運行速度：
Precision Calibration
TensorRT 支持 FP16 和 INT8 的精度，我們知道深度學習在訓練時的精度一般都是FP32，而我們透過降低權重的精度已達到加速推論的目的，而 TensorRT INT8 需經過特殊的量化處理才能保證其準確度。
Layer & Tensor Fusion
TensorRT 對訓練網路的結構進行重構，把一些能夠合併的運算合併在一起，針對 GPU 的特性做了優化。
Kernel Auto-Tuning
模型在推論計算時是調用 GPU 的 CUDA 核進行計算的，而 TensorRT 可以針對不同的算法、網路架構、GPU 平台進行 CUDA 核的調整，以保證當前模型能在特定平台上以最優的性能進行推論。
Dynamic Tensor Memory
在每個 tensor 的使用期間，TensorRT 會為其指定顯存，避免顯存重複申請，減少內存占用和提高重複使用效率。
Multi-Stream Execution
經過 Layer & Tensor Fusion 後，可以看到兩個框框彼此之間是不相關的，因此我們可以單獨啟用兩個計算流 (stream)，分別運算。
以下為將 ONNX 或 TensorFlow 轉換成 TensorRT 的範例，更多範例可以參考我整理的 https://github.com/d246810g2000/tensorrt：
1. ONNX to TensorRT
01_train.py
from tensorflow.keras.applications.mobilenet import MobileNet
from tensorflow.keras import backend
backend.set_image_data_format('channels_first')
print(backend.image_data_format())
model = MobileNet(weights='imagenet',
include_top=True,
input_shape=(3, 224, 224))
# save to keras model
model.save('mobilenet.h5')02_keras2onnx.py
import onnx
from tensorflow.keras.models import load_model
import onnxmltools
keras_path = './Model/mobilenet.h5'
onnx_path = './Model/mobilenet.onnx'
keras_model = load_model(keras_path)
onnx_model = onnxmltools.convert_keras(keras_model)
onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1
onnx_model.graph.output[0].type.tensor_type.shape.dim[0].dim_value = 1
onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, onnx_path)03_onnx2trt.py
import os
import tensorrt as trt
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)onnx_file = 'mobilenet.onnx'
trt_file = 'mobilenet.trt'
batch_size = 1"""Takes an ONNX file and creates a TensorRT engine to run inference with"""
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
builder.max_workspace_size = 1<<28 # 256MiB
builder.max_batch_size = batch_size
builder.fp16_mode = True # fp32_mode -> False
# Parse model file
with open(onnx_file, 'rb') as model:
print('Beginning ONNX file parsing')
if not parser.parse(model.read()):
print ('ERROR: Failed to parse the ONNX file.')
for error in range(parser.num_errors):
print (parser.get_error(error))
print('Completed parsing of ONNX file')
engine = builder.build_cuda_engine(network)
print("Completed creating Engine")
with open(trt_file, "wb") as f:
f.write(engine.serialize())
04_inference.py
import common
import tensorflow as tf
import tensorrt as trt
import numpy as np
from tensorflow.keras.applications.mobilenet import decode_predictions
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)x = tf.random.uniform([1, 3, 224, 224])# load trt engine
trt_path = 'mobilenet.trt'
with open(trt_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())# 分配 buffers 給 inputs 和 outputs
inputs, outputs, bindings, stream = common.allocate_buffers(engine)
inputs[0].host = tf.reshape(x, [-1])# inference
with engine.create_execution_context() as context:
trt_outputs = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
preds = trt_outputs[0].reshape(1, -1)
print('Predicted: {}'.format(decode_predictions(preds, top=3)[0]))
common.py 為 NVIDIA Jetson Nano 內建的 code，檔案路徑為 /usr/src/tensorrt/samples/python/common.py，也可以在這邊找到 https://github.com/d246810g2000/tensorrt/blob/main/common.py
TF -> ONNX -> TRT
- 優點：消耗更少的內存，推論速度通常比 TF-TRT 快
- 缺點：若遇到某些圖層不支持 TRT，我們必須為這些圖層使用插件 (plugin) 或自定義代碼來實現
2. TensorFlow to TensorRT
01_trian.py
import tensorflow as tf
from tensorflow.keras.applications.mobilenet import MobileNet
from tensorflow.keras import backend
backend.set_image_data_format('channels_first')
print(backend.image_data_format())
model = MobileNet(weights='imagenet',
include_top=True,
input_shape=(3, 224, 224))
# save to saved_model
tf.saved_model.save(model, 'mobilenet_saved_model')02_tf2trt.py
import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt
saved_file = 'mobilenet_saved_model'
TRT_file = 'mobilenet_saved_model_TFTRT_FP16'
print('Converting to TF-TRT FP16...')
conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS
conversion_params = conversion_params._replace(
max_workspace_size_bytes=(1<<28),
precision_mode="FP16",
maximum_cached_engines=1)
converter = trt.TrtGraphConverterV2(
input_saved_model_dir=saved_file,
conversion_params=conversion_params)
converter.convert()
converter.save(output_saved_model_dir=TRT_file)
print('Done Converting to TF-TRT FP16')03_inference.py
import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt
from tensorflow.python.saved_model import tag_constants
from tensorflow.keras.applications.mobilenet import preprocess_input, decode_predictionsx = tf.random.uniform([1, 3, 224, 224])input_saved_model = 'mobilenet_saved_model_TFTRT_FP16'
saved_model_loaded = tf.saved_model.load(
input_saved_model,
tags=[tag_constants.SERVING])
signature_keys = list(saved_model_loaded.signatures.keys())
print(signature_keys)infer = saved_model_loaded.signatures['serving_default']
print(infer.structured_outputs)
keys = list(infer.structured_outputs.keys())preds = infer(x)[keys[0]].numpy()
print('Predicted: {}'.format(decode_predictions(preds, top=3)[0]))
TF -> TRT
- 優點：API 易於使用，若遇到某些圖層不支持 TRT 會自動由 TF 執行，無需擔心插件 (plugin)。
- 缺點：需要將整個 TF 庫存儲在 HDD 中（這對部署環境不利）；需要在運行時將 TF 加載到內存中，通常比純 TRT 引擎運行慢。
