How to use shape tensor in custom plugin?

AI & Data Science Deep Learning (Training & Inference) TensorRT
1

Description

Hi

I want to create a custom plugin to do adaptive pooling with dynamic shape support. And it would be cool if I can feed a shape tensor as the pooling output size. like below:

layer = network.add_plugin_v2(inputs=[input_trt, shape_trt], plugin=plugin)

It seems that getOutputDimensions() can only get the “shape” of the shape tensor, not the value. What should I do to implement this?

Environment

TensorRT Version : 7.0
GPU Type : 2080ti
Nvidia Driver Version : 440
CUDA Version : 10.0
CUDNN Version :
Operating System + Version : ubuntu18.04
Python Version (if applicable) : 3.7
TensorFlow Version (if applicable) :
PyTorch Version (if applicable) : 1.4
Baremetal or Container (if container which image + tag) :

2

Hi @TJWindows,
You can take reference from the below link to build a custom layer.

Please allow me some time to check on this.
Thanks!

3

Hi @TJWindows,
I hope this links helps you in the implementation of getOutputDimensions() -

Thanks!

4

Thanks for your anwser.

I did something like the link you provided. Of cause it works fine.

But I expect something like IResizeLayer, I can set the second input as a int32 shape tensor like below

reshape = network_definition.add_resize(y)
reshape.set_input(1, network_definition.add_shape(X)->get_output(0))

add_shape() give me a tensor of shape (4,) with value (n,c,h,w), I want to use the value (n,c,h,w) (the value might pass through other layer like elementwise or slice, generate a new int32 tensor) to decide the output of my custom layer, but I can only get (4,) in getOutputDimensions. How did IResizeLayer and IShuffleLayer implement this? Can I do the same with my custom layer?

by the way, here is my plugin code, for now I use another big tensor as the shape provider. Hope there is a better way :

5

Hi @TJWindows,
Hope you will get the details here.

Thanks!