onnx-tensorrt: Gather in Upsample problem

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.

@jinfagang try

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