Cannot convert EfficientNet with torch2trt

The issue - the EfficientNet implementation from the efficientnet-pytorch package is not convertible to TensorRT format with torch2trt.

It's very easy to reproduce - am testing this within the most recent pytorch container from NVIDIA - nvcr.io/nvidia/pytorch:21.07-py3.

Set up the environment:

pip install efficientnet-pytorch
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

These are the versions I'm working with, for reference:

efficientnet-pytorch          0.7.1
nvidia-dlprof-pytorch-nvtx    1.3.0
pytorch-quantization          2.1.0
pytorch-transformers          1.1.0
torch                         1.10.0a0+ecc3718
torch2trt                     0.3.0
torchtext                     0.11.0a0
torchvision                   0.11.0a0
tensorrt                      8.0.1.6

Try to convert:

import torch
from torch2trt import torch2trt
from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b3').eval().cuda()
x = torch.ones((1, 3, 224, 224)).cuda()
model_trt = torch2trt(model, [x])

Result:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/conda/lib/python3.8/site-packages/torch2trt-0.3.0-py3.8.egg/torch2trt/torch2trt.py", line 552, in torch2trt
    outputs = module(*inputs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1056, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/efficientnet_pytorch/model.py", line 314, in forward
    x = self.extract_features(inputs)
  File "/opt/conda/lib/python3.8/site-packages/efficientnet_pytorch/model.py", line 296, in extract_features
    x = block(x, drop_connect_rate=drop_connect_rate)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1056, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/efficientnet_pytorch/model.py", line 109, in forward
    x = self._depthwise_conv(x)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1056, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/efficientnet_pytorch/utils.py", line 275, in forward
    x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
  File "/opt/conda/lib/python3.8/site-packages/torch2trt-0.3.0-py3.8.egg/torch2trt/torch2trt.py", line 300, in wrapper
    converter["converter"](ctx)
  File "/opt/conda/lib/python3.8/site-packages/torch2trt-0.3.0-py3.8.egg/torch2trt/converters/conv_functional.py", line 46, in convert_Conv_trt7_functional
    layer.stride_nd = stride
TypeError: (): incompatible function arguments. The following argument types are supported:
    1. (arg0: tensorrt.tensorrt.IConvolutionLayer, arg1: tensorrt.tensorrt.Dims) -> None

Invoked with: <tensorrt.tensorrt.IConvolutionLayer object at 0x7f28bd52da70>, ([1, 1], [1, 1])

It's unclear what this error means or what needs to be done to resolve this. It's surprising to see issues with conv2d, as that is a very mature & commonly used operation in PyTorch.

Expected behavior:

The model should convert with no issues and without any changes needed to efficientnet-pytorch. Are there any suggestions for how to get this to convert?

NVIDIA-AI-IOT / torch2trt

Cannot convert EfficientNet with torch2trt #610