ONNX Export Placeholder for unsupported PyTorch ops?

username: 5had3z
cooked: <p>I’m trying to use ONNX as an intermediate bewteeen Pytorch and TensorRT. When exporting, onnx complains about grid_sample not being supported. Is there a way to get it to throw in a placeholder node with the same args etc for a PyTorch op?</p>
<p>I have previously added ONNX support for a correlation function which somewhat exports correctly. ONNX complains about exporting the autograd function, but I am able to just comment out the autograd call and use the custom op directly just when exporting for it to work correctly. I tried to be cheeky about it and use the raw function directly during self.eval() during exporting instead of the autograd function but it still tries to go through autograd.</p>
<p>I have already written both grid_sample and correlation IPluginV2ExtIO, so this intermediate step is currently the bottleneck (until I start debugging them haha).</p>

username: 5had3z
cooked: <p>Finally figured out that I needed to register the custom op without a domain namespace (I’m not sure if that’s strictly the right term) and it was recognised and replaced it correctly. I was trying things like torch::, functional:: or _VariableFunctions:: before and was uncertain if it was just the name I didn’t know or it wasn’t possible.</p>
<pre><code class="lang-auto">def grid_sample_op(g, input1, input2, mode, padding_mode, align_corners):
    return g.op("torch::grid_sampler", input1, input2, mode, padding_mode, align_corners)

torch.onnx.register_custom_op_symbolic('::grid_sampler', grid_sample_op, 11)
</code></pre>

username: briankosw
cooked: <p>Hey <a class="mention" href="/u/5had3z">@5had3z</a>, were you able to do inference when you registered <code>grid_sampler</code>? I was able to export a model that uses <code>grid_sample</code> to the ONNX format, but I wasn’t able to do inference with that same model.</p>

username: 5had3z
cooked: <p>I use TensorRT for inference and implemented the required plugins from the PyTorch souce code. You can find my source code for Tensor RT below. If you use ONNX Runtime for inference, I presume you will have to make and register a custom plugin for that. Unfortunately I don’t have any experience with ONNX other than using the graphs as an intermediate representation for Tensor RT to parse so I can’t really help with that.</p>
<p>As an aside (as an improvement on my previous solution), you can ensure correct typing and keyword arguments with the function defined below.</p>
<pre><code class="lang-auto">@parse_args('v', 'v', 'i', 'i', 'b')
def grid_sample_op(g, input1, input2, mode, padding_mode, align_corners):
    return g.op("torch::grid_sampler", input1, input2, interpolation_mode_i=mode,
                padding_mode_i=padding_mode, align_corners_i=align_corners)
</code></pre>
<aside class="onebox githubfolder">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="16" height="16">
      <a href="https://github.com/5had3z/CerberusNet/tree/master/runtime/cerberus_net/trt_plugins" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <img src="https://avatars.githubusercontent.com/u/59598353?s=400&amp;amp;v=4" class="thumbnail onebox-full-image" width="60" height="60">

<h3><a href="https://github.com/5had3z/CerberusNet/tree/master/runtime/cerberus_net/trt_plugins" target="_blank" rel="noopener nofollow ugc">5had3z/CerberusNet</a></h3>

<p><a href="https://github.com/5had3z/CerberusNet/tree/master/runtime/cerberus_net/trt_plugins" target="_blank" rel="noopener nofollow ugc">master/runtime/cerberus_net/trt_plugins</a></p>

  <p><span class="label1">The aim of this project is to generate a pointcloud with motion vectors and semantic segmentation of stereoscopic image sequences. Targeting real-time deployment on an nVIDIA Xavier. - 5had3z/Cerbe...</span></p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>