Trying to avoid onnx:NonZero operation

username: Andrea_Rosasco
cooked: <p>Hello,<br>
I’m trying to export a model in onnx and to run it with TensorRT.</p>
<pre><code class="lang-python">from polygraphy.backend.trt import EngineFromNetwork, NetworkFromOnnxPath
import torch

class Model(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.x2 = torch.zeros((2048, 1)).cuda()

    def forward(self, x1):
        x2 = self.x2
        idx = x2 &lt; x1
      x1[idx] = x2[idx]
        return x1


if __name__ == '__main__':
    onnx_file = 'test.onnx'

    model = Model()
    x = torch.zeros((2048, 1)).cuda()
    torch.onnx.export(model, x, onnx_file, input_names=['input'], output_names=['output'], opset_version=11)

    build_engine = EngineFromNetwork(NetworkFromOnnxPath(onnx_file))
    engine = build_engine()
</code></pre>
<p>While the above model is correctly exported as an onnx, TensorRT has trouble parsing it:</p>
<pre><code class="lang-auto">[02/02/2022-17:38:36] [TRT] [E] ModelImporter.cpp:773: While parsing node number 3 [NonZero -&gt; "4"]:
[02/02/2022-17:38:36] [TRT] [E] ModelImporter.cpp:774: --- Begin node ---
[02/02/2022-17:38:36] [TRT] [E] ModelImporter.cpp:775: input: "3"
output: "4"
name: "NonZero_3"
op_type: "NonZero"

[02/02/2022-17:38:36] [TRT] [E] ModelImporter.cpp:776: --- End node ---
[02/02/2022-17:38:36] [TRT] [E] ModelImporter.cpp:779: ERROR: builtin_op_importers.cpp:4870 In function importFallbackPluginImporter:
[8] Assertion failed: creator &amp;&amp; "Plugin not found, are the plugin name, version, and namespace correct?"
[E] In node 3 (importFallbackPluginImporter): UNSUPPORTED_NODE: Assertion failed: creator &amp;&amp; "Plugin not found, are the plugin name, version, and namespace correct?"
</code></pre>
<p>I’d like to rewrite the code in such a way that this NonZero operation is not present in the onnx graph, is that possible?</p>
<p>This is the onnx graph:</p>
<pre><code class="lang-auto">graph(%input : Float(2048, 1, strides=[1, 1], requires_grad=0, device=cuda:0),
      %23 : Long(1, strides=[1], requires_grad=0, device=cpu),
      %24 : Long(1, strides=[1], requires_grad=0, device=cpu)):
  %1 : Float(2048, 1, strides=[1, 1], requires_grad=0, device=cuda:0) = onnx::Constant[value=&lt;Tensor&gt;]() # C:/Users/arosasco/PycharmProjects/pcr/delete2.py:13:0
  %2 : Float(2048, 1, strides=[1, 1], requires_grad=0, device=cuda:0) = onnx::Constant[value=&lt;Tensor&gt;]()
  %3 : Bool(2048, 1, strides=[1, 1], requires_grad=0, device=cuda:0) = onnx::Less(%2, %input) # C:/Users/arosasco/PycharmProjects/pcr/delete2.py:13:0
  %4 : Long(2, *, device=cpu) = onnx::NonZero(%3)
  %5 : Long(*, 2, device=cpu) = onnx::Transpose[perm=[1, 0]](%4)
  %6 : Float(*, strides=[1], requires_grad=0, device=cuda:0) = onnx::GatherND(%1, %5) # C:/Users/arosasco/PycharmProjects/pcr/delete2.py:14:0
  %7 : Long(2, strides=[1], device=cpu) = onnx::Shape(%input)
  %8 : Bool(2048, 1, device=cpu) = onnx::Expand(%3, %7)
  %9 : Long(2, *, device=cpu) = onnx::NonZero(%8)
  %10 : Long(*, 2, device=cpu) = onnx::Transpose[perm=[1, 0]](%9)
  %11 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={-1}]()
  %12 : Float(*, device=cpu) = onnx::Reshape(%6, %11)
  %13 : Long(2, strides=[1], device=cpu) = onnx::Shape(%10)
  %14 : Long(device=cpu) = onnx::Constant[value={0}]()
  %15 : Long(device=cpu) = onnx::Gather[axis=0](%13, %14)
  %18 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[axes=[0]](%15)
  %21 : Float(*, device=cpu) = onnx::Slice(%12, %23, %18, %24)
  %output : Float(2048, 1, strides=[1, 1], requires_grad=0, device=cuda:0) = onnx::ScatterND(%input, %10, %21) # C:/Users/arosasco/PycharmProjects/pcr/delete2.py:14:0
  return (%output)
</code></pre>
<p>It looks like the problem is around lines 13 and 14 of the above scripts:</p>
<pre><code class="lang-python">idx = x2 &lt; x1
x1[idx] = x2[idx]
</code></pre>
<p>I’ve tried to change the first line with <code>torch.zeros_like(x1).to(torch.bool)</code> but the problem persists so I’m thinking the issue is with the second one.</p>
<p>I have no clue on how to solve it, can anyone help?</p>

username: Andrea_Rosasco
cooked: <p>Apparently changing it with <code>torch.where(x2 &lt; x1, x2, x1)</code> solved the problem.</p>