Custom C++/CUDA from Torch to ONNX to TRT

username: 5had3z
cooked: <p>I was wondering what the best way to go forward from a Custom C++/CUDA Pytorch operation to onnx and then to TensorRT (I want to end up running realtime on an AGX Xavier).<br>
The custom module is done inline with “Custom C++ and CUDA Extensions”, however the torch.onnx Custom operations section mentions using torch script. Is this the proper way of doing it? Or can I just export/register the forward static method from the custom autograd function or something?</p>
<p>The actual error I get is below<br>
ONNX export failed: Couldn’t export Python operator CorrelationFunction</p>
<p>Aside: the onnx exporter also has issues with grid_sample, but I just added the ATen fallback, is this okay, or is it best to also make this a custom plugin for better performance?</p>
<p>Source code for reference<br>
</p><aside class="onebox allowlistedgeneric">
  <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/stereo-to-all/tree/master/nnet_training" target="_blank" rel="nofollow noopener">GitHub</a>
  </header>
  <article class="onebox-body">
    <img src="https://avatars0.githubusercontent.com/u/59598353?s=400&amp;v=4" class="thumbnail onebox-avatar" width="60" height="60">

<h3><a href="https://github.com/5had3z/stereo-to-all/tree/master/nnet_training" target="_blank" rel="nofollow noopener">5had3z/stereo-to-all</a></h3>

<p>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/stere...</p>


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

<p>Cheers</p>

username: 5had3z
cooked: <p>Uhh bump…?</p>

username: ofek
cooked: <p>Have you managed to export a model with grid_sampler to onnx?</p>

username: 5had3z
cooked: <p>Yeah, I had another go at it yesterday and I figured out how to correctly define a placeholder. You will still have to add your own runtime implementation for onnx. Since I’m using TensorRT I made an implementation for that (although I hasn’t been validated yet, I only just started working on it again).</p>
<aside class="quote quote-modified" data-post="1" data-topic="99695">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/5had3z/48/28066_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/onnx-export-placeholder-for-unsupported-pytorch-ops/99695">ONNX Export Placeholder for unsupported PyTorch ops?</a> 
  </div>
  <blockquote>
    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? 
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 …
  </blockquote>
</aside>