[Relax][Frontend][TFLite] Add soft-NMS support for TFLite NON_MAX_SUPPRESSION_V5#19426
[Relax][Frontend][TFLite] Add soft-NMS support for TFLite NON_MAX_SUPPRESSION_V5#19426Aharrypotter wants to merge 4 commits intoapache:mainfrom
Conversation
… TFLite NMSv5 This commit adds soft-NMS support to relax.vision.non_max_suppression and the TFLite frontend's NonMaxSuppressionV5 converter. Key changes: - Extended NonMaximumSuppressionAttrs with soft_nms_sigma and score_threshold. - Implemented Gaussian soft-NMS decay in TOPI _classic_nms_ir: score *= exp(-iou^2 / sigma) for overlapping boxes instead of hard suppression. - Introduced conditional return type: soft-NMS returns a 3-tuple (out_data, box_indices, valid_box_count) so the TFLite frontend can extract decayed scores; hard NMS keeps the original 2-tuple. - Updated Relax struct info inference to match the conditional return type. - Removed the soft_nms_sigma != 0.0 OpNotImplemented guard in TFLite frontend. - Added tests: struct info inference, legalization, TVMScript parser, and TFLite frontend E2E smoke tests for soft-NMS.
This commit fixes a critical bug in the soft-NMS path where decayed boxes
whose scores dropped below the threshold were not properly invalidated,
causing misalignment between selected_indices and selected_scores.
Fixes in TOPI _classic_nms_ir:
1. After Gaussian decay, invalidate boxes with score <= threshold by
setting their out_box_indices to -1.
2. During the compaction phase for return_indices=True, also compact
out_data so that the first num_valid entries match the selected boxes
in order. Fill remaining entries with -1.
Fixes in the NumPy reference implementation (nms_python.py):
- Synchronize the same post-decay invalidation logic to keep the
reference behavior consistent with the TIR kernel and TensorFlow.
Without this fix, the TFLite frontend would extract wrong scores from
processed_data[:max_output_size] when a middle box was decayed below
threshold but later boxes were still valid.
Align the soft-NMS path in non_max_suppression with LiteRT's reference behavior. - Rewrite soft-NMS candidate selection in TOPI TIR and the NumPy reference to re-pick the current best score after each decay step - Match LiteRT's Gaussian decay formula and keep decayed scores/data aligned with the returned indices for TFLite NMSv5 - Update conditional soft-NMS return documentation and add Relax/TFLite tests for reordered outputs and forwarded soft-NMS attributes Validation: - python -m pytest -n 1 tests/python/relax/test_op_vision.py -k "all_class_non_max_suppression or get_valid_counts or nms" -v Result: - 32 passed
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements Soft-NMS support across the TVM Relax stack, including the TFLite frontend, Relax operators, and TOPI TIR implementations. The changes introduce soft_nms_sigma and score_threshold attributes to the non_max_suppression operator and update the return signature to handle the additional output data required for soft-NMS. Feedback focuses on refactoring redundant code in the TFLite frontend, removing unreachable logic in the Python reference implementation, and consolidating duplicated IoU calculation logic within the TIR implementation to improve maintainability.
| num_valid_boxes_buf = T.alloc_buffer((1,), "int32", scope="local") | ||
| num_valid_boxes = T.buffer_proxy(num_valid_boxes_buf) | ||
| num_valid_boxes[0] = T.int32(0) | ||
| if is_soft_nms: |
There was a problem hiding this comment.
The IoU calculation logic is duplicated in both the soft-NMS (if is_soft_nms, lines 311-401) and hard-NMS (else, lines 483-537) branches. This makes the code harder to maintain. Consider refactoring this logic into a shared helper function or macro if possible within TIR script to avoid duplication.
There was a problem hiding this comment.
The IoU blocks operate in different indexing contexts and are self-contained within their respective branches. Extracting a shared TIR helper would add complexity without clear benefits for this already-large diff. I'd prefer to keep them as-is.
… impl - Deduplicate squeeze/strided_slice/reshape for selected_indices and num_valid across soft-NMS and hard-NMS branches - Remove redundant `if not is_soft_nms` check in the reference implementation (Greedy NMS section is only reached when False)
|
cc @tlopex |
tlopex
left a comment
tlopex
left a comment
There was a problem hiding this comment.
A few things still need to be cleaned up before this is ready:
-
Please add a note in
python/tvm/relax/op/vision/nms.pythat the return tuple shape depends onsoft_nms_sigma: it is a 2-tuple whensoft_nms_sigma == 0, and a 3-tuple (with decayedout_dataprepended) whensoft_nms_sigma > 0. That is easy to miss from the signature alone. -
In both
python/tvm/relax/op/vision/nms.pyandpython/tvm/topi/vision/nms.py, please clarify that thisscore_thresholdis a post-decay floor used only in the soft-NMS path, which is different fromget_valid_counts.score_threshold(pre-filter). Since the TFLite frontend passes both, it would be good to spell that out clearly. -
In
python/tvm/topi/vision/nms.py, the IoU computation block is duplicated between the hard-NMS and soft-NMS branches of_classic_nms_ir. Please consider pulling it into a local helper so the two paths do not drift on future fixes. -
In
python/tvm/topi/vision/nms.py,best_idx_buf,best_score_buf,tmp_idx_buf, andtmp_val_bufare currently allocated inside nestedT.If/T.serialbodies. Please move them up next tonum_valid_boxes_bufunderwith T.parallel(0, batch_size). -
In
python/tvm/relax/frontend/tflite/tflite_frontend.py, the soft-NMS path takes scores fromprocessed_data, which uses-1.0for invalid slots. TensorFlowNonMaxSuppressionV5expects non-negative padding. Please double-check that this is really okay, and add aclip(min=0.0)if needed.
2 participants
Summary
This PR completes the TFLite
NON_MAX_SUPPRESSION_V5implementation in Relax by adding support forsoft_nms_sigma != 0.It extends
relax.vision.non_max_suppressionwith soft-NMS attributes, updates the TFLite frontend to consume the soft-NMS outputs correctly, and aligns the TOPI implementation with LiteRT's reference behavior.Relates to #19412.
Changes
Relax / TOPI soft-NMS support
NonMaximumSuppressionAttrswithsoft_nms_sigmaandscore_threshold.relax.vision.non_max_suppressioncan represent theNON_MAX_SUPPRESSION_V5behavior.TFLite frontend support for
NON_MAX_SUPPRESSION_V5soft_nms_sigma != 0unsupported-path guard in the TFLite frontend.soft_nms_sigmaandscore_thresholdintorelax.vision.non_max_suppression.Soft-NMS correctness fixes
reference behavior.
Test coverage
NON_MAX_SUPPRESSION_V5withsoft_nms_sigma != 0.soft_nms_sigmaandscore_thresholdare forwarded correctly.Testing
Result:
NON_MAX_SUPPRESSION_V5coverage added for both hard-NMS and soft-NMS paths