在DeepStream上使用自己的Pytorch模型

Posted By Lin Felix on 9 月 23, 2021 in DeepStream, Edge AI開發地圖, Edge AI開發環境, Jetson Nano, NVIDIA, PyTorch, 教學文 | 0 comments

作者：Felix Lin

DeepStream是NVIDIA專為處理多個串流影像，並進行智慧辨識而整合出的強大工具。開發語言除了原先的C++，從DeepStream SDK 5.1也支援基於原先安裝，再掛上Python套件的方式，讓較熟悉Python程式語言的使用者也能使用DeepStream。

本文主要將其應用在Jetson Nano上，並於DeepStream導入自己的模型執行辨識。

在Jetson Nano上面安裝DeepStream

筆者使用的硬體為Jetson Nano 2GB/4GB，參照官方提供的步驟與對應的版本，幾乎可以說是無痛安裝。對比同樣採用乾淨映像檔，使用源碼或是Docker安裝的JetBot與Jetson Inference要快上許多。

執行官方範例

DeepStream有提供不少範例，不論是從CSI或USB介面的攝影機取得畫面，或是多影像辨識結果顯示，都能經由查看這些範例，學習如何設定。

透過下列指令執行一個配置檔，查看DeepStream是否安裝成功，這個配置檔會開啟一部mp4影片，並模擬產生8個輸入來源，經模型推論處理過後於同一個畫面顯示，點擊單一個區塊可以顯示該來源的詳細資訊。實際應用上可以將各部攝影機的畫面同時輸出並進行處理。

deepstream-app -c source8_1080p_dec_infer-resnet_tracker_tiled_display_fp16_nano.txt

使用自己的模型

如果您與筆者一樣是NVIDIA官方課程小粉絲，從擁有Jetson Nano開始，就按部就班的跟著課程學習，那您一定看過下列三種不同主題的課程。

從入門的Section 1開始，到執行Section 2的JetBot自駕車專案，最後Section 3 Hello AI World。經過三個Section，您應該稍微熟悉Pytorch，並且也訓練了不少自己的模型，特別是在Hello AI World有訓練了Object Detection模型。既然都有自己的模型，何不放到DeepStream上面製作專屬的串流辨識專案，針對想要辨識的項目導入適合的模型。

在Hello AI World專案訓練Object Detection模型的時候，我們使用的是SSD-Mobilenet，在DeepStream的物件辨識範例中有提供使用自己的SSD模型方法，可在下列路徑找到參考文件，文件中使用的例子是使用coco資料集預訓練的SSD-Inception。

/opt/nvidia/deepstream/deepstream-5.1/sources/objectDetector_SSD

可惜的是文件中使用的是從Tensorflow訓練的模型，經由轉換.uff再餵給DeepStream，與官方課程使用的Pytorch是不同路線。筆者在網上尋找解決方法，看是否有DeepStream使用Pytorch模型的方案，也於NVIDIA開發者論壇找到幾個同樣的提問，但最終都是導到上述提到的參考文件。

從Hello AI World訓練的Object Detection模型，經過執行推論的步驟，您應該會有三個與模型有關的檔案，分別是用Pytorch訓練好的.pth，以及為了使用TensorRT加速而將.pth轉換的.onnx，最後是執行過程中產生的.engine。既然Pytorch模型找不到解決方案，那就從ONNX模型下手吧，所幸經過一番折騰，終於讓筆者找到方法。

neilyoung提供的方法主要是能產生動態函式庫，以便我們能在DeepStream使用ONNX模型，除了準備好自己訓練的ONNX模型檔案與Labels檔案，只要再新增設定模型路徑與類型的config檔案，與deepstream配置檔就能實現使用自己的模型進行推論囉！

STEP 1：
首先於以下路徑底下新增執行ONNX專案的資料夾，筆者命名為objectDetector_ONNX。
/opt/nvidia/deepstream/deepstream-5.1/sources
STEP 2：
新增專案資料夾後，請clone方才的nvdsinfer_custom_impl_onnx專案到資料夾內。
STEP 3：
打開Terminal進到nvdsinfer_custom_impl_onnx專案裡面，透過sudo make指令產生動態函式庫。
STEP 4：
接著將自己從Hello AI World專案訓練的Object Detection模型與標籤複製到objectDetector_ONNX專案資料夾。
STEP 5：
從別的專案資料夾複製config檔案與deepstream配置檔到我們的資料夾內，這邊複製objectDetector_SSD，因為模型類型相近，只要稍微修改即可。
STEP 6：
首先修改config檔案，如下圖所示，將模型路徑與標籤路徑，修正為自己的模型與標籤名稱，engine檔案的部份與Hello AI World專案一樣，在執行ONNX檔案進行TensorRT加速時會自動產生，只需給路徑與名稱即可。對於classes的部份，切記在Hello AI World專案訓練的模型會加上BACKGROUND這一個類別，所以若是您辨識的物件有三種，就得在classes這邊填上3+1。
下方三項的設定則依照nvdsinfer_custom_impl_onnx專案github上的說明，記得動態函式庫的路徑請改成自己的路徑。

output-blob-names="boxes;scores"
parse-bbox-func-name="NvDsInferParseCustomONNX"
custom-lib-path="/path/to/lib/libnvdsinfer_custom_impl_onnx.so"

接著依照個人需求設定辨識的參數，例如希望信心指數達多少%才認定物件類別，可以修改threshold。
STEP 7：
接著修改deepstream配置檔，筆者在這邊設定為USB Webcam輸入，並輸出單一視窗顯示，除了正常調整輸入與輸出之外，請將config檔案與Labels檔案導引至自己的路徑，engine的部份與config設定相同即可，如下圖所示。