專案使用 YOLOv8 物件檢測模型部屬在邊緣裝置上,在了解模型與測試後發現仍然有可以優化地方存在,而關鍵在於 YOLOv8 在 Decode bbox 過程中是將每一個特徵網格(grid) 的每一個 bbox 都進行解碼,而這部分能夠將其移除,並在邊緣端進行實作,進而減少 decode box 時間
在將影像輸入模型並進行後續操作後,針對不同部署模型格式 (如 ONNX、TensorRT、TFLite 等) 進行 export 時,後處理操作可能會影響計算圖 (computational graph) 的生成,而 decode bbox 過程也包含在裡面,因此我們需要將其拔除並另外實現在邊緣裝置上
我們從 YOLOv8 (現行為 YOLOv11 了) 來進行部分程式碼的解讀,參考 ultralytics
的 github 開源程式碼 ultralytics/nn/modules/head.py
在 Detect
類別中負責了以下的任務
雖然 YOLOv8 是 anchor-free 的模型,但實際上訓練與 inference 也都需要 anchor
self.cv2
self.cv3
self.reg_max
self.cv2 = nn.ModuleList(
nn.Sequential(
Conv(x, c2, 3),
Conv(c2, c2, 3),
nn.Conv2d(c2, 4 * self.reg_max, 1)
) for x in ch
)
self.cv3 = nn.ModuleList(
nn.Sequential(
nn.Sequential(DWConv(x, x, 3), Conv(x, c3, 1)),
nn.Sequential(DWConv(c3, c3, 3), Conv(c3, c3, 1)),
nn.Conv2d(c3, self.nc, 1),
) for x in ch
)
在 YOLOv8 的架構中,forward
函數需要處理來自不同尺度的特徵圖,而這些特徵圖被儲存在一個列表 x
中。這些特徵圖分過 self.cv2
和 self.cv3
層進行向前傳播,可以看出 YOLOv8 在輸出預測上進行了 解偶 (decoupling) 的
特徵圖 x[i]
經過 self.cv2[i]
和 self.cv3[i]
的處理後,tensor 的形狀如下:
self.cv2[i](x[i])
的 tensor shape 為 self.cv3[i](x[i])
的 tensor shape 為 最終,特徵圖 x[i]
的 tensor shape 會是
這樣,我們就完成了第一階段的 tensor 處理。由於這過程涉及許多 tensor 操作,對初學者來說可能較難理解,但基本上是將不同的預測任務(如邊界框和分類)進行分離後,再合併回輸出中。
def forward(self, x):
"""Concatenates and returns predicted bounding boxes and class probabilities."""
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.training: # Training path
return x
y = self._inference(x)
return y if self.export else (y, x)
_inference
在 _inference
函數中,負責處理 decode bbox與 anchor 處理,我們將核心的程式碼來進行解讀
x_cat
make_anchors(feats, strides, grid_cell_offset)
feat[0][0]
生成的話就會是在特徵點的中心有一個 anchor,strides
則是相對於原圖其特徵圖縮小的倍率,通常就是 Detect
的特徵圖相對於原圖的縮放比例self.dlf(x)
reshape() -> permute()
積分操作藉由 Conv 1x1
來完成,等價於積分操作
def _inference(self, x):
x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
self.anchors, self.strides = (
x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5)
)
self.shape = shape
box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides
return torch.cat((dbox, cls.sigmoid()), 1)
藉由 NCNN 將模型部屬,並使用 C++ 進行撰寫
model.export(format="ncnn", dynamic=True, simplify=True)
將模型轉換## ultralytics/nn/modules/head.py 58-68
def forward(self, x):
"""Concatenates and returns predicted bounding boxes and class probabilities."""
batch_size = x[0].shape[0]
for i in range(self.nl):
# x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
box = self.cv2[i](x[i]).view(batch_size, self.reg_max * 4, -1)
cls = self.cv3[i](x[i]).view(batch_size, self.nc, -1)
x[i] = torch.cat((box, cls), 1)
if self.training: # Training path
return x
# y = self._inference(x)
y = torch.cat(x, dim=2)
return y if self.export else (y, x)
在export後會得到4個檔案,我們需要的是.bin
與.param
預定義的框,用於協助模型學習,讓模型學習的目標是如何調整預定義的框,以去符合標籤所標記的框
YOLOv8 Anchor-Free 模型預測偏移量(
模型輸出參數:
邊界框計算公式:
由於 YOLOv8 是密集地在一張影像上進行bbox的預測,因此需要透過nms來消除分數較低以及重複預測的bbox
在不更改原始程式碼的情況下,我們可以透過 export.py
來將 forward
函數進行改寫後,替換掉原生 Detect
的類別,這樣我們便可以將bbox decode的過程省略掉
class PatchDetect(nn.Module):
def forward(self, x):
"""Concatenates and returns predicted bounding boxes and class probabilities."""
print("Use Patch Forward!")
if self.end2end:
return self.forward_end2end(x)
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.training: # Training path
return x
y = torch.cat([xi.view(x[0].shape[0], self.no, -1) for xi in x], 2)
# y = self._inference(x)
return y # i self.export else (y, x)
# Replace original Package Detect head class
ultralytics.nn.modules.head.Detect = PatchDetect
from ultralytics import YOLO
import argparse
def run(args):
if args.patch:
import export_patch
model = YOLO(args.path)
model.export(format="ncnn", dynamic=True, simplify=True)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run YOLO model with optional patch.")
parser.add_argument('--path', type=str, help="Export Model Path")
parser.add_argument('--patch', action='store_true', help='Enable export patch')
args = parser.parse_args()
run(args)
Export model時不包含decode box的功能
python .\export_ncnn.py --path best.pt --patch
Export model時包含decode box的功能
python .\export_ncnn.py --path best.pt