嘿!您似乎在 United States,您想使用我们的 English 网站吗?
Switch to English site
Skip to main content

NVIDIA CUDA核心GPU實做:Jetson Nano 運用TensorRT加速引擎 – 下篇

使用TensorRT運行ONNX

要來運行TensorRT了,來複習一下TensorRT的流程:

  1. ONNX parser:將模型轉換成ONNX的模型格式。
  2. Builder:將模型導入TensorRT並且建構TensorRT的引擎。
  3. Engine:引擎會接收輸入值並且進行Inference跟輸出。
  4. 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

執行的結果如下,其實沒有耗費很多時間:

01_build_trt_engine_9678fb620887d74185c2a95cfa7f6cfd36a173bd.png

可以看到已經有一個 “ Alexnet.trt “ 生成出來了,或許因為經過序列化處理,所以檔案少了非常多的容量:

02_file_size_1d1fa661cb7677023f5c9715ad50d20c2e0ebb34.png

接下來就是重頭戲了,使用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)]}")

03_trt_infer_0a0dbeac84d1f45edf6d99f487488d29f94d5652.png

接著稍微比較了所有的框架,但可能因為同時進行三種不同的框架,Jetson Nano負荷不來所以時間都被拉長了,但我們仍然可以從比例看的出來彼此之間的差距,PyTorch:ORT:TRT大約是7:2:1,代表運行PyTorch所耗費的時間是開啟TRT引擎的7倍!

04_Compare_AI_framework_f3a680abaf895a4a7809f2fcc947024e568ad6a6.png

從下圖可以看到蠻有趣的一點是理論上精度最高的是PyTorch的版本,結果在信心指數的部分卻是最低的。

05_CMD_Compare_AI_framework_f7127f79c6f07e927a5c8e13477d87e28c0e43d1.png

PyTorch使用TensorRT最簡單的方式

torch2trt套件

接下來來介紹一下這個套件,可以只用少少的程式碼將PyTorch的模型轉換成trt的模型,對於torch的愛好者來說實在是太棒了:

0551_bdf5d86f7f1bdb7d694119a43aa2fa8eeeea44e8.png

安裝方法如下,這邊我是採用我的映像檔並且開啟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:

06_torch2trt_convert_c2596ba6e389b31efd145f7a1963edba46d03195.png

接著我稍微改動了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))

獲得的結果如下:

07_torch2trt_test_6ba5cfd043f82d6244174b677547b91a948bd418.png

接著我們來嘗試運作物件辨識的範例吧!我第一個想到的就是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

08_yolov5_export_f339214b5c45d9a42b3fe3879101a83fbc5a3dfd.png

預設是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' 

09_tensorrtx_yolov5_006462032ebab0f9a101e805c8af6684b4c607fb.png

建構完之後就可以直接透過下列指令執行,會輸出圖檔 _bus.jpg、_zidane.jpg:

(yolov5) $ sudo ./yolov5 -d  ~yolov5/data/images  

10_tensorrtx_compiler_a6d194147f796c5da077d24b049bd43a8d6fb25d.png

接下來還可以使用Python來進行Inference,需安裝tensorrt跟pycuda:

(yolov5) $ pip install pycuda
(yolov5) $ python yolov5_trt.py

執行結果如下:

11_tensorrtx_infer_fb8fe408719ba3c4707b35c76d1183cd6401d079.png

接著使用原生的yolov5進行inference,耗費時間約為0.549s、0.244s:

12_yolov5_infer_159c2cc9c7fa9e7079f0ec43fc4917ec191817c3.png

不曉得為什麼跑了那麼多範例,yolov5的速度沒有提升反而下降,這個部分還需要研究一下…如果廣大讀者們知道的話麻煩在留言告訴我~

結語

從圖片分類的範例來看的話,TensorRT還是非常的厲害的!但是目前支援度還是不太高,入門的難易度很高,所以如果有要使用自己的模型要好好研究一番,但如果是用GPU模型並且是直接使用常見模型的話TensorRT絕對一個大趨勢,畢竟加速的效果真的不錯。

參考文章

What is the difference between tflite and tensorRT?

TensorRT(1)-介绍-使用-安装

TensorRT-Optimization-Principle

Speeding up Deep Learning Inference Using TensorFlow, ONNX, and TensorRT

CAVEDU Education is devoted into robotics education and maker movement since 2008, and is intensively active in teaching fundamental knowledge and skills. We had published many books for readers in all ages, topics including Deep Learning, edge computing, App Inventor, IoT and robotics. Please check CAVEDU's website for more information: http://www.cavedu.com, http://www.appinventor.tw
DesignSpark Electrical Logolinkedin