[Relax][Frontend][TFLite] Add Conv3D support#19523
Conversation
![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 the CONV_3D operator in the TFLite frontend for Relax, including the convert_conv3d method which handles attribute extraction, layout mapping (NDHWC/DHWIO), and padding calculations. Unit tests for "VALID" and "SAME" padding modes have also been added. The review feedback suggests removing commented-out code, ensuring consistency by using a 6-element padding list for 3D convolutions, and correcting a typo in an error message.
| # "kernel_size": [kernel_d, kernel_h, kernel_w], | ||
| "strides": [stride_d, stride_h, stride_w], |
There was a problem hiding this comment.
Remove the commented-out code to keep the implementation clean and maintainable.
| # "kernel_size": [kernel_d, kernel_h, kernel_w], | |
| "strides": [stride_d, stride_h, stride_w], | |
| "strides": [stride_d, stride_h, stride_w], |
| # "kernel_size": [kernel_d, kernel_h, kernel_w], | ||
| "strides": [stride_d, stride_h, stride_w], | ||
| "dilation": [dilation_d, dilation_h, dilation_w], | ||
| "padding": [0, 0, 0], |
There was a problem hiding this comment.
For 3D convolution, it is more consistent to use a 6-element padding list (front, top, left, back, bottom, right) even for the zero-padding case. This aligns with the 6-element format used in the SAME padding logic and matches the expected IR structure in the unit tests.
| "padding": [0, 0, 0], | |
| "padding": [0, 0, 0, 0, 0, 0], |
| params["kernel_layout"] = "DHWIO" | ||
| if input_c != in_channels: | ||
| assert input_c % in_channels == 0, ( | ||
| "Input channels is not divisible of kernel in_channels." |
There was a problem hiding this comment.
Typo in the error message: "divisible of" should be "divisible by".
| "Input channels is not divisible of kernel in_channels." | |
| "Input channels is not divisible by kernel in_channels." |
This PR adds support for the CONV_3D operator in the TFLite frontend for Relax. It includes: - Implementation of convert_conv3d in OperatorConverter. - Mapping of CONV_3D builtin operator. - Unit tests to verify correctness with various padding and stride configurations.
weicheng-hsu
force-pushed
the
relax/tflite-conv3d
branch
from
9da3b2f to
9a361cc
Compare
|
Done. Updated the implementation to use consistent 6-element padding, removed redundant comments, and fixed the typo in the error message. Thanks! |
2 participants
Description
This PR adds support for the CONV_3D operator in the TFLite frontend for Relax.
Key Changes
Testing:
python3 -m pytest tests/python/relax/test_frontend_tflite.py -k "test_conv3d"Notes for Reviewers
The implementation follows the existing pattern used for CONV_2D but extends it to the 5D case (NDHWC layout). I've ensured that the kernel layout mapping aligns with TVM's R.nn.conv3d requirements.
Related to: #19519