From 1539b5a678b388d2b678184a3e1431c91d5f10c7 Mon Sep 17 00:00:00 2001 From: wa22 Date: Mon, 3 Jun 2024 11:21:49 +0800 Subject: [PATCH] add calculation for flops --- flops.py | 8 ++++++++ ultralytics/nn/modules/head.py | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) create mode 100644 flops.py diff --git a/flops.py b/flops.py new file mode 100644 index 00000000..baa052bf --- /dev/null +++ b/flops.py @@ -0,0 +1,8 @@ +from ultralytics import YOLOv10 + +model = YOLOv10('yolov10n.yaml') +model.model.model[-1].export = True +model.model.model[-1].format = 'onnx' +del model.model.model[-1].cv2 +del model.model.model[-1].cv3 +model.fuse() \ No newline at end of file diff --git a/ultralytics/nn/modules/head.py b/ultralytics/nn/modules/head.py index 5bc7c068..b55942c2 100644 --- a/ultralytics/nn/modules/head.py +++ b/ultralytics/nn/modules/head.py @@ -496,7 +496,7 @@ class RTDETRDecoder(nn.Module): class v10Detect(Detect): - max_det = -1 + max_det = 300 def __init__(self, nc=80, ch=()): super().__init__(nc, ch)