F.grid_sample extremely slow

qiminchen: <p>I found that <code>F.grid_sample</code> in my code is extremely slow, for example, the following block takes about <code>0.9s</code> on GPU with Pytorch 1.6.0</p>
<pre><code class="lang-auto">B, H, W, D, C = coords.size()
coords = coords.view(B, -1, C).unsqueeze(1)  # (B, 1, N, 3), N = 262,144
# yz_plane with size of (1, 32, 256, 256)
sample_yz_feat = F.grid_sample(yz_plane.repeat(B, 1, 1, 1), coords[..., [1, 2]].flip(-1), align_corners=True).permute(0, 2, 3, 1).squeeze(1)  # B N 32
</code></pre>
<p>However, if I add two more <code>F.grid_sample</code>, for example, the following block <strong>STILL</strong> takes about <code>0.9s</code> on GPU. I wonder why in my case <code>F.grid_sample</code> is so slow, am I missing something since the other two <code>F.grid_sample</code> barely cost time.</p>
<pre><code class="lang-auto">B, H, W, D, C = coords.size()
coords = coords.view(B, -1, C).unsqueeze(1)  # (B, 1, N, 3), N = 262,144
# yz_plane, xz_plane, xy_plane with size of (1, 32, 256, 256)
sample_yz_feat = F.grid_sample(yz_plane.repeat(B, 1, 1, 1), coords[..., [1, 2]].flip(-1), align_corners=True).permute(0, 2, 3, 1).squeeze(1)  # B N 32
sample_xz_feat = F.grid_sample(xz_plane.repeat(B, 1, 1, 1), coords[..., [0, 2]].flip(-1), align_corners=True).permute(0, 2, 3, 1).squeeze(1)
sample_xy_feat = F.grid_sample(xy_plane.repeat(B, 1, 1, 1), coords[..., [0, 1]].flip(-1), align_corners=True).permute(0, 2, 3, 1).squeeze(1)
</code></pre>
<p>Updated:</p>
<p>This is how I compute the time</p>
<pre><code class="lang-auto">import time

coords = torch.randn((1, 64, 64, 64, 3)).cuda()
yz_plane = torch.randn((1, 32, 256, 256)).cuda()
xy_plane = torch.randn((1, 32, 256, 256)).cuda()
xz_plane = torch.randn((1, 32, 256, 256)).cuda()
start_time = time.time()
# above code block
print(time.time() - start_time)
</code></pre>

kingsj0405: <p>It seems benchmark with <code>torch.utils.benchmark</code> show it requires honest linear time.</p>
<pre data-code-wrap="txt"><code class="lang-nohighlight">&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f02998ded00&gt;
example(coords, yz_plane)
setup: from __main__ import example
  8.28 ms
  1 measurement, 100 runs , 1 thread
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f02998dee50&gt;
example(coords, yz_plane)
setup: from __main__ import example
  5.80 ms
  1 measurement, 100 runs , 1 thread
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f02998ded00&gt;
example2(coords, yz_plane, xy_plane, xz_plane)
setup: from __main__ import example2
  19.76 ms
  1 measurement, 100 runs , 1 thread
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f02998dee50&gt;
example2(coords, yz_plane, xy_plane, xz_plane)
setup: from __main__ import example2
  21.34 ms
  1 measurement, 100 runs , 1 thread

</code></pre>
<pre><code class="lang-python">import torch
import torch.nn.functional as F

coords = torch.randn((1, 512, 512, 2, 3)).cuda()
yz_plane = torch.randn((1, 32, 256, 256)).cuda()
xy_plane = torch.randn((1, 32, 256, 256)).cuda()
xz_plane = torch.randn((1, 32, 256, 256)).cuda()

def example(coords, yz_plane):
    B, H, W, D, C = coords.size()
    coords = coords.view(B, -1, C).unsqueeze(1)  # (B, 1, N, 3), N = 262,144
    # yz_plane with size of (1, 32, 256, 256)
    sample_yz_feat = F.grid_sample(yz_plane.repeat(B, 1, 1, 1), coords[..., [1, 2]].flip(-1), align_corners=True).permute(0, 2, 3, 1).squeeze(1)  # B N 32
    return sample_yz_feat

def example2(coords, yz_plane, xz_plane, xy_plane):
    B, H, W, D, C = coords.size()
    coords = coords.view(B, -1, C).unsqueeze(1)  # (B, 1, N, 3), N = 262,144
    # yz_plane, xz_plane, xy_plane with size of (1, 32, 256, 256)
    sample_yz_feat = F.grid_sample(yz_plane.repeat(B, 1, 1, 1), coords[..., [1, 2]].flip(-1), align_corners=True).permute(0, 2, 3, 1).squeeze(1)  # B N 32
    sample_xz_feat = F.grid_sample(xz_plane.repeat(B, 1, 1, 1), coords[..., [0, 2]].flip(-1), align_corners=True).permute(0, 2, 3, 1).squeeze(1)
    sample_xy_feat = F.grid_sample(xy_plane.repeat(B, 1, 1, 1), coords[..., [0, 1]].flip(-1), align_corners=True).permute(0, 2, 3, 1).squeeze(1)
    return sample_yz_feat


import torch.utils.benchmark as benchmark

t0 = benchmark.Timer(
    stmt='example(coords, yz_plane)',
    setup='from __main__ import example',
    globals={'coords': coords, 'yz_plane': yz_plane})

t1 = benchmark.Timer(
    stmt='example2(coords, yz_plane, xy_plane, xz_plane)',
    setup='from __main__ import example2',
    globals={'coords': coords, 'yz_plane': yz_plane, 'xy_plane': xy_plane, 'xz_plane': xz_plane})

print(t0.timeit(100))
print(t0.timeit(100))
print(t1.timeit(100))
print(t1.timeit(100))
</code></pre>

qiminchen: <p>sorry for being unclear, <code>N=262144</code></p>

kingsj0405: <p>I updated but it’s same.<br>
Can you reproduce the problem like provided code and log?</p>
<p>(Updated)<br>
I’ve also experienced similar problem on my complex training code.<br>
But the problem is resolved at some time while large chunk modification without clarifying the exact problem.<br>
So I’m interested in your case to clarify the problem.</p>

qiminchen: <p>It seems like testing <code>F.grid_sample</code> individually works fine, however, it is quite slow when I test it on my training code, so weird. Let me see if I can find out. Thanks anyway</p>