[feat] YOLO detection head: export merged (boxes, class_scores) ONNX outputs#9
Open
Go-MinSeong wants to merge 3 commits intodevfrom
Open
[feat] YOLO detection head: export merged (boxes, class_scores) ONNX outputs#9Go-MinSeong wants to merge 3 commits intodevfrom
Go-MinSeong wants to merge 3 commits intodevfrom
Conversation
…outputs
Previously the ONNX export produced three separate 4-D feature-map outputs
(one per detection scale). Downstream inference pipelines expect a single
decoded pair of 3-D tensors:
- output : (B, total_anchors, 4) xyxy box coordinates
- class_scores : (B, total_anchors, C) per-class probabilities (sigmoid)
Changes
-------
* yolo_head.py
- Add `prepare_export(input_size, feat_sizes)` to `YOLODetectionHead`.
When `feat_sizes` is supplied, anchor grids are pre-computed and stored
as constant buffers so no Range/Cast ONNX ops appear in the graph.
- Add `_decode_outputs()`: concatenates all scales, applies a *single*
Anchor2Vec (one Softmax + Conv3d), decodes ltrb distances to xyxy, and
returns `(boxes, class_scores)`.
- `forward()` delegates to `_decode_outputs()` when `_export` flag is set.
* utils/onnx.py
- `_save_onnx` now accepts a list of output names and builds dynamic-axes
entries for each output automatically.
- `save_onnx` detects export mode via `head._export` and passes
`output_names=['output', 'class_scores']` accordingly.
* tools/onnx_convert_its.py
- Before calling `prepare_export`, run backbone+neck once with a dummy
input to obtain per-scale feature map sizes; pass them to
`prepare_export` so anchor buffers are baked as ONNX constants.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Merge prepare_export logic from onnx_convert_its.py into onnx_convert.py - Replace hardcoded TEMP_NUM_CLASSES with --num-classes CLI argument (default 80) - Fix --sample-size to use nargs=2 / int for correct parsing - Remove onnx_convert_its.py (duplicate) and tools/model_eval.yaml (local config) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
onnx_convert.py로 통합 (refactor)Changes
src/netspresso_trainer/heads/detection/experimental/yolo_head.py— YOLO head에 merged (boxes, class_scores) ONNX output 지원 추가src/netspresso_trainer/utils/onnx.py— ONNX 유틸리티 개선tools/onnx_convert.py— ONNX export 스크립트 통합 및 개선.github/CODEOWNERS— CODEOWNERS 수정.gitignore— gitignore 항목 추가Dockerfile— Dockerfile 수정Test plan
onnx_convert.py통합 스크립트 정상 동작 확인🤖 Generated with Claude Code