[TensorRT] ERROR: INVALID_ARGUMENT: Cannot find binding of given name: input_0

I implement a very simple network and try to convert it using torch2trt:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        ...... #omitted
    def forward(self, x):
        x = F.tanh(self.conv1(x))
        x = F.tanh(self.conv2(x))
        ......

I only use two conv2d layers and nothing else. The conversion code is like:

   model.load_state_dict(torch.load('epochs/' + MODEL_NAME))
    x = torch.randn([1, 1, 270, 480]).cuda()
   model_trt = torch2trt(model, [x], fp16_mode=True)

Simple, right?

When I use the model:

out = model_trt(image)
print(out.size())

it gives me this fucking error. I check the dimensions of “out”, and it shows: torch.Size([1, 1, 1, 2160, 3840]) If I do not convert the model, it works just fine, by giving: torch.Size([1, 1, 2160, 3840]), which is the target output I want.

I cannot just figure out why such simple code cannot be converted, and I do not know what this error means.

I use ubuntu18.04 cuda10.2, cudnn7.6.5 with trt 7.0.11, and the lateset torch2trt with plugin. ``

[luhang-CCL]

I have an unsupported layer called: pixel_shuffle in my model, and I think it causes the problem: Warning: Encountered known unsupported method torch.nn.functional.pixel_shuffle If I comment out this layer, torch2trt works fine. It seems that this layer has added an extra dimension. Then, the problem is how I can live this unsupported layer?

[owoshch]

@hive-cas Hi! I have faced the same problem with pixel_shuffle layer in my network. Have you found the ready-to-use solution for this layer?

Thank you!