Onnx-tensorrt: Gather in Upsample problem

[aidonchuk]

Hi! Cant export model from onnx to tensorrt.

`----------------------------------------------------------------
Input filename: model.onnx
ONNX IR version: 0.0.4
Opset version: 9
Producer name: pytorch
Producer version: 1.1
Domain:
Model version: 0

Doc string:

WARNING: ONNX model has a newer ir_version (0.0.4) than this parser was built against (0.0.3).
Parsing model
WARNING: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
Successfully casted down to INT32.
While parsing node number 69 [Gather -> "208"]:
ERROR: /home/alex/tools/onnx-tensorrt/onnx2trt_utils.hpp:335 In function convert_axis:
[8] Assertion failed: axis >= 0 && axis < nbDims
%206 : Long() = onnx::Constantvalue={2}, scope: ResNet18_OneConvDecoder/DecoderBlock[center]/Sequential[block]/Upsample[0]
%207 : Tensor = onnx::Shape(%205), scope: ResNet18_OneConvDecoder/DecoderBlock[center]/Sequential[block]/Upsample[0]
%208 : Long() = onnx::Gatheraxis=0, scope: ResNet18_OneConvDecoder/DecoderBlock[center]/Sequential[block]/Upsample[0]
%209 : Tensor = onnx::Constantvalue={2}
%210 : Tensor = onnx::Mul(%208, %209)`

[aidonchuk]

Please help!!!

[yalagamsrinivas]

I have same issue , please suggest me solution

[zimenglan-sysu-512]

any idea to solve it?

[aidonchuk]

I'm replace UpSample with ConTanspose2D but i don't like such workaround.

[zimenglan-sysu-512]

i solve this problem by only exporting the backbone without fpn to onnx.

[tibistrat]

I also had this issue with a model coming from PyTorch. Here's an explanation of what I did to work around the problem:

This PyTorch model, when exported to ONNX, fails when importing in TensorRT because of the Gather operation:

class ShapeModel(nn.Module):
    def __init__(self):
        super(ShapeModel, self).__init__()
    def forward(self, x):
        return x.shape

Assertion failed: axis >= 0 && axis < nbDims

However, this model works:

class ShapeModel(nn.Module):
    def __init__(self):
        super(ShapeModel, self).__init__()
    def forward(self, x):
        return torch.tensor(x.shape)

with just a warning during export to onnx that the trace might not generalize to other inputs.

A single PyTorch upsampling by a factor of 2 gets traced like this:

class ResizeModel(nn.Module):
    def __init__(self):
        super(ResizeModel, self).__init__()
    def forward(self, x):
        return F.interpolate(x, scale_factor=(2, 2), mode='nearest')

in which a lot of work is done to determine the desired size of the output tensor (and in which Gather appears).

A PyTorch interpolate function will also work if you supply not the upsampling factor, but the already-known future size of your tensor. Below an upsampler for (batch_size x channels x H x W) tensors:

class ResizeModel(nn.Module):
    def __init__(self):
        super(ResizeModel, self).__init__()
    def forward(self, x):
        sh = torch.tensor(x.shape)
        return F.interpolate(x, size=(sh[2] * 2, sh[3] * 2), mode='nearest')

Which gets traced to ONNX like this:

thus avoiding the Gather and which functions in TensorRT.

[aidonchuk]

@tibistrat you are absolulty magician! Thx a lot!

[jinfagang]

@tibistrat Thanks for your explaination. But actually, with this code:

        return F.interpolate(x, size=(sh[2] * 2, sh[3] * 2), mode='nearest')

It can not onnx2trt either, it will throw such an error

[8] Assertion failed: axis >= 0 && axis < nbDims

I have tested with a similar code:

up3 = F.interpolate(output3, size=(output2.size(2), output2.size(3)), mode="nearest")
        output2 = output2 + up3

Once I commented out this line, it works. Uncomment, it will raise error.

I still don't know why it hits and how to solve it.

Any suggestions?

[backtime92]

@jinfagang Have you solved it?

[jinfagang]

No, I don't know what's the reason exactly.

[kealennieh]

@jinfagang try

up3 = F.interpolate(output3, size=(int(output2.size(2)), int(output2.size(3))), mode="nearest") 

[freetown113]

@jinfagang You know your output exact size, right? You have to enter these exact values, for example:
F.interpolate(output3, size=[64,64], mode="nearest")
As well as you need to provide a distinct output shape in "forward" function for example:
x = view(1,-1,1024,1024)
as mentioned here: https://github.com/pytorch/pytorch/issues/16908

[DJMeng]

@jinfagang Have you solved it ? thanks

[DJMeng]

@aidonchuk have you solved it as the @tibistrat saids？ why i still have errors? Pls help me , thanks

[banhbaomocmeo]

@jinfagang try

up3 = F.interpolate(output3, size=(int(output2.size(2)), int(output2.size(3))), mode="nearest") 

It works perfectly, thank you

[LeviViana]

@tibistrat Thanks a lot !

[codeslord]

You may use onnx-simplifier to do the same
https://github.com/daquexian/onnx-simplifier

[daixiangzi]

I also had this issue with a model coming from PyTorch. Here's an explanation of what I did to work around the problem:

This PyTorch model, when exported to ONNX, fails when importing in TensorRT because of the Gather operation:

class ShapeModel(nn.Module):
    def __init__(self):
        super(ShapeModel, self).__init__()
    def forward(self, x):
        return x.shape

Assertion failed: axis >= 0 && axis < nbDims

However, this model works:

class ShapeModel(nn.Module):
    def __init__(self):
        super(ShapeModel, self).__init__()
    def forward(self, x):
        return torch.tensor(x.shape)

with just a warning during export to onnx that the trace might not generalize to other inputs.

A single PyTorch upsampling by a factor of 2 gets traced like this:

class ResizeModel(nn.Module):
    def __init__(self):
        super(ResizeModel, self).__init__()
    def forward(self, x):
        return F.interpolate(x, scale_factor=(2, 2), mode='nearest')

in which a lot of work is done to determine the desired size of the output tensor (and in which Gather appears).

A PyTorch interpolate function will also work if you supply not the upsampling factor, but the already-known future size of your tensor. Below an upsampler for (batch_size x channels x H x W) tensors:

class ResizeModel(nn.Module):
    def __init__(self):
        super(ResizeModel, self).__init__()
    def forward(self, x):
        sh = torch.tensor(x.shape)
        return F.interpolate(x, size=(sh[2] * 2, sh[3] * 2), mode='nearest')

Which gets traced to ONNX like this:

thus avoiding the Gather and which functions in TensorRT.

it look like cool!

[kevinch-nv]

Closing since this has been resolved.