onnx转TensorRT遇到Assertion failed: axis ＞= 0 && axis ＜ nbDims问题

最新推荐文章于 2024-03-27 18:25:48 发布

foso1994

最新推荐文章于 2024-03-27 18:25:48 发布

阅读量2.3k

点赞数 4

分类专栏：日常Bug处理 Deep Learning 文章标签： pytorch

本文链接：https://blog.csdn.net/foso1994/article/details/108881940

版权

在将ONNX模型转换为TensorRT时遇到断言失败错误，问题根源在于网络中的Gather操作。通过Netron分析，发现插值操作与预期不符，原因是输入Tensor的shape属性包含1*1的Tensor而不是int。解决方案是将shape中1*1的Tensor转换为int，然后正确组合成新的shape，避免了额外的Gather分支，成功解决了问题。

摘要由CSDN通过智能技术生成

While parsing node number xxx [Gather]:
ERROR: onnx2trt_utils.hpp:277 In function convert_axis:
[8] Assertion failed: axis >= 0 && axis < nbDims

这个问题我在网上看了很多帖子,基本都是reshape的问题,但是我的网络结构中并不含有reshape函数,因此与其他的解决方式不同.

我将onnx转入到netron中,发现了一些端倪:

报错的位置Gather操作,源自于

F.interpolate(x, size=x.shape[-2:], mode='nearest')

插值操作,但是另我不解的是,明明一个简单上采样,为什么会多出两个Gather分支?由网络结构中可以看到这个两个分支是用来取到x.shape的第2和第3个索引的,于是我打开pycharm的调试,发现对于x这个Tensor的shape属性,并不是我以为的[a,b,c,d]这种形式,而是[Tensor(a),Tensor(b),Tensor(c),Tensor(d)]!也就是说,对于shape中的每一个维度并不是int类型,而是每个int类型的数构成的1*1的Tensor.

这也就解释了为什么网络结构中会取出两个索引后,还需要进行Unsqueeze操作,因为需要将取出的1*1Tensor展开到一个int形式,再进行concat组成一个[c,d]的shape形式.

于是就可以先将shape中1*1的Tensor转为int类型后,再组成shape,这样再进行插值操作,即:</