Skip to content

Commit 0c67c53

Browse files
committed
fix(object-detection): don't export models with NMS and pin opset
- don't export models with NMS: - so we can tweak NMS settings on Robotoff side if needed - exported models didn't work out of the box (0 detections) when tested - pin opset to 20 for ONNX export
1 parent 8ca9558 commit 0c67c53

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

packages/train-yolo/main.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -344,10 +344,12 @@ def main(
344344
model.export(
345345
format="onnx",
346346
# Include NMS in the exported model
347-
nms=True,
348347
# Ultralytics tweaks the ONNX opset when exporting for GPUs
349348
# to prevent compatibility issues
350349
device="gpu",
350+
# We use opset 20, as it's the latest supported by our version of
351+
# Triton (25.02)
352+
opset=20,
351353
)
352354
# Rename the exported model to a standard name
353355
(run_dir / "weights/best.onnx").rename(run_dir / "weights/model.onnx")

0 commit comments

Comments
 (0)