Dynamic output shape question for nonzero operation

Description

The newest TensorRT7.1 does not support NonZero opreation. So I want to implement it by using nvinfer1::IPluginV2DynamicExt API.

But I don’t konw how to set the output shapes!

class NonZeroPluginDynamic : public nvinfer1::IPluginV2DynamicExt
{
public:
    ......
    nvinfer1::DimsExprs NonZeroPluginDynamic::getOutputDimensions(
        int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder)
    {
        int totalNonZeroCount = ???;
        return nvinfer1::Dims2(2, totalNonZeroCount);
    }
    ......
}

The totalNonZeroCount parameter can not be calculated by input shapes, but only counted in enqueue time.

For example, the output shapes for ONNX NonZero op is [inputDims, totalNonZeroCount].

import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto
import onnxruntime as ort
import numpy as np

x = helper.make_tensor_value_info('x', TensorProto.FLOAT, shape=[3, 3])
y = helper.make_tensor_value_info('y', TensorProto.INT64, shape=[2, None])

node_def = helper.make_node('NonZero', inputs=['x'], outputs=['y'])
graph_def = helper.make_graph([node_def], 'onnx-graph', [x], [y])

model_def = helper.make_model(graph_def)
onnx.checker.check_model(model_def)
onnx.save(model_def, 'nonzero.onnx')
# There are seven numbers are "none zero", so totalNonZeroCount=7
x_np = np.array([[1, 2, 0], [3, 4, 5], [6, 0, 7]], dtype=np.float32)
print(x_np.shape) # (3, 3)  inputDims= 2

ort_session = ort.InferenceSession('nonzero.onnx')
ort_inputs = {ort_session.get_inputs()[0].name: x_np}
ort_outputs = ort_session.run(None, ort_inputs)
y_np = ort_outputs[0]
print(y_np.shape) # (2, 7) inputDims=2, totalNonZeroCount=7

image

So, how do I implement getOutputDimensions function ?

Environment

TensorRT Version: 7.1.3.4

Hi @1034092330,
Request you to check the below link for reference.
https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#add_custom_layer

Thanks!

This link only explained how to calculate output shapes by input.

DimsExprs BarPlugin::getOutputDimensions(int outputIndex, 
    const DimsExprs* inputs, int nbInputs, 
    IExprBuilder& exprBuilder)
{
    switch (outputIndex)
    {
    case 0: 
    {
        // First dimension of output is sum of input 
        // first dimensions.
        DimsExprs output(inputs[0]);
        output.d[0] = 
            exprBuilder.operation(DimensionOperation::kSUM, 
                inputs[0].d[0], inputs[1].d[0]);
	   return output;
    }
    case 1:
        return inputs[0];
    default:
       throw std::invalid_argument(“invalid output”);
} 

In the above example, the shapes of the second output are calculated with input shapes.

But now NonZero output shapes cannot be calculated by input!
We don’t konw how many none zero numbers in input before runing and it will change with different inputs.
Even the output shapes will also change with different inputs.

image

Look forward to your favourable reply.
Thanks.

Can someone help me solve this issue ???

@ AakankshaS
Could you help me? Wish for your reply…

Hi @1034092330,
I am checking on this. Please allow me some time.

Thanks!

Thank you very much @AakankshaS
Other ONNX operations(such as NonMaxSuppression and Unique) also have the same issus that output shapes may change with different fixed shape inputs …


Onnx unique: https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Unique-11
Onnx NonMaxSuppression: https://github.com/onnx/onnx/blob/master/docs/Changelog.md#NonMaxSuppression-11