Description
When I tranform a dummy model from Pytorch to ONNX with the op torch.size(), it generates onnx::Gather and onnx::ConstantOfShape together with some other ops. An example shows below:
feat_size = 32
in_channel = 64
out_channel = 128
kernel_size = 3
batch = 2
feat = torch. randn(batch, in_channel, feat_size, feat_size, device="cuda")
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(in_channel, out_channel, kernel_size, padding=(1, 1))
def forward(self, x):
b, c, h, w = x.size()
y = torch.ones((b, c, h, w), dtype=torch.float32, device="cuda")
x = self. conv1(y)
model = Net().to("cuda")
torch.onnx.export(model, feat, "model.onnx", verbose=True)
When I tried to generate TensorRT engine, the following error occured:
In node 2 (importGather): UNSUPPORTED_NODE: Assertion failed: (data->getType() == nvinfer1::DataType::kINT32 && nbDims == 1) && "Cannot perform gather on a shape tensor!"
Since onnx::Gather operates the same manner of np.take, I tried to add a plugin layer to do the same thing. It failed because the input to Gather was a shape tensor and the plugin layer cannot take shape tensor as input.
On the other hand, I observed that the Split op defined in /TensorRT/parsers/onnx/builtin_op_imports.cpp
used exactly a shape tensor as input in gatherDimension()
(defined in onnx2trt_utils.cpp
), so I wonder if Gather op is properly handled in onnx_tensorrt and I comment out the assertion in the implementation in builtin_op_imports.cpp
.
Furthermore, I found that onnx::ConstantOfShape is not supported in TensorRT6.0 but it is supported in TensorRT7.0 and there is no additional op involved in the implementation not supported by TensorRT6.0. So I just copy the implementation in TensorRT7.0 to TensorRT6.0:
DEFINE_BUILTIN_OP_IMPORTER(ConstantOfShape)
{
OnnxAttrs attrs(node);
nvinfer1::ITensor* shape = &convertToTensor(inputs.at(0), ctx);
ShapedWeights zeroWeights
= ctx->createTempWeights(::ONNX_NAMESPACE::TensorProto_DataType_FLOAT, nvinfer1::Dims{1, 1});
static_cast<float*>(zeroWeights.values)[0] = 0.f;
auto valueWeights = TensorOrWeights{attrs.get("value", zeroWeights)};
nvinfer1::ITensor* value = &convertToTensor(valueWeights, ctx);
return {{constantOfShape(ctx, value, shape)}};
}
After doing the above two steps, there is a shape mismatch some where in ConstantOfShape.
Would you please add support to torch.size() and more documentation on addGather?
Environment
TensorRT Version: 6.0.1.5
GPU Type: V100
Nvidia Driver Version:
CUDA Version: 10.1
CUDNN Version: 7.6
Operating System + Version: Ubuntu18.04
Python Version (if applicable):
TensorFlow Version (if applicable):
PyTorch Version (if applicable): 1.4
Baremetal or Container (if container which image + tag):