Nn.functional.one_hot taking exceptionally long

import-antigravity: <p>I’m training a GAN on a computing cluster with pytorch lightning and I’m trying to investigate what’s causing some bottlenecks with a profiler. As you can see, <code>one_hot()</code> is taking a very long time, with each call averaging 0.03 seconds:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/3/0/30b9d5af98fc9bbf1a22498af717275fdc72a8e8.png" alt="Screen Shot 2021-03-04 at 9.10.15 PM" data-base62-sha1="6X34bU2PfIo8qfijBvvqdBVcjEs" width="578" height="17"></p>
<p>Here is the relevant code, sampling a conditional probability distribution from the generator:</p>
<pre><code class="lang-auto">    def sample_G(self, n) -&gt; Tuple[Tensor, Tensor]:
        """
        Generates n generator samples. In this conditional GAN case we're sampling G with a uniform class label prior.

        @param n: Number of samples
        @return: Tensor of generator samples
        """
        z = self.sample_z(n)
        y = tr.floor(tr.rand(n, device=next(self.G.parameters()).device) * 10).long()
        y_one_hot = F.one_hot(y, num_classes=10).float()

        return self.G(z, y_one_hot)
</code></pre>
<p>Any help would be greatly appreciated.</p>

ptrblck: <p>Could you post the input shape to <code>F.one_hot</code>?</p>

import-antigravity: <p>Here:<br>
<code>torch.Size([128])</code></p>
<p>It’s just the batch size</p>

ptrblck: <p>I cannot reproduce the 30ms using the PyTorch <code>1.8.0</code> binaries:</p>
<pre><code class="lang-python">import torch
import torch.nn.functional as F
import time

y = torch.randint(0, 10, (128,), device='cuda')

# warmup
for _ in range(10):
    y_one_hot = F.one_hot(y, num_classes=10)

nb_iters = 1000
torch.cuda.synchronize()
t0 = time.time()
for _ in range(nb_iters):
    y_one_hot = F.one_hot(y, num_classes=10)
torch.cuda.synchronize()
t1 = time.time()
print('{:.3f}us/iter'.format(((t1 - t0)/nb_iters)*1e6))
</code></pre>
<p>Output</p>
<pre><code class="lang-python">83.642us/iter
</code></pre>

import-antigravity: <p>I’ll try re-running the code with <code>time.time()</code></p>

import-antigravity: <p>Getting between 100-6000 us per call. It looks like some calls are taking really long and others are not.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/5/65e9b3766d002b973e613cd996254be60a34763d.png" alt="Screen Shot 2021-03-06 at 1.18.15 PM" data-base62-sha1="exyVs3z8inQRMjEngRchnYnnkPr" width="251" height="363"></p>
<p>Also I’m on pytorch 1.7.1</p>
<p>Here’s the code again:</p>
<pre><code class="lang-auto">def sample_G(self, n) -&gt; Tuple[Tensor, Tensor]:
    """
    Generates n generator samples. In this conditional GAN case we're sampling G with a uniform class label prior.

    @param n: Number of samples
    @return: Tensor of generator samples
    """
    z = self.sample_z(n)
    y = tr.floor(tr.rand(n, device=next(self.G.parameters()).device) * 10).long()
    t0 = time.time()
    y_one_hot = F.one_hot(y, num_classes=10).float()
    t1 = time.time()
    print(t1 - t0)

    return self.G(z, y_one_hot)
</code></pre>

KFrank: <p>Hi Robbie!</p>
<aside class="quote no-group" data-username="import-antigravity" data-post="6" data-topic="113776" data-full="true">
<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/import-antigravity/48/30128_2.png" class="avatar"> import-antigravity:</div>
<blockquote>
<p>Getting between 100-6000 us per call. It looks like some calls are taking really long and others are not.</p>
</blockquote>
</aside>
<p>Please note the <code>torch.cuda.synchronize()</code> in <a class="mention" href="/u/ptrblck">@ptrblck</a>’s sample<br>
code.  Without it you will get unreliable gpu timings.</p>
<p>(Also, please avoid posting screenshots of textual information.  It<br>
breaks accessibility, searchability, and copy-paste.)</p>
<p>Best.</p>
<p>K. Frank</p>

import-antigravity: <p>Getting 100-200 us which appears to be pretty normal</p>
<pre><code class="lang-auto">(pid=14209) 172.13821411132812
(pid=14209) 195.50323486328125
(pid=14204) 175.23765563964844
(pid=14209) 151.3957977294922
(pid=14204) 189.30435180664062
(pid=14204) 145.43533325195312
(pid=14204) 207.1857452392578
(pid=14209) 212.4309539794922
(pid=14209) 130.8917999267578
(pid=14209) 186.44332885742188
(pid=14204) 169.27719116210938
(pid=14209) 171.661376953125
(pid=14204) 152.34947204589844
(pid=14204) 128.5076141357422
(pid=14204) 190.25802612304688
(pid=14209) 181.43653869628906
(pid=14209) 203.60946655273438
(pid=14209) 154.4952392578125
(pid=14204) 209.808349609375
(pid=14204) 244.140625
(pid=14209) 182.15179443359375
(pid=14204) 129.22286987304688
(pid=14204) 150.44212341308594
(pid=14209) 229.35867309570312
</code></pre>