Files
apache--tvm/python
HoYi b915cac0bf [Relax][Frontend][TFLite] Add soft-NMS support for TFLite NON_MAX_SUPPRESSION_V5 (#19426)
## Summary

This PR completes the TFLite `NON_MAX_SUPPRESSION_V5` implementation in
Relax by adding support for `soft_nms_sigma != 0`.

It extends `relax.vision.non_max_suppression` with 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

1. **Relax / TOPI soft-NMS support**
- Extend `NonMaximumSuppressionAttrs` with `soft_nms_sigma` and
`score_threshold`.
- Thread the new attributes through Relax op registration, Python
wrapper, and legalization.
- Add soft-NMS handling to TOPI classic NMS so
`relax.vision.non_max_suppression` can represent the
`NON_MAX_SUPPRESSION_V5` behavior.

2. **TFLite frontend support for `NON_MAX_SUPPRESSION_V5`**
- Remove the previous `soft_nms_sigma != 0` unsupported-path guard in
the TFLite frontend.
- Forward `soft_nms_sigma` and `score_threshold` into
`relax.vision.non_max_suppression`.
- Handle the soft-NMS return path explicitly so the frontend reads
decayed scores from the processed NMS output instead of re-reading the
original score tensor.

3. **Soft-NMS correctness fixes**
- Fix the soft-NMS path so boxes whose scores fall below the threshold
after decay are invalidated consistently.
- Keep returned indices and decayed scores aligned in both the TOPI TIR
implementation and the NumPy reference implementation.
- Update the soft-NMS candidate selection logic to re-pick the current
best candidate after each decay step, matching LiteRT's
     reference behavior.
   - Align the Gaussian decay formula with LiteRT.

4. **Test coverage**
- Add Relax tests for soft-NMS struct-info inference and legalization.
- Add Relax E2E tests covering reordered outputs after score decay and
other soft-NMS follow-up cases.
- Add TFLite frontend tests for `NON_MAX_SUPPRESSION_V5` with
`soft_nms_sigma != 0`.
- Add IR checks to verify that `soft_nms_sigma` and `score_threshold`
are forwarded correctly.

## Testing

```bash
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
python -m pytest tests/python/relax/test_frontend_tflite.py -k "nms_v5" -v
```

## Result:
- Relax vision tests passed locally
- TFLite `NON_MAX_SUPPRESSION_V5` coverage added for both hard-NMS and
soft-NMS paths
2026-04-27 16:42:34 -04:00
..