Avoid onnx:NonZero operation when filltering pytorch

1

I’m trying to convert my PyTorch model to ONNX and after that to TensorRT and found out that there is no NonZero op for TensorRT
I found the part in my code that produces this operation and it is filtering tensors:

keep = (conf_scores > conf_thresh)
sc = conf_scores[keep]

When I apply my mask keep (it contains 1 and 0) to my tensor this operation produces NonZero tensor when I execute torch.onnx.export
I tried to a little bit another code:

mask = conf_scores > conf_thresh
indices = torch.nonzero(mask)
sc = conf_scores[indices]

The result is the same.
How to filter tensor and not to produce NonZero operation when converting in onnx?

UPD: tried the other code:

mask = conf_scores > conf_thresh
sc = torch.masked_select(conf_scores, mask)

and the other:

sc = conf_scores[conf_scores > conf_thresh]

and again there is NonZero op in ONNX

2

I met the same question, have you solved it?

3

No, whatever I do it always come with NonZero op.
I had to convert pythorch model to tensorrt and such code was in postprocessing (nms and so on), so I decided to remove postprocess from tensorrt and write it on pytorch. If this is the case I advice you to do the same

4

Thanks, I also met it in nms, but I found the latest onnx-tensorrt add NonZero op(just this mounth, lucky :grin:), but I haven’t try it, I met some questions when build, if it works, I will reply you.

5

Yes, it’s a long time issue from 2020 (No importer registered for op: NonZero · Issue #401 · onnx/onnx-tensorrt · GitHub)
If it works, let me know)