Description
I’m exporting a model that uses a custom operator (from pytorch_scatter), so I wrote a custom TRT plugin that implements the operator and I’m now trying to export a custom model with it.
torch_scatter.scatter_sum (the function that I am trying to export) has dynamic output shape based on the data in one of its parameters and, since this doesn’t seem to be supported in TRT 8.4 (especially so for a TRT plugin), I thought of precalculating the value at which the output needs to be sliced, running the plugin op with a fixed output shape and following it with a slice
To keep things extremely simple and only test the bare minumum, I wrote the following pytorch model:
class M(torch.nn.Module):
def __init__(self, dim) -> None:
super().__init__()
self.dim = dim
def forward(self, src, index):
assert self.dim == 1, "The code as is now expects dim to be 1"
index_max = index.max()+1
out = torch.ops.torch_scatter.scatter_sum(src, index, self.dim, None, src.shape[self.dim])
return out[:, :index_max]
This can be exported successfully to ONNX. After applying the necessary node changes to place the TRT_PluginV2 ONNX node, we get this onnx model.
Finally, I try to export the model with:
CUDA_MODULE_LOADING=LAZY trtexec --plugins=libscatter_plugin.so --onnx=scatter_sum_minimal.onnx --saveEngine=engine.trt
This raises the following error:
[04/21/2023-11:38:51] [E] Error[2]: [shapeContext.cpp::volumeOfShapeTensor::470] Error Code 2: Internal Error (Assertion t.extent.nbDims == 0 || t.extent.nbDims == 1 failed. )
[04/21/2023-11:38:51] [E] [TRT] ModelImporter.cpp:773: While parsing node number 4 [Slice -> "13"]:
[04/21/2023-11:38:51] [E] [TRT] ModelImporter.cpp:774: --- Begin node ---
[04/21/2023-11:38:51] [E] [TRT] ModelImporter.cpp:775: input: "val_out_1"
input: "/Constant_1_output_0"
input: "/Unsqueeze_output_0"
input: "/Constant_2_output_0"
input: "/Constant_2_output_0"
output: "13"
name: "/Slice"
op_type: "Slice"
[04/21/2023-11:38:51] [E] [TRT] ModelImporter.cpp:776: --- End node ---
[04/21/2023-11:38:51] [E] [TRT] ModelImporter.cpp:779: ERROR: ModelImporter.cpp:180 In function parseGraph:
[6] Invalid Node - /Slice
[shapeContext.cpp::volumeOfShapeTensor::470] Error Code 2: Internal Error (Assertion t.extent.nbDims == 0 || t.extent.nbDims == 1 failed. )
[04/21/2023-11:38:51] [E] Failed to parse onnx file
[04/21/2023-11:38:51] [I] Finish parsing network model
[04/21/2023-11:38:51] [E] Parsing model failed
[04/21/2023-11:38:51] [E] Failed to create engine from model or file.
[04/21/2023-11:38:51] [E] Engine set up failed
What is the meaning of the error? Am I doing something wrong or misunderstanding how Slice should be used?
Environment
TensorRT Version: 8.4
GPU Type: Tesla V100
Nvidia Driver Version: 515.65
CUDA Version: 11.2
CUDNN Version: 8.1.1
Baremetal or Container (if container which image + tag): Baremetal