*本文由RS components 贊助發表，轉載自DesignSpark部落格原文連結
使用TensorRT運行ONNX
要來運行TensorRT了，來複習一下TensorRT的流程：
- ONNX parser：將模型轉換成ONNX的模型格式。
- Builder：將模型導入TensorRT並且建構TensorRT的引擎。
- Engine：引擎會接收輸入值並且進行Inference跟輸出。
- Logger：負責記錄用的，會接收各種引擎在Inference時的訊息。
第一個我們已經完成了，接下來的部分要建構TensorRT引擎，這個部分可以參考於NVIDIA的官網文件，主要程式如下，總共三個副函式build_engine、save_engine、load_engine，就如字面上的意思一樣是建置、儲存、載入，而log函式庫是我自己寫的用來顯示狀態以及計時：
import tensorrt as trt
from log import timer, logger
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt_runtime = trt.Runtime(TRT_LOGGER)
if __name__ == "__main__":
onnx_path = 'alexnet.onnx'
trt_path = 'alexnet.trt'
input_shape = [1, 224, 224, 3]
build_trt = timer('Parser ONNX & Build TensorRT Engine')
engine = build_engine(onnx_path, input_shape)
build_trt.end()
save_trt = timer('Save TensorRT Engine')
save_engine(engine, trt_path)
save_trt.end()
build_engine的程式碼，max_workspace是指GPU的使用暫存最大可到多少，因為TensorRT可以支援到Float Point 16，這次模式選擇fp16，在建構引擎之前需要先解析 (Parser) ONNX模型，接著再使用build_cuda_engine來建構：
def build_engine(onnx_path, shape = [1,224,224,3]):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(1) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
builder.max_workspace_size = (256 << 20)
# 256MiB
builder.fp16_mode = True
# fp32_mode -> False
with open(onnx_path, 'rb') as model:
parser.parse(model.read())
engine = builder.build_cuda_engine(network)
return engine
save_engine的程式碼，儲存的時候需要將引擎給序列化以供儲存以及載入：
def save_engine(engine, engine_path):
buf = engine.serialize()
with open(engine_path, 'wb') as f:
f.write(buf)
load_engine的程式碼，主要在於要反序列化獲得可運行的模型架構：
def load_engine(trt_runtime, engine_path):
with open(engine_path, 'rb') as f:
engine_data = f.read()
engine = trt_runtime.deserialize_cuda_engine(engine_data)
return engine
執行的結果如下，其實沒有耗費很多時間：
![]()
可以看到已經有一個 “ Alexnet.trt “ 生成出來了，或許因為經過序列化處理，所以檔案少了非常多的容量：
![]()
接下來就是重頭戲了，使用TensorRT進行Inference，先導入函式庫，這邊要注意common是從 /usr/src/tensorrt/samples/python/common.py複製出來的，engine是剛剛建構引擎的程式，log是我另外寫的用來計時跟顯示，其中trt的logger跟runtime也都先定義好了方便之後的呼叫：
import tensorrt as trt
from PIL import Image
import torchvision.transforms as T
import numpy as np
import common
from engine import load_engine
from log import timer, logger
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt_runtime = trt.Runtime(TRT_LOGGER)
載入資料的副函式，這邊需要轉成numpy格式，因為trt的引擎只吃numpy：
def load_data(path):
trans = T.Compose([
T.Resize(256), T.CenterCrop(224), T.ToTensor()
])
img = Image.open(path)
img_tensor = trans(img).unsqueeze(0)
return np.array(img_tensor)
接著載入引擎並且分配內存，binding是存放input、output所要佔用的空間大小，stream則是pycuda的東西 ( cuda.Stream() )，是cuda計算缺一不可的成員；這邊將inputs的內容替換成我要inference的照片，.host的意思是input的內存空間。
# load trt engine
load_trt = timer("Load TRT Engine")
trt_path = 'alexnet.trt'
engine = load_engine(trt_runtime, trt_path)
load_trt.end()
# allocate buffers
inputs, outputs, bindings, stream = common.allocate_buffers(engine)
# load data
inputs[0].host = load_data('../test_photo.jpg')
推論的部分則是透過common.do_inference來進行，create_execution_context是必要的旦沒有開源所以不太清楚裡面的內容，：
# inference
infer_trt = timer("TRT 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]
infer_trt.end()
最後取得標籤以及對應的信心指數：
# Get Labels
f = open('../imagenet_classes.txt')
t = [ i.replace('\n','') for i in f.readlines()]
logger(f"Result : {t[np.argmax(preds)]}")
![]()
接著稍微比較了所有的框架，但可能因為同時進行三種不同的框架，Jetson Nano負荷不來所以時間都被拉長了，但我們仍然可以從比例看的出來彼此之間的差距，PyTorch:ORT:TRT大約是7:2:1，代表運行PyTorch所耗費的時間是開啟TRT引擎的7倍！
![]()
從下圖可以看到蠻有趣的一點是理論上精度最高的是PyTorch的版本，結果在信心指數的部分卻是最低的。
![]()
PyTorch使用TensorRT最簡單的方式
torch2trt套件
接下來來介紹一下這個套件，可以只用少少的程式碼將PyTorch的模型轉換成trt的模型，對於torch的愛好者來說實在是太棒了：
![]()
安裝方法如下，這邊我是採用我的映像檔並且開啟torch虛擬環境：
$ git clone https://github.com/NVIDIA-AI-IOT/torch2trt
$ cd torch2trt
$ workon torch
(torch) $ python3 setup.py install
torch2trt 將 PyTorch的模型轉換成tensorrt
使用方法如下，透過torch2trt就可以直接轉換，一樣需要宣告輸入的維度大小；接著也可以使用pytorch的方式進行儲存；導入則需要使用TRTModule：
import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet
# Load Alexnet Model
model = alexnet(pretrained=True).eval().cuda()
# TRT Model
x = torch.ones((1, 3, 224, 224)).cuda()
model_trt = torch2trt(model, [x])
# Save Model
alexnet_trt_pth = 'alexnet_trt.pth'
torch.save(model_trt.state_dict(), alexnet_trt_pth)
# Load Model
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict( torch.load('alexnet_trt.pth'))
完整程式碼如下：
import torch
from torch2trt import torch2trt
from torchvision import transforms as T
from torchvision.models.alexnet import alexnet
import time
# Use to print info and timing
from print_log import log
# Load Model
alexnet_pth = 'alexnet.pth'
load_model = log("Load Model...")
model = alexnet(pretrained=True).eval().cuda()
torch.save(model.state_dict(), alexnet_pth, _use_new_zipfile_serialization=False)
load_model.end()
# TRT Model
convert_model = log("Convert Model...")
x = torch.ones((1, 3, 224, 224)).cuda()
model_trt = torch2trt(model, [x])
convert_model.end()
# Save Model
alexnet_trt_pth = 'alexnet_trt.pth'
save_model = log("Saving TRT...")
torch.save(model_trt.state_dict(), alexnet_trt_pth)
save_model.end()
在 JetsonNano上的運作結果如下，轉換的時間為127秒左右，儲存的時間為160秒左右，基本上時間會因各個裝置不同而改變，執行完之後就會看到多了兩個檔案 alexnet以及alexnet_trt：
![]()
接著我稍微改動了Github的範例，使用自己的貓咪照片來當作預測資料，首先載入模型、資料、標籤檔：
import torch
import torch.nn as nn
from torchvision import transforms as T
from torchvision.models.alexnet import alexnet
from torch2trt import torch2trt
from torch2trt import TRTModule
import os
import cv2
import PIL.Image as Image
import time
# Use to print info and timing
from print_log import log
def load_model():
model_log = log('Load {} ... '.format('alexnet & tensorrt'))
model = alexnet().eval().cuda()
model.load_state_dict(torch.load('alexnet.pth'))
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))
model_log.end()
return (model, model_trt)
def load_data(img_path):
data_log = log('Load data ...')
img_pil = Image.open(img_path)
trans = T.Compose([T.Resize(256),T.CenterCrop(224), T.ToTensor()])
data_log.end()
return trans(img_pil).unsqueeze(0).cuda()
def load_label(label_path):
f = open( label_path, 'r')
return f.readlines()
接著先寫好了inference的過程：
def infer(trg_model, trg_label, trg_tensor, info = 'Normal Model'):
softmax = nn.Softmax(dim=0)
infer_log = log('[{}] Start Inference ...'.format(info))
with torch.no_grad():
predict = trg_model(trg_tensor)[0]
predict_softmax = softmax(predict)
infer_log.end()
label = trg_label[torch.argmax(predict_softmax)].replace('\n',' ')
value = torch.max(predict_softmax)
return ( label, value)
最後就是整個運作的流程了：
if __name__ == "__main__":
# Load Model
model, model_trt = load_model()
# Input Data
img_path = 'test_photo.jpg'
img_tensor = load_data(img_path)
# Label
label_path = 'imagenet_classes.txt'
labels = load_label(label_path)
# Predict : Normal
label, val = infer(model, labels, img_tensor, "Normal AlexNet")
print('\nResult: {} {}\n'.format(label, val))
# Predict : TensorRT
label_trt, val_trt = infer(model_trt, labels, img_tensor, "TensorRT AlexNet")
print('\nResult: {} {}\n'.format(label_trt, val_trt))
獲得的結果如下：
![]()
接著我們來嘗試運作物件辨識的範例吧！我第一個想到的就是torchvision中常見的物件辨識fasterrcnn，但是他轉換成trt過程問題很多，從這點也看的出來trt的支援度還沒有非常高，常常因為一些沒支援的層而無法轉換這時候你就要自己去重新定義，對於新手而言實在是非常辛苦，所以Fasterrcnn這部分我就先跳過了，轉戰YOLOv5去嘗試運行TensorRT看看。
YOLOv5使用TensorRT引擎方式
其實排除上述介紹的簡單方式，正規的方式應該是先轉成ONNX再轉成TensorRT，其中yolov5就有提供轉換成ONNX的方式
1. 轉換成 ONNX格式再導入TensorRT ( 僅到匯出 )
這邊我們直接使用YOLOv5來跑範例，先下載YOLOv5的Github並且開啟虛擬環境，如果你是使用自己的環境則可以透過安裝YOLOv5相依套件來完成：
$ git clone https://github.com/ultralytics/yolov5
$ cd yolov5
$ workon yolov5
Nano 上安裝ONNX、coremltools：
$ sudo apt-get install protobuf-compiler libprotoc-div
$ pip install onnx
$ pip install coremltools==4.0
進行ONNX的轉換：
$ source ~/.bashrc
$ workon yolov5
(yolov5) $ python models/export.py
![]()
預設是yolov5s.pt，執行完之後可以發現多了yolov5s.onnx以及yolov5.torchscript.pt，接著就可以使用ONNX的方法去導入使用，不過在yolo系列很少人會這樣做，主要是因為yolo有自定義的層，可能會導致trt無法轉換，但是也因為yolo已經很出名了，所以轉換的部分已經有人整合得很好，可以直接拿來使用。
1. 使用tensorrtx直接轉換 ( 可執行 )
YOLOv5有提供直接從.pt轉成trt的方式，這時候就要參考另外一個github了 https://github.com/wang-xinyu/tensorrtx，裡面有個yolov5的資料夾，按照ReadMe進行即可完成轉換並且進行inferece：
1.複製兩個Github並複製py到yolov5資料夾，運行gen_wts.py來生成yolov5s.wpt，這個是讓TensorRT引擎運行的權重檔。
(yolov5) $ git clone https://github.com/wang-xinyu/tensorrtx.git
(yolov5) $ git clone https://github.com/ultralytics/yolov5.git
(yolov5) $ cd yolov5
(yolov5) $ cp yolov5s.pt weights/
(yolov5) $ cp ~/tensorrtx/yolov5/gen_wts.py .
(yolov5) $ python gen_wts.py
2.由於tensorrt引擎是基於C++，所以常常會使用cmake來建構成執行檔。
(yolov5) $ cd ~/tensorrtx/yolov5
(yolov5) $ cp ~/yolov5/yolov5s.wts .
(yolov5) $ mkdir build
(yolov5) $ cd build
(yolov5) $ cmake ..
(yolov5) $ make
(yolov5) $ sudo ./yolov5 -s // serialize model to plan file i.e. 'yolov5s.engine'
![]()
建構完之後就可以直接透過下列指令執行，會輸出圖檔 _bus.jpg、_zidane.jpg：
(yolov5) $ sudo ./yolov5 -d ~yolov5/data/images
![]()
接下來還可以使用Python來進行Inference，需安裝tensorrt跟pycuda：
(yolov5) $ pip install pycuda
(yolov5) $ python yolov5_trt.py
執行結果如下：
![]()
接著使用原生的yolov5進行inference，耗費時間約為0.549s、0.244s：
![]()
不曉得為什麼跑了那麼多範例，yolov5的速度沒有提升反而下降，這個部分還需要研究一下…如果廣大讀者們知道的話麻煩在留言告訴我～
結語
從圖片分類的範例來看的話，TensorRT還是非常的厲害的！但是目前支援度還是不太高，入門的難易度很高，所以如果有要使用自己的模型要好好研究一番，但如果是用GPU模型並且是直接使用常見模型的話TensorRT絕對一個大趨勢，畢竟加速的效果真的不錯。
參考文章
*本文由RS components 贊助發表，轉載自DesignSpark部落格原文連結 (本篇文章完整範例程式請至原文處下載)
