博客 / 詳情

返回

邊緣部署第一章 YOLO如何通過ONNX部署在Jetson orin NX/NANO

最近負責一個公司對內開發的一個項目，幫助庫管做一個可以數零件的桌面軟件，其中就需要將訓練好的模型部署到小型化的嵌入式設備上，公司經費充足，直接給了我一塊Jetson Orin NX 16G版來做邊緣部署平台。根據我的計算20TOPS左右的算力就是足夠的，所以100TOPS的Orin NX 16G性能遠遠溢出。因此，無需考慮性能，考慮到將來可能會有更多的設備型號部署任務，沒有選擇在英偉達GPU上效率最高的TensoerRT，使用兼容性更好的ONNX格式來做模型部署。
這是我第一次做ONNX邊緣部署，不太熟悉整體流程，打算先在服務器上跑通一次Pytorch模型轉換ONNX，再利用ONNX RunTime來推理的全流程，最後根據服務器的環境將Orin NX刷機到合適的版本完成部署。
領導指定要求TensorRT也要部署一下，所以後續會更新一下TensorRT的部署筆記。

參考文章：
非常好的ONNX部署教程

ONNX框架理解

何為ONNX？ONNX是一個與平台無關的格式，由Meta推出的一個開源項目，目的就是將不同格式的深度學習模型轉換為同樣的格式表達即ONNX（Open Neural Network Exchange）。通過導出到ONNX格式，可以顯著提升在CPU上的運行效率，根據YOLO官方的説法，在CPU上使用ONNX的模型推理速度可以提升3倍。
image_4.png
除了推理速度上的提升，ONNX格式還有專屬的ONNX Runtime推理引擎，該推理引擎分CPU和GPU版本，我需要在GPU上運行，所以肯定是安裝GPU版本的。使用ONNX Runtime進行推理能擺脱對Pytorch的依賴，極大的減少了打包後的體積，光Pytorch一個包就用2個多G，而ONNX Runtime GPU只有一百多MB（CPU版本更小），這已經是非常大的提升了。

配置環境

ONNX和ONNX Runtime GPU的環境和CUDA與cuDNN版本有很強的耦合性，通過下面的鏈接來查看對應關係：

https://onnxruntime.ai/docs/execution-providers/CUDA-Executio...

我服務器上的環境通過nvcc -v查看是CUDA Version: 11.8, 注意不要通過nvidia-smi來看哦，那個是驅動最高的支持版本。

cuDNN的版本查看需要去這個目錄下自己去看了，沒有方便的命令，版本如下圖所示，是8.9.2：

image_5.png

到官網看了一下11.8CUDA對應的表格：
image_6.png

我環境上的torch是2.4.1, 感覺適配的應該是1.18.x版本的ONNX Runtime, conda沒這個版本的包，只能用pip了，於是安裝命令如下：

pip install onnxruntime-gpu==1.18.1

還需要選擇ONNX的版本，這裏又涉及到一個概念：ONNX Opset。Opset是ONNX操作集的版本號，不同的Opset版本支持的操作和功能有所不同。每個ONNX版本都有自己的Opset版本，不過所有的ONNX RunTime都具有Opset 7版本以上的向後兼容性，因此二者並不需要完全對應。我就直接不指定版本了。

conda install onnx

模型導出

模型導出非常簡單，用yolo自己帶的api就能輕鬆完成，代碼如下：

from ultralytics import YOLO

model = YOLO("./best.pt")

model.export(format="onnx")

導出後會生成best.onnx文件，接下來就可以用ONNX RunTime來進行推理了。這裏還可以用

https://netron.app

來可視化ONNX模型結構，方便我們進行調試和分析，另外這個網站的前端設計也非常好看，還是開源的，以後可以偷一偷。

ONNX RunTime推理

這部分核心就是對上輸入和輸出的尺寸，可以用上上邊提到的Netron來查看模型的輸入和輸出節點信息如下圖所示：

image_7.png

可以看到我們模型的輸入是1x3x640x640，這裏注意我們需要吧tensor轉換成numpy的array格式才能輸入到ONNX RunTime中，輸出是1x7x8400的維度，代表8400個預測框，每個框有7個值，分別是[batch_index, x1, y1, x2, y2, score, class]。
圖片也需要做歸一化再送進去，完整代碼如下：

import cv2
import onnxruntime
import numpy as np
model_path = "./best.onnx"
# onnxruntime.InferenceSession用於獲取一個 ONNX Runtime 推理器
ort_session = onnxruntime.InferenceSession(model_path)

input_img = "./image.png"
img = cv2.imread(input_img)
# 轉為RGB
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 根據模型要求resize
img = cv2.resize(img, (640, 640))

np_img = img.astype(np.float32) / 255.0  # 轉為float32類型並歸一化
np_img = np.transpose(np_img, (2, 0, 1))
# 增加batch維度
np_img = np.expand_dims(np_img, axis=0)


# 通過 get_inputs() 方法獲取模型的輸入節點信息，並將輸入圖像傳遞給推理器
ort_inputs = {ort_session.get_inputs()[0].name: np_img}
output = ort_session.run(None, ort_inputs)

# 後處理
outputs = output[0]  # shape: (1, 7, 8400)
outputs = np.transpose(np.squeeze(outputs))  # shape: (8400, 7)

boxes = []
scores = []
class_ids = []

# 類別名稱
class_names = ['Gasket', 'Screw', 'Nut']

# 遍歷預測結果
for i in range(outputs.shape[0]):
    classes_scores = outputs[i][4:]
    max_score = np.amax(classes_scores)
    if max_score > 0.5:  # 置信度閾值
        class_id = np.argmax(classes_scores)
        x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
        # 轉換為左上角座標
        left = int(x - w / 2)
        top = int(y - h / 2)
        width = int(w)
        height = int(h)
        
        boxes.append([left, top, width, height])
        scores.append(float(max_score))
        class_ids.append(class_id)

# 非極大值抑制
indices = cv2.dnn.NMSBoxes(boxes, scores, 0.5, 0.45)

if len(indices) > 0:
    indices = np.array(indices).flatten()
    for i in indices:
        box = boxes[i]
        left, top, width, height = box[0], box[1], box[2], box[3]
        
        # 畫框
        cv2.rectangle(img, (left, top), (left + width, top + height), (0, 255, 0), 2)
        
        # 標籤
        label = f"{class_names[class_ids[i]]}: {scores[i]:.2f}"
        cv2.putText(img, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# 轉回BGR顯示
img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite("result_onnx.jpg", img_bgr)
cv2.imshow("Result", img_bgr)
cv2.waitKey(0)
cv2.destroyAllWindows()

查看結果圖片，結果可以正常輸出了：

result_onnx.jpg

接下來就是往Orin NX上邊遷移了。

Orin Nano 與NX刷機

在拿到我們的開發板後，第一件事就是要給開發板刷一個系統，最佳實踐就是通過英偉達提供的開發者套件JetPack來完成。Jetpack中包含英偉達調試好的對應Jetson系列的系統鏡像，CUDA，cuDNN，TensorRT等一系列工具包。

那麼怎樣向開發板刷入Jetpack呢？這裏有兩種方式，第一種是通過英偉達提供的刷機工具SDK Manager來完成，這個工具要求宿主機系統環境嚴格適配要求。第二種方式是直接下載JetPack對應的鏡像文件，然後通過跳線來讓開發板進入刷機模式，再通過命令行工具將鏡像刷入開發板中。我剛好想把公司的服務器重裝一個Ubuntu系統，於是就選擇了第一種方式。

首先通過英偉達的官網找到硬件對應的JetPack版本：
鏈接

image_8.png

如圖所示，我們要使用的Orin系列可以兼容5.1到6.2版本的JetPack。再查看對應JetPack版本的SDK manager的宿主機要求：鏈接
如圖所示，我們儘量選擇新一點的Ubuntu系統（我最開始裝了個18.04，主板上的網卡驅動都不支持，連網都上不了），這裏就選擇了22.04版本的系統。
image_9.png
打開SDK Manager後安裝就比較簡單了，選擇對應的JetPack版本就可以看到對應的組件列表，我這裏選擇了jetpack 6.2.1版本，對應的組件版本可以點what's new查看。
image_11.png
下載好後會彈出一個對話框要求輸入遠程連接的賬號密碼，以及保存各個組件的設備，選擇好後就開始刷機了。
image_10.png
刷好機後，這裏我設置的ip地址都無效了，只能重新搬顯示器鍵盤鼠標過來設置一遍，折騰了好久才把遠程連接弄好，具體過程見我的遠程連接筆記，裏面介紹了rdp的最佳實踐。

接下來就是裝ONNX RunTime了，因為pip 上提供的ONNXRuntime-GPU只支持x86架構的，能夠在jetson上運行的ONNXRuntime-GPU版本需要去jetsonZoo下載對應版本安裝，鏈接如下：jetson zoo。
jetson zoo中最高只提供到6.0.0版本的jetpack支持，想要更高版本的ONNXRuntime GPU需要去下面這個官方源下載：
鏈接

接下來我分別用CPU和GPU都測試了一下，測試結果如下：

CPU 結果
(yolo) lzz@orinnano:~/Desktop/egdeTest$ python runtimeOnnx.py
可用的執行提供程序: ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
注意: 使用 CPU 執行
實際使用的執行提供程序: ['CPUExecutionProvider']
--------------------------------------------------
開始性能測試...
測試時長: 3.23 秒
推理幀數: 4 幀
平均FPS: 1.24 幀/秒
平均單幀耗時: 808.44 毫秒

GPU 結果
開始詳細性能測試...
==================================================
1. 圖像讀取: 18.84 ms
2. 顏色轉換 (BGR->RGB): 2.39 ms
3. 圖像Resize (640x640): 2.88 ms
4. 數據預處理 (歸一化+轉置): 3.46 ms

5. GPU推理測試 (預熱3次, 測試20次):
   平均推理時間: 82.82 ms
   最快推理時間: 80.68 ms
   最慢推理時間: 102.60 ms
   標準差: 5.11 ms
   理論最大FPS: 12.07 FPS

6. 輸出解析 (轉置): 0.07 ms
7. 閾值過濾 (遍歷8400個框): 104.59 ms
   檢測到 160 個候選框
8. NMS非極大值抑制: 0.23 ms
   最終保留 32 個框
9. 繪製結果: 2.32 ms
10. 保存圖像: 4.98 ms

看來TOPS和FLOPS換算關係不是簡單的係數關係啊，如果要達到30FPS的目標，NANO的性能還不怎麼夠啊。

yolo , 邊緣計算 , Python
user avatar
0 位用戶收藏了這個故事！

發佈 評論

Some HTML is okay.