Skip to content

[Relax][Frontend][TFLite] Add NON_MAX_SUPPRESSION_V4 converter#19464

Merged
tlopex merged 1 commit intoapache:mainfrom
as4230:feat/tflite-nms-v4
Apr 29, 2026
Merged

[Relax][Frontend][TFLite] Add NON_MAX_SUPPRESSION_V4 converter#19464
tlopex merged 1 commit intoapache:mainfrom
as4230:feat/tflite-nms-v4

Conversation

@as4230
Copy link
Copy Markdown
Contributor

@as4230 as4230 commented Apr 28, 2026

Adds the missing TFLite NonMaxSuppressionV4 frontend handler. The
underlying relax.op.vision.non_max_suppression already covers V4's
behavior with soft_nms_sigma at the default 0.0 (hard-NMS path). The
handler bridges TFLite's tensor format to the Relax op, following the
same pattern as convert_nms_v5 (#19426) but without its soft-NMS
branching.

Tests cover conversion and IR structural assertions, run with
pytest tests/python/relax/test_frontend_tflite.py -k nms_v4. E2E
correctness runs on the nightly gate (CI_ENV_NIGHTLY).

Relates to #19412.

Adds the missing TFLite NonMaxSuppressionV4 frontend handler. The
underlying relax.op.vision.non_max_suppression already covers V4's
behavior with soft_nms_sigma at the default 0.0 (hard-NMS path). The
handler bridges TFLite's tensor format to the Relax op, following the
same pattern as convert_nms_v5 (apache#19426) but without its soft-NMS
branching.

Tests cover conversion and IR structural assertions, run with
`pytest tests/python/relax/test_frontend_tflite.py -k nms_v4`. E2E
correctness runs on the nightly gate (CI_ENV_NIGHTLY).

Relates to apache#19412.
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for the NON_MAX_SUPPRESSION_V4 operator in the TFLite frontend for TVM Relax. The changes include the implementation of the convert_nms_v4 method and the addition of comprehensive unit tests covering IR verification and end-to-end correctness. Feedback was provided to address a potential shape mismatch issue that occurs when the requested max_output_size exceeds the number of input boxes, suggesting that the output indices should be padded to ensure a consistent tensor shape.

Comment on lines +3643 to +3646
selected_indices = relax.op.squeeze(nms_ret[0], axis=[0])
selected_indices = relax.op.strided_slice(
selected_indices, axes=[0], begin=[0], end=[max_output_size]
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The current implementation of strided_slice may result in an incorrect output shape when max_output_size is greater than the number of input boxes (num_boxes). In TFLite, NON_MAX_SUPPRESSION_V4 (and V5) typically pads the output to exactly max_output_size if pad_to_max_output_size is enabled (which is common).

Since relax.op.vision.non_max_suppression returns a tensor of shape [batch_size, num_anchors], the squeezed selected_indices has length num_boxes. If max_output_size > num_boxes, strided_slice with end=[max_output_size] will clip to num_boxes, resulting in a shape mismatch if the rest of the graph expects max_output_size.

Consider padding the indices to ensure the output shape is consistently max_output_size regardless of the input size.

Suggested change
selected_indices = relax.op.squeeze(nms_ret[0], axis=[0])
selected_indices = relax.op.strided_slice(
selected_indices, axes=[0], begin=[0], end=[max_output_size]
)
num_boxes = int(self.get_tensor_shape(input_tensors[0])[0])
if max_output_size > num_boxes:
selected_indices = relax.op.nn.pad(
selected_indices, pad_width=[0, max_output_size - num_boxes], pad_value=-1
)
selected_indices = relax.op.strided_slice(
selected_indices, axes=[0], begin=[0], end=[max_output_size]
)

Copy link
Copy Markdown
Member

@tlopex tlopex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@tlopex tlopex merged commit c1415d6 into apache:main Apr 29, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants