Skip to content

Commit

Permalink
[ONNX] Fix an assertion failure involving Slice (#71965)
Browse files Browse the repository at this point in the history
Before this change, exporting a model to ONNX involving Slice crashes at `axes[i]` in line 153 if C++ assertions are enabled:
```
/usr/include/c++/11.1.0/bits/stl_vector.h:1045: std::vector<_Tp, _Alloc>::reference std::vector<_Tp, _Alloc>::operator[](std::vector<_Tp, _Alloc>::size_type) [with _Tp = long int; _Alloc = std::allocator<long int>; std::vector<_Tp, _Alloc>::reference = long int&; std::vector<_Tp, _Alloc>::size_type = long unsigned int]: Assertion '__n < this->size()' failed.
```
The relevant check is https://github.com/gcc-mirror/gcc/blob/releases/gcc-11.1.0/libstdc++-v3/include/bits/stl_vector.h#L1045, which checks the vector index.

The issue can be reproduced by exporting Mask R-CNN or similar ones. For example,
```Python
import io
import torch
import torchvision as tv

model = tv.models.detection.maskrcnn_resnet50_fpn(pretrained=False)
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
with io.BytesIO() as f:
    torch.onnx.export(model, x, f, opset_version=11)
```
(extracted from [onnxoptimizer tests](https://github.com/onnx/optimizer/blob/master/onnxoptimizer/test/optimizer_test.py))

Tested environment: Arch Linux x86_64 with pytorch and torchvisoin installed from [the official repo](https://github.com/archlinux/svntogit-community/blob/packages/python-pytorch/trunk/PKGBUILD) and [AUR](https://aur.archlinux.org/cgit/aur.git/tree/PKGBUILD?h=python-torchvision), respectively.

Pull Request resolved: pytorch/pytorch#72989
  • Loading branch information
BowenBao authored and cyyever committed Feb 21, 2022
1 parent 76fa986 commit 6984ac7
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion torch/csrc/jit/passes/onnx/constant_fold.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ c10::optional<at::Tensor> runTorchSlice_opset10(
return c10::nullopt;
}
auto axes_a = inputTensorValues[3].accessor<int64_t, 1>();
axes.reserve(inputTensorValues[3].sizes()[0]);
axes.resize(inputTensorValues[3].sizes()[0]);
// ONNX slice accepts negative axis, fix this for aten op
for (const auto i : c10::irange(inputTensorValues[3].sizes()[0])) {
axes[i] = axes_a[i] < 0 ? axes_a[i] + inputTensorValues[0].sizes().size()
Expand Down

0 comments on commit 6984ac7

Please sign in to comment.