Dealing with data-dependent slices in TensorRT 8.4


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:
        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 --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?


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

Please refer to below links related custom plugin implementation and sample:

While IPluginV2 and IPluginV2Ext interfaces are still supported for backward compatibility with TensorRT 5.1 and 6.0.x respectively, however, we recommend that you write new plugins or refactor existing ones to target the IPluginV2DynamicExt or IPluginV2IOExt interfaces instead.


I had already seen the Sample Support Guide (the one you link is for version 8.6 EA, while I’m on 8.4, but they seem to be equivalent for the problem at hand).

The custom plugin is written for IPluginV2DynamicExt. From what I read in the developer guide, the *IOExt interface is intended for plugins needing INT8 support, which is not the case for me.