Add ONNX support to KRCNNConvDeconvUpsampleHead#4315
Add ONNX support to KRCNNConvDeconvUpsampleHead#4315thiagocrepaldi wants to merge 1 commit intofacebookresearch:mainfrom
Conversation
f2bc516 to
d9a65ee
Compare
d9a65ee to
b2aa17f
Compare
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
cb7c08b to
d4d3ec2
Compare
a7aa3b9 to
b92ca14
Compare
|
@ppwwyyxx https://github.com/facebookresearch/detectron2/pull/4205/files#diff-7d768d0cd0f46e5999dcab0cb5f2f7a101611e771ff0483869568aa01e0b5240R95-R142 will test the correct dynamic axes behavior once both PRs are merged. Local testing verified it |
b92ca14 to
5c356fe
Compare
5c356fe to
d969d8f
Compare
|
@wat3rBro Could you review this one when you have some time? Thanks |
d969d8f to
ac7b413
Compare
|
Rebasing to allow up-to-date review |
ac7b413 to
dab1cf2
Compare
dab1cf2 to
88a738c
Compare
88a738c to
f0e2bb7
Compare
f0e2bb7 to
8a8ee8d
Compare
|
@thiagocrepaldi Sorry for the delay, we're trying to ping people internally so they can have a look. |
8a8ee8d to
b774488
Compare
|
@malfet FYI |
|
|
||
| # Although semantically equivalent, `reshape` is used instead of `squeeze` due | ||
| # to limitation during ONNX export of `squeeze` in scripting mode | ||
| roi_map = roi_map.reshape(roi_map.shape[1:]) # keypoints x H x W |
There was a problem hiding this comment.
Would appreciate it if comment could link to an issue in pytorch or onnx
|
@wat3rBro has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@thiagocrepaldi has updated the pull request. You must reimport the pull request before landing. |
|
@thiagocrepaldi The newly added test fails in Meta's prod environment (pytorch 2.0 alpha + probably outdated ONNX version) due to Detailed Error MessageThe error message comes from opset version 9 for |
eee10cb to
7757576
Compare
|
@thiagocrepaldi has updated the pull request. You must reimport the pull request before landing. |
The fix involves changes to both PyTorch and detectron2: * Pytorch had a bug which prevented some tensors to be identified as float (refer to pytorch/pytorch#81386) * `detectron2/structures/keypoints.py::heatmaps_to_keypoints` internally does advanced indexing on a `squeeze`d `roi_map`. The aforentioned `squeeze` fails rank inference due to the presence of `onnx::If` on its composition to support dynamic dims. By replacing `squeeze` by `replace` on detectron's `heatmaps_to_keypoints`, shape inference succeeds, allowing ONNX export to succeed with dynamic axes support.
7757576 to
1a56083
Compare
|
@thiagocrepaldi has updated the pull request. You must reimport the pull request before landing. |
|
Yup, this is a regression from newer pytorch. pytorch/pytorch#94427 should fix it Since yesterday, 2/7/2023, Torch uses ONNX 1.13.0, so the ONNX version shouldn't be an issue |
|
Oh great! pytorch/pytorch#94427 actually solves other tests in |
|
It seems CI is green. Should we merge this before it breaks again? lol |
|
@wat3rBro has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
Saw the other fix in pytorch/pytorch#94427, now the entire |
|
This pull request has been merged in 1523b3e. |
In order to export
KRCNNConvDeconvUpsampleHeadto ONNX usingtorch.jit.script, changes to both PyTorch and detectron2:Pytorch has a bug which prevents a tensor wrapped by a list as float. Refer to the required fix
detectron2/structures/keypoints.py::heatmaps_to_keypointsinternally does advanced indexing on asqueezed tensor. The aforementionedsqueezefails rank inference due to the presence ofonnx::Ifon its implementation (to support dynamic dims). The fix is replacingsqueezebyreshape. A possible fix tosqueezeon PyTorch side might be done too (TBD and would take some time), but the proposed change here does not bring any consequence to detectron2 while it enables ONNX support with scriptableKRCNNConvDeconvUpsampleHead.After the proposed changes, the
KRCNNConvDeconvUpsampleHeaddoes include aLoopnode to represent a for-loop inside the model anddynamic outputs, as shown below:This PR has been tested with ONNX Runtime (this PR) to ensure the ONNX output matches PyTorch's for different
gen_input(X, Y)combinations and it succeeded. The model was converted to ONNX once with a particular input and tested with inputs of different shapes and compared to equality to PyTorch'sDepends on: pytorch/pytorch#81386 and #4291