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

*本文由RS components 贊助發表,轉載自DesignSpark部落格原文連結

使用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

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

可以看到已經有一個 “ 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部落格原文連結 ()

發佈留言

發佈留言必須填寫的電子郵件地址不會公開。 必填欄位標示為 *