Interpolation of network output extremely slow

kevinpl07: <p>Hey there,</p>
<p>I trained a human portrait segmentation model which I then wrapped in another model so it does alpha blending to replace the background (similar to MS Teams or Zoom).</p>
<p>The segmentation itself is running on a 512x256 img. After segmentation and alpha blending I tried to use <a href="https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html" rel="noopener nofollow ugc">torch.nn.functional.interpolate</a> to get it back to at least 720p.</p>
<p>Now to my problem: the segmentation model including alpha blending is extremely efficient it needs like 8ms. But the interpolation brings the model to an inference time of 35ms.</p>
<p>Here’s the code of my wrapper model:</p>
<pre><code class="lang-python">class OverlayWrapperModel(torch.nn.Module):
    def _init_(self, seg_model):
        super(OverlayWrapperModel, self)._init_()
        self.net = seg_model
        
    def forward(self, input_stacked):
        # I stacked the webcam img and the overlay it should get 
        input_img, overlay, alpha = torch.split(input_stacked, [3, 3, 1], dim=1)
        
        # segmentation inference
        output_logits = self.net(input_img).sigmoid()
        
        # cutout the mask from alpha channel
        alpha[output_logits &gt; 0.5] = 0
        
        # combine webcam img and upscale
        return torch.nn.functional.interpolate(overlay * alpha + input_img * (1 - alpha),  (720, 1280), mode="nearest")
</code></pre>
<p>What confuses me:</p>
<ul>
<li>Why is the interpolation making my model so slow?</li>
<li>is there a better way to replace the background and get 1280x720 output img?</li>
</ul>
<p>Thanks in advance!</p>
<p>Regards<br>
Kev</p>

eqy: <p>A few questions:<br>
Is the interpolation being done on the GPU or CPU?<br>
If the model is being done on the GPU, have you explored jit scripting the model e.g., <a href="https://pytorch.org/docs/stable/generated/torch.jit.script.html" class="inline-onebox" rel="noopener nofollow ugc">torch.jit.script — PyTorch 1.13 documentation</a> as it could potentially fuse the pointwise operations done before the interpolation.</p>
<p>If you can tolerate a bleeding-edge user experience, I would also check if Torch 2.0’s compile function could also fuse some pointwise ops and offer some speedup as well:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://pytorch.org/get-started/pytorch-2.0/">
  <header class="source">
      <img src="https://pytorch.org/favicon.ico?" class="site-icon" width="32" height="32">

      <a href="https://pytorch.org/get-started/pytorch-2.0/" target="_blank" rel="noopener nofollow ugc">pytorch.org</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/361;"><img src="https://pytorch.org/assets/images/featured-img-pytorch-2.png" class="thumbnail" width="690" height="361"></div>

<h3><a href="https://pytorch.org/get-started/pytorch-2.0/" target="_blank" rel="noopener nofollow ugc">PyTorch 2.0</a></h3>

  <p>Overview</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


kevinpl07: <p>Sorry for the missing info:</p>
<p>Everything is being done on CPU.<br>
I’ll try torch’s 2.0 compile function and report back the improvements.</p>
<p>Thanks for your help!</p>

eqy: <p>If everything is being done on the CPU, I would be skeptical of <code>compile</code>/graph-mode improving performance. It may be that interpolation is just an expensive operation on CPU due to bandwidth requirements.</p>