YOLOv8模型pytorch格式轉為onnx格式的步驟詳解
更新時間：2024年12月02日 08:52:09 作者：相識已是上上簽
這篇文章主要介紹了YOLOv8模型pytorch格式轉為onnx格式的相關資料,本文介紹了YOLOv8的Pytorch網(wǎng)絡結構和轉換為ONNX的過程,包括自定義轉換和使用EfficientNMS_TRT插件進行后處理優(yōu)化,需要的朋友可以參考下
一、YOLOv8的Pytorch網(wǎng)絡結構
model DetectionModel(
(model): Sequential(
(0): Conv(
(conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(act): SiLU(inplace=True)
)
(1): Conv(
(conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(act): SiLU(inplace=True)
)
(2): C2f(
(cv1): Conv(
(conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(320, 128, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): ModuleList(
(0-2): 3 x Bottleneck(
(cv1): Conv(
(conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
)
)
)
(3): Conv(
(conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(act): SiLU(inplace=True)
)
(4): C2f(
(cv1): Conv(
(conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): ModuleList(
(0-5): 6 x Bottleneck(
(cv1): Conv(
(conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
)
)
)
(5): Conv(
(conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(act): SiLU(inplace=True)
)
(6): C2f(
(cv1): Conv(
(conv): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): ModuleList(
(0-5): 6 x Bottleneck(
(cv1): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
)
)
)
(7): Conv(
(conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(act): SiLU(inplace=True)
)
(8): C2f(
(cv1): Conv(
(conv): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(1280, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): ModuleList(
(0-2): 3 x Bottleneck(
(cv1): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
)
)
)
(9): SPPF(
(cv1): Conv(
(conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1, ceil_mode=False)
)
(10): Upsample(scale_factor=2.0, mode='nearest')
(11): Concat()
(12): C2f(
(cv1): Conv(
(conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(1280, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): ModuleList(
(0-2): 3 x Bottleneck(
(cv1): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
)
)
)
(13): Upsample(scale_factor=2.0, mode='nearest')
(14): Concat()
(15): C2f(
(cv1): Conv(
(conv): Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(640, 256, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): ModuleList(
(0-2): 3 x Bottleneck(
(cv1): Conv(
(conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
)
)
)
(16): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(act): SiLU(inplace=True)
)
(17): Concat()
(18): C2f(
(cv1): Conv(
(conv): Conv2d(768, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(1280, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): ModuleList(
(0-2): 3 x Bottleneck(
(cv1): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
)
)
)
(19): Conv(
(conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(act): SiLU(inplace=True)
)
(20): Concat()
(21): C2f(
(cv1): Conv(
(conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(1280, 512, kernel_size=(1, 1), stride=(1, 1))
(act): SiLU(inplace=True)
)
(m): ModuleList(
(0-2): 3 x Bottleneck(
(cv1): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(cv2): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
)
)
)
(22): PostDetect(
(cv2): ModuleList(
(0): Sequential(
(0): Conv(
(conv): Conv2d(256, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(1): Conv(
(conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
)
(1-2): 2 x Sequential(
(0): Conv(
(conv): Conv2d(512, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(1): Conv(
(conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
(cv3): ModuleList(
(0): Sequential(
(0): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(1): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(2): Conv2d(256, 35, kernel_size=(1, 1), stride=(1, 1))
)
(1-2): 2 x Sequential(
(0): Conv(
(conv): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(1): Conv(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(act): SiLU(inplace=True)
)
(2): Conv2d(256, 35, kernel_size=(1, 1), stride=(1, 1))
)
)
(dfl): DFL(
(conv): Conv2d(16, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
)
)
)
)yolov8網(wǎng)絡從1-21層與pt文件相對應是BackBone和Neck模塊，22層是Head模塊。
二、轉ONNX步驟
2.1 yolov8官方
"""
代碼解釋
pt模型轉為onnx格式
"""
import os
from ultralytics import YOLO
model = YOLO("weights/best.pt")
success = model.export(format="onnx")
print("導出成功！")
將pytorch轉為onnx后，pytorch支持的一系列計算就會轉為onnx所支持的算子，若沒有相對應的就會使用其他方式進行替換（比如多個計算替換其單個）。比較常見是conv和SiLU合并成一個Conv模塊進行。
其中，1*4*8400表示每張圖片預測 8400 個候選框，每個框有 4 個參數(shù)邊界框坐標 (x,y,w,h)。 1*35*8400類同，1和4800代表意義相同，35是類別屬性包含了其置信度概率值。
最后兩個輸出Concat操作，得到1*39*8400。最后根據(jù)這個結果去進行后續(xù)操作。
2.2 自定義轉換
所謂的自定義轉換其實是在轉onnx時，對1*39*8400多加了一系列自定義操作例如NMS等。
2.2.1 加載權重并優(yōu)化結構
YOLOv8 = YOLO(args.weights) #替換為自己的權重 model = YOLOv8.model.fuse().eval()
2.2.2 后處理檢測模塊
def gen_anchors(feats: Tensor,
strides: Tensor,
grid_cell_offset: float = 0.5) -> Tuple[Tensor, Tensor]:
"""
生成錨點，并計算每個錨點的步幅。
參數(shù)：
feats (Tensor): 特征圖，通常來自不同的網(wǎng)絡層。
strides (Tensor): 每個特征圖的步幅（stride）。
grid_cell_offset (float): 網(wǎng)格單元的偏移量，默認為0.5。
返回：
Tuple[Tensor, Tensor]: 錨點的坐標和對應的步幅張量。
"""
anchor_points, stride_tensor = [], []
assert feats is not None # 確保輸入的特征圖不為空
dtype, device = feats[0].dtype, feats[0].device # 獲取特征圖的數(shù)據(jù)類型和設備
# 遍歷每個特征圖，計算錨點
for i, stride in enumerate(strides):
_, _, h, w = feats[i].shape # 獲取特征圖的高（h）和寬（w）
sx = torch.arange(end=w, device=device,
dtype=dtype) + grid_cell_offset # 計算 x 軸上的錨點位置
sy = torch.arange(end=h, device=device,
dtype=dtype) + grid_cell_offset # 計算 y 軸上的錨點位置
sy, sx = torch.meshgrid(sy, sx) # 生成網(wǎng)格坐標
anchor_points.append(torch.stack((sx, sy), -1).view(-1, 2)) # 將 x 和 y 組合成坐標點
stride_tensor.append(
torch.full((h * w, 1), stride, dtype=dtype, device=device)) # 生成步幅張量
return torch.cat(anchor_points), torch.cat(stride_tensor) # 返回合并后的錨點和步幅
class customize_NMS(torch.autograd.Function):
"""
繼承torch.autograd.Function
用于TensorRT的非極大值抑制（NMS）自定義函數(shù)。
"""
@staticmethod
def forward(
ctx: Graph,
boxes: Tensor,
scores: Tensor,
iou_threshold: float = 0.65,
score_threshold: float = 0.25,
max_output_boxes: int = 100,
background_class: int = -1,
box_coding: int = 0,
plugin_version: str = '1',
score_activation: int = 0
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
"""
正向計算NMS輸出，模擬真實的TensorRT NMS過程。
參數(shù)：
boxes (Tensor): 預測的邊界框。
scores (Tensor): 預測框的置信度分數(shù)。
其他參數(shù)同樣為NMS的超參數(shù)。
返回：
Tuple[Tensor, Tensor, Tensor, Tensor]: 包含檢測框數(shù)量、框坐標、置信度分數(shù)和類別標簽。
"""
batch_size, num_boxes, num_classes = scores.shape # 獲取批量大小、框數(shù)量和類別數(shù)
num_dets = torch.randint(0,
max_output_boxes, (batch_size, 1),
dtype=torch.int32) # 隨機生成檢測框數(shù)量（僅為模擬）
boxes = torch.randn(batch_size, max_output_boxes, 4) # 隨機生成預測框
scores = torch.randn(batch_size, max_output_boxes) # 隨機生成分數(shù)
labels = torch.randint(0,
num_classes, (batch_size, max_output_boxes),
dtype=torch.int32) # 隨機生成類別標簽
return num_dets, boxes, scores, labels # 返回模擬的結果
@staticmethod
def symbolic(
g,
boxes: Value,
scores: Value,
iou_threshold: float = 0.45,
score_threshold: float = 0.25,
max_output_boxes: int = 100,
background_class: int = -1,
box_coding: int = 0,
score_activation: int = 0,
plugin_version: str = '1') -> Tuple[Value, Value, Value, Value]:
"""
計算圖的符號函數(shù)，供TensorRT使用。
參數(shù)：
g: 計算圖對象
boxes (Value), scores (Value): 傳入的邊界框和得分
其他參數(shù)是用于配置NMS的參數(shù)。
返回：
經(jīng)過NMS處理的檢測框、得分、類別標簽及檢測框數(shù)量。
"""
out = g.op('TRT::EfficientNMS_TRT',
boxes,
scores,
iou_threshold_f=iou_threshold,
score_threshold_f=score_threshold,
max_output_boxes_i=max_output_boxes,
background_class_i=background_class,
box_coding_i=box_coding,
plugin_version_s=plugin_version,
score_activation_i=score_activation,
outputs=4) # 使用TensorRT的EfficientNMS插件
nums_dets, boxes, scores, classes = out # 獲取輸出的檢測框數(shù)量、框坐標、得分和類別
return nums_dets, boxes, scores, classes # 返回結果
class Post_process_Detect(nn.Module):
"""
用于后處理的檢測模塊，執(zhí)行檢測后的非極大值抑制（NMS）。
"""
export = True
shape = None
dynamic = False
iou_thres = 0.65 # 默認的IoU閾值
conf_thres = 0.25 # 默認的置信度閾值
topk = 100 # 輸出的最大檢測框數(shù)量
def __init__(self, *args, **kwargs):
super().__init__()
def forward(self, x):
"""
執(zhí)行后處理操作，提取預測框、置信度和類別。
參數(shù)：
x (Tensor): 輸入的特征圖。
返回：
Tuple[Tensor, Tensor, Tensor]: 預測框、置信度和類別。
"""
shape = x[0].shape # 獲取輸入的形狀
b, res, b_reg_num = shape[0], [], self.reg_max * 4
# b為特征列表第一個元素的批量大小，表示處理的樣本數(shù)量，
# res聲明一個空列表存儲處理過的特征圖
# b_reg_num為回歸框的數(shù)量
#遍歷特征層（self.nl表示特征層數(shù)），將每一層的框預測和分類預測拼接。
for i in range(self.nl):
res.append(torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)) # 特征拼接
# 調(diào)用
# make_anchors
# 生成錨點和步幅，用于還原邊界框的絕對坐標。
if self.dynamic or self.shape != shape:
self.anchors, self.strides = (x.transpose(
0, 1) for x in gen_anchors(x, self.stride, 0.5)) # 生成錨點和步幅
self.shape = shape # 更新輸入的形狀
x = [i.view(b, self.no, -1) for i in res] # 調(diào)整特征圖形狀
y = torch.cat(x, 2) # 拼接所有特征圖
boxes, scores = y[:, :b_reg_num, ...], y[:, b_reg_num:, ...].sigmoid() # 提取框和分數(shù)
boxes = boxes.view(b, 4, self.reg_max, -1).permute(0, 1, 3, 2) # 變換框的形狀
boxes = boxes.softmax(-1) @ torch.arange(self.reg_max).to(boxes) # 對框進行softmax處理
boxes0, boxes1 = -boxes[:, :2, ...], boxes[:, 2:, ...] # 分離框的不同部分
boxes = self.anchors.repeat(b, 2, 1) + torch.cat([boxes0, boxes1], 1) # 合并框坐標
boxes = boxes * self.strides # 乘以步幅
return customize_NMS.apply(boxes.transpose(1, 2), scores.transpose(1, 2),
self.iou_thres, self.conf_thres, self.topk) # 執(zhí)行NMS
def optim(module: nn.Module):
setattr(module, '__class__', Post_process_Detect)
for item in model.modules():
optim(item)
item.to(args.device) #輸入cpu或者gpu的卡號
自定義這里是在yolo官方得到的1*4*8400和1*35*8400進行矩陣轉換2<->3，最后引入EfficientNMS_TRT插件后處理，可以有效加速NMS處理。
2.2.3 EfficientNMS_TRT插件
EfficientNMS_TRT 是 TensorRT 中的一個高效非極大值抑制 (NMS) 插件，用于快速過濾檢測框。它通過優(yōu)化的 CUDA 實現(xiàn)來執(zhí)行 NMS 操作，特別適合于深度學習推理階段中目標檢測任務的后處理。支持在一個批次中對多個圖像同時執(zhí)行 NMS。
輸出結果為num_dets， detection_boxes， detection_scores， detection_classes ，分別代表經(jīng)過 NMS 篩選后保留的邊界框數(shù)，每張圖片保留的檢測框的坐標，每張圖片中保留下來的檢測框的分數(shù)（由高到低），每個保留下來的邊界框的類別索引。
三、結語
到此這篇關于YOLOv8模型pytorch格式轉為onnx格式的文章就介紹到這了,更多相關YOLOv8模型pytorch轉onnx格式內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家！
相關文章
Python 中的判斷語句,循環(huán)語句,函數(shù)
這篇文章主要介紹了Python 中的判斷語句,循環(huán)語句,函數(shù)，文章圍繞主題展開詳細的內(nèi)容介紹，具有一定的參考價值，需要的小伙伴可以參考一下2022-08-08
今天小編就為大家分享一篇pytorch使用指定GPU訓練的實例，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2019-08-08
這篇文章主要介紹了Python反射用法,結合實例形式簡單分析了Python反射的概念、原理及使用方法,需要的朋友可以參考下2017-12-12
Python使用微信itchat接口實現(xiàn)查看自己微信的信息功能詳解
這篇文章主要介紹了Python使用微信itchat接口實現(xiàn)查看自己微信的信息功能,結合實例形式分析了Python微信itchat模塊常見功能與操作技巧,需要的朋友可以參考下2019-08-08
在 Python 中,有多種方法可以實現(xiàn)數(shù)據(jù)合并,無論是合并列表、合并字典、合并字符串、合并CSV文件還是合并多個文件夾中的文件,都可以使用簡單而強大的Python技巧來實現(xiàn),通過合并數(shù)據(jù),可以更方便地進行數(shù)據(jù)處理和分析2024-03-03
這篇文章主要為大家詳細介紹了Python如何通過腳本實現(xiàn)音頻和視頻格式轉換,文中的示例代碼講解詳細,感興趣的小伙伴可以跟隨小編一起學習一下2025-03-03
這篇文章主要為大家介紹了Blender Python編程快速入門教程示例詳解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進步，早日升職加薪2022-08-08
Django 是 Python Web 開發(fā)最常用的框架之一，跟進它的最新變化絕對是必須的。下面這篇文章主要給大家介紹了關于django 2.0更新的10條注意事項，文中通過示例代碼介紹的非常詳細，需要的朋友可以參考借鑒，下面來一起看看吧。2018-01-01
對DataFrame數(shù)據(jù)中的重復行,利用groupby累加合并的方法詳解
今天小編就為大家分享一篇對DataFrame數(shù)據(jù)中的重復行,利用groupby累加合并的方法詳解，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2019-01-01
Python網(wǎng)絡請求模塊urllib與requests使用介紹
網(wǎng)絡爬蟲的第一步就是根據(jù)URL，獲取網(wǎng)頁的HTML信息。在Python3中，可以使用urllib和requests進行網(wǎng)頁數(shù)據(jù)獲取，這篇文章主要介紹了Python網(wǎng)絡請求模塊urllib與requests使用2022-10-10