Add topk min function for trace and onnx #5310
Conversation
💊 CI failures summary and remediationsAs of commit 859aec2 (more details on the Dr. CI page): 💚 💚 Looks good so far! There are no failures yet. 💚 💚 This comment was automatically generated by Dr. CI (expand for details).Please report bugs/suggestions to the (internal) Dr. CI Users group. |
|
@shubhambhokare1 Thanks for the contribution. TorchVision's detection models do not properly support tracing. Instead we recommend users to JIT-script them. @fmassa @prabhat00155 Could you weight in on this one as you are more knowledgeable on ONNX? I also find it strange that despite not properly supporting tracing, our RPN code has some kind of mitigation about it. Any thoughts on why this is the case? |
|
Thanks @shubhambhokare1 for the PR, I have a couple of comments. |
shubhambhokare1
force-pushed
the
sbhokare/topk_onnx
branch
from
45756ed to
d64af49
Compare
datumbox
dismissed
their stale review
The user force pushed the branch and resolved the main issue.
Think it should be fixed now |
shubhambhokare1
force-pushed
the
sbhokare/topk_onnx
branch
from
d64af49 to
c96a819
Compare
shubhambhokare1
force-pushed
the
sbhokare/topk_onnx
branch
from
3b11dfa to
b02c5e6
Compare
|
Hey @prabhat00155! You merged this PR, but no labels were added. The list of valid labels is available at https://github.com/pytorch/vision/blob/main/.github/process_commit.py |
| """ | ||
| axis_dim_val = torch._shape_as_tensor(input)[axis].unsqueeze(0) | ||
| min_kval = torch.min(torch.cat((torch.tensor([orig_kval], dtype=axis_dim_val.dtype), axis_dim_val), 0)) | ||
| return min_kval # type: ignore[arg-type] |
There was a problem hiding this comment.
@prabhat00155 @shubhambhokare1 There might be performance implications here that we need to benchmark thoroughly. If we have such benchmarks we should post them in this PR to make them easily accessible.
If we don't I propose either to revert the PR until we measure it. Another alternative is to do the onnx friendly approach only when we are tracing and do the previous simpler estimation in all other cases.
Summary: * Add topk minimizer function to _utils * Apply ufmt formatting * Apply min function for tracing and scripting * Add type ignore to avoid cast * fix flake * Fix python_type_check Reviewed By: jdsgomes Differential Revision: D34475307 fbshipit-source-id: 0dfa9d59a0fc85d43247ab0f4024deda9640609c Co-authored-by: Prabhat Roy <[email protected]>
ONNX spec requires the k-value to be less than or equal to the number of inputs along provided dim. Certain models use the number of elements along a particular axis instead of K if K exceeds the number of elements along that axis. Previously, python's min() function was used to determine whether to use the provided k-value or the specified dim axis value.
However in cases where the model is being exported in tracing mode, python min() is static causing the model to be traced incorrectly and eventually fail at the topk node. In order to avoid this situation, in tracing mode, torch.min() is used instead.