Support Exporting MultiScaleRoiAlign to ONNX#1324
Conversation
fmassa
left a comment
There was a problem hiding this comment.
This is pretty clean, thanks!
I have a few comments, let me know what you think.
Codecov Report
@@ Coverage Diff @@
## master #1324 +/- ##
=========================================
Coverage ? 63.05%
=========================================
Files ? 78
Lines ? 6193
Branches ? 846
=========================================
Hits ? 3905
Misses ? 1993
Partials ? 295
Continue to review full report at Codecov.
|
fmassa
left a comment
There was a problem hiding this comment.
LGTM, just waiting until ONNX support OrderedDict to merge this PR so that tests pass.
…idar/export_poolers
|
@fmassa the dict PR was merged an hour ago, but the test fails with unsupported dict error on the CI. |
…idar/export_poolers
fmassa
left a comment
There was a problem hiding this comment.
the dict PR was merged an hour ago, but the test fails with unsupported dict error on the CI.
Does the CI run the test with the latest pytorch master?
TorchVision CI runs with the latest PyTorch nightly, so it's expected that tests would fail if they had been merged in master 1h before.
I have one more comment regarding your latest changes, can you have a look?
Also, let me restart the CI for the failing jobs
| results = [] | ||
| for level, (per_level_feature, scale) in enumerate(zip(x, self.scales)): | ||
| idx_in_level = torch.nonzero(levels == level).squeeze(1) | ||
| idx_in_level = torch.nonzero(levels.to(torch.int64) == level).squeeze(1) |
There was a problem hiding this comment.
Can you explain why this cast is necessary? From
vision/torchvision/ops/poolers.py
Line 40 in e8d3291
levels are already cast to int64, so this seems unecessary?
There was a problem hiding this comment.
self.k_min is float so the returned value was a torch.float.
I added a cast on (target_lvls.to(torch.int64) - self.k_min) and reverted the one on levels.
There was a problem hiding this comment.
Oh I see, this is a regression introduced by type promotion, thanks a lot for the catch!
In this case, can you instead change
vision/torchvision/ops/poolers.py
Lines 111 to 112 in 17e355f
to do instead
lvl_min = int(-torch.log2(torch.tensor(scales[0], dtype=torch.float32)))
lvl_max = int(-torch.log2(torch.tensor(scales[-1], dtype=torch.float32)))i.e., replace .item() with a int()? This way it fixes it at the root cause of it
…idar/export_poolers
Modifications to export MultiScaleRoiAlign to ONNX.
Requires dictionary support added in pytorch/pytorch#25889.