thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Minor fixes

Browse Source
This commit is contained in:
Marcos Luciano
2023-09-01 12:52:07 -03:00
parent fb6729c382
commit 0f4377f540
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
utils/export_yoloV8.py
View File

@@ -8,7 +8,7 @@ import torch.nn as nn
from copy import deepcopy from copy import deepcopy
from ultralytics import YOLO from ultralytics import YOLO
from ultralytics.yolo.utils.torch_utils import select_device from ultralytics.yolo.utils.torch_utils import select_device
from ultralytics.nn.modules import C2f, Detect from ultralytics.nn.modules import C2f, Detect, RTDETRDecoder
class DeepStreamOutput(nn.Module): class DeepStreamOutput(nn.Module):
@@ -38,7 +38,7 @@ def yolov8_export(weights, device):
model.float() model.float()
model = model.fuse() model = model.fuse()
for k, m in model.named_modules(): for k, m in model.named_modules():
if isinstance(m, Detect): if isinstance(m, (Detect, RTDETRDecoder)):
m.dynamic = False m.dynamic = False
m.export = True m.export = True
m.format = 'onnx' m.format = 'onnx'