Error [Convolution]: number of channels in input tensor to a convolution layer must not be dynamic

[Kracozebr]

Hi, I'm trying to use torch2trt_dynamic to convert model from pytorch to tensorrt and get the error that number of channels in input tensor to a convolution layer must not be dynamic To replicate the error I write some example code with lenet5. Here it is:

from torch2trt_dynamic import torch2trt_dynamic
import torch
import torch.nn as nn

class LeNet5(nn.Module):
    def __init__(self):
        super().__init__()

        self._body = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2),

            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2),
        )

        self._head = nn.Sequential(
            nn.Linear(in_features=16 * 5 * 5, out_features=120),
            nn.ReLU(inplace=True),
            nn.Linear(in_features=120, out_features=84),
            nn.ReLU(inplace=True),
            nn.Linear(in_features=84, out_features=10)
        )

    def forward(self, x):
        x = self._body(x)
        x = x.view(x.size()[0], -1)
        x = self._head(x)
        return x

model = LeNet5()
model_path = 'models/lenet5_mnist.pt'
model_to_save = 'models/lenet5_mnist_trt.pt'

# loading the model and getting model parameters by using load_state_dict
model.load_state_dict(torch.load(model_path))
model.eval().cuda()

# create example data
x = torch.ones((1, 1, 32, 32)).cuda()

# convert to TensorRT feeding sample data as input
opt_shape_param = [
    [
        [1, 1, 16, 16],   # min
        [1, 1, 32, 32],   # opt
        [1, 1, 64, 64]    # max
    ]
]
model_trt = torch2trt_dynamic(model, [x], fp16_mode=False, opt_shape_param=opt_shape_param)

x = torch.rand(1,1,32,32).cuda()
with torch.no_grad():
    y = model(x)
    y_trt = model_trt(x)

# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))

And there is the error:

[TensorRT] ERROR: (Unnamed Layer* 31) [Convolution]: number of channels in input tensor to a convolution layer must not be dynamic
[TensorRT] ERROR: Builder failed while analyzing shapes.
Traceback (most recent call last):
  File "example_converter.py", line 38, in <module>
    y_trt = model_trt(x)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/appuser/torch2trt_dynamic/torch2trt_dynamic/torch2trt_dynamic.py", line 471, in forward
    idx = self.engine.get_binding_index(input_name)
AttributeError: 'NoneType' object has no attribute 'get_binding_index'
appuser@ccaf751bbcf0:/yolact_edge/CustomPlugins$ python example_converter.py 
using lenet5
[TensorRT] ERROR: (Unnamed Layer* 31) [Convolution]: number of channels in input tensor to a convolution layer must not be dynamic
[TensorRT] ERROR: Builder failed while analyzing shapes.
Traceback (most recent call last):
  File "example_converter.py", line 39, in <module>
    y_trt = model_trt(x)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/appuser/torch2trt_dynamic/torch2trt_dynamic/torch2trt_dynamic.py", line 471, in forward
    idx = self.engine.get_binding_index(input_name)
AttributeError: 'NoneType' object has no attribute 'get_binding_index'

Ubuntu 18.04 Cuda 10.2 cudnn 8.0.5 Tensorrt 7.1.3 torch 1.8.1

[grimoire]

Hi ,This seem like a bug of tensorrt7, the code works when I switch the tensorrt version to 8.

[grimoire]

And, wait. Since you have flattened the tensor after self._body(x) and do linear after that, I think the opt_shape_param should not be dynamic.

[Kracozebr]

My last comment was deleted somehow =( Yes if I use opt_shpe_param like this:

opt_shape_param = [
    [
        [1, 1, 32, 32],   # min
        [1, 1, 32, 32],   # opt
        [1, 1, 32, 32]    # max
    ]
]

It works properly on tensorrt 7.