EmbeddingBag vs Embedding performance

mkula: <p>I’ve been trying to use the new <code>EmbeddingBag</code> layer to improve the performance of parts of my models where I first perform indexing into an <code>Embedding</code> layer, then sum or mean operations on the resulting embeddings.</p>
<p>Unfortunately, I have found using <code>EmbeddingBag</code> to make my models <em>slower</em>, often more than three times slower. Running simple indexing operations in a loop suggests that, for the simple case of embedding indexing followed by a sum, the <code>EmbeddingBag</code> layer is 40% slower than <code>Embedding</code> then sum on a CPU, and about 25% slower on a GPU.</p>
<p>I used the following snippet (on PyTorch 0.2.0):</p>
<pre><code class="lang-auto">import time

import torch
from torch.nn import Embedding, EmbeddingBag
from torch.autograd import Variable


def time_layer(layer_class, repetitions=100):

    layer = layer_class(10000, 32)
    indices = Variable(torch.ones(1000).long().view(100, 10))

    offsets = Variable(torch.arange(0, indices.numel(), indices.size(1)).long())

    if torch.cuda.is_available():
        layer = layer.cuda()
        indices = indices.cuda()
        offsets = offsets.cuda()

    start = time.time()
    for _ in range(repetitions):

        if isinstance(layer, EmbeddingBag):
            result = layer(indices.view(-1), offsets).view(indices.size(0), -1)
        else:
            result = layer(indices)
            result = result.sum(1)

        loss = result.sum()
        loss.backward()

    print(result.size())

    return time.time() - start


if __name__ == '__main__':

    embedding_time = time_layer(Embedding)
    embedding_bag_time = time_layer(EmbeddingBag)

    print('Embedding: {}, EmbeddingBag: {}, ratio: {}'.format(
        embedding_time,
        embedding_bag_time,
        embedding_time / embedding_bag_time))
</code></pre>
<p>I generated these timings:</p>
<ul>
<li>CPU: <code>Embedding: 0.10044050216674805, EmbeddingBag: 0.1490950584411621, ratio: 0.6736675461741425</code>
</li>
<li>GPU: <code>Embedding: 0.18432116508483887, EmbeddingBag: 0.2421271800994873, ratio: 0.7612576374494734</code>
</li>
</ul>
<p>Am I doing something that’s obviously wrong?</p>

smth: <p>Currently, EmbeddingBag is only optimized on the GPU. The CPU implementation is quite naive. I’ll work on optimizing the CPU implementation for the next release.</p>

mkula: <p>Thank you for the speedy response!</p>
<p>I gleaned from the source comments that the CPU implementation is a temporary solution; however, even on the GPU <code>EmbeddingBag</code> is substantially slower. In the snippet I posted above, it’s ~25% slower; if you increase the size of the embedding layer and the number of indices used to</p>
<pre><code class="lang-auto">    layer = layer_class(100000, 256)
    indices = Variable(torch.ones(10 ** 5).long().view(10 ** 3, 10 ** 2))
</code></pre>
<p>it becomes twice as slow as the <code>Embedding</code> layer followed by a sum (measured on the Amazon <code>p2.xlarge</code> instance).</p>

mkula: <p><a class="mention" href="/u/smth">@smth</a> To add a little bit more detail, I ran <code>nvprof</code> on both embedding and embedding bag layers; the results are in this <a href="https://gist.github.com/maciejkula/cc5acc610af511c53489ba77fe29c010" rel="nofollow noopener">gist</a>.</p>
<p>The curx of the problem seems to be that the <code>cunn_LookupTableBag_accGradParametersKernel</code> is about twice as slow as <code>cunn_LookupTable_accGradParametersKernel</code>, far outweighing any gains from not allocating intermediate results when summing over some dimensions of embedding layers.</p>

egilm-ragulpr: <p>Any updates on this? Is it worth using EmbeddingBag?</p>
<p>For reference to googlers like me, see</p>
<aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/pull/5433" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
      <a href="https://github.com/martinraison" target="_blank" rel="nofollow noopener">
    <img alt="martinraison" src="https://avatars3.githubusercontent.com/u/2560662?v=4" class="thumbnail onebox-avatar" width="90" height="90">
  </a>

<h4>
  <a href="https://github.com/pytorch/pytorch/pull/5433" target="_blank" rel="nofollow noopener">speed up CPU EmbeddingBag (indexSelectAdd op)</a>
</h4>

<div class="date">
  by <a href="https://github.com/martinraison" target="_blank" rel="nofollow noopener">martinraison</a>
  on <a href="https://github.com/pytorch/pytorch/pull/5433" target="_blank" rel="nofollow noopener">02:53PM - 27 Feb 18</a>
</div>

<div class="github-commit-stats">
  <strong>5 commits</strong>
  changed <strong>3 files</strong>
  with <strong>71 additions</strong>
  and <strong>13 deletions</strong>.
</div>

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

<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/4441" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <a href="https://github.com/SsnL" rel="nofollow noopener">
<img src="https://avatars1.githubusercontent.com/u/5674597?v=2&amp;s=96" class="thumbnail onebox-avatar" width="96" height="96">
</a>

<h4><a href="https://github.com/pytorch/pytorch/issues/4441" target="_blank" rel="nofollow noopener">Issue: Implement CPU EmbeddingBag in C</a></h4>

<div class="date" style="margin-top:10px;">
	<div class="user" style="margin-top:10px;">
	opened by <a href="https://github.com/SsnL" target="_blank" rel="nofollow noopener">SsnL</a>
	on <a href="https://github.com/pytorch/pytorch/issues/4441" target="_blank" rel="nofollow noopener">2018-01-02</a>
	</div>
	<div class="user">
	closed by <a href="https://github.com/soumith" target="_blank" rel="nofollow noopener">soumith</a>
	on <a href="https://github.com/pytorch/pytorch/issues/4441" target="_blank" rel="nofollow noopener">2018-02-15</a>
	</div>
</div>

<pre class="content" style="white-space: pre-wrap;">to replace current slow Python implementation.</pre>

<div class="labels">
</div>

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


SimonW: <p>CPU embedding bag has been improved a lot by <span class="mention">@cpuhrsch</span> and many others.</p>
<p>Since you linked those issue/PRs, you should see that they are already solved/merged…  So I’m not sure why you asked.</p>

egilm-ragulpr: <p>Thanks, well the merged shows improvement but no comparison to see whether the claim in docs of large performance improvement over regular <code>Embedding</code> is true.</p>

SimonW: <p>I didn’t do benchmarks personally, but from the code it should be more efficient than achieving the same thing with Embedding. Let me know what you find if you do a benchmark yourself <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

xvdp: <p>I do realize that Im waking up an old thread. Apologies. However I dont get results as I expected: see code below</p>
<p>Running a slight modification of this benchmark code ( to load the gpu sufficientlly and sync cuda) on pytorch 1.3,  Driver Version: 435.21       CUDA Version: 10.1; TitanRTX  does not justify the claim that EmbeddingBag is more efficient. In CPU it is 50% better, in GPU marginally worse.<br>
caveat, Ive been working with vision mostly, only just now started looking at Embedings to combine text and vision. If I am using embeddings incorrectly please let me know.<br>
Perhaps to get the full utility of embedding/bag i need to run deeper network? Or?</p>
<pre><code class="lang-python">import time
import torch
import torch.nn as nn
def time_layer(layer_class, repetitions=100, device='cpu'):
    mul = 1
    if device == "cuda":
        mul = 10

    layer = layer_class(100000*mul, 512)
    layer.to(device=device)
    indices = torch.ones(10000*mul, device=device, dtype=torch.int64).view(1000*mul, 10)
    offsets = torch.arange(0, indices.numel(), indices.size(1), device=device, dtype=torch.int64)

    start = time.time()
    for _ in range(repetitions):
        if isinstance(layer, nn.EmbeddingBag):
            result = layer(indices.view(-1), offsets).view(indices.size(0), -1)
        else:
            result = layer(indices)
            result = result.sum(1)

        loss = result.sum()
        loss.backward()

    print(result.size(), result.device)
    if device == "cuda":
        torch.cuda.synchronize()
    return time.time() - start

if __name__ == '__main__':
    DEVICE = 'cpu'
    print(DEVICE)
    embedding_time = time_layer(nn.Embedding, device=DEVICE)
    embedding_bag_time = time_layer(nn.EmbeddingBag, device=DEVICE)

    print('Embedding:\t{},\nEmbeddingBag:\t{}'.format(
        embedding_time,
        embedding_bag_time))
    
    Z = torch.zeros(1, device='cuda') # isolate cuda init from timing
    DEVICE = 'cuda'
    print(DEVICE)

    embedding_time = time_layer(nn.Embedding, device=DEVICE)
    embedding_bag_time = time_layer(nn.EmbeddingBag, device=DEVICE)

    print('Embedding:\t{},\nEmbeddingBag:\t{}'.format(
        embedding_time,
        embedding_bag_time))
python embedbag.py 
``` Results
cpu
torch.Size([1000, 512]) cpu
torch.Size([1000, 512]) cpu
Embedding:      6.203256607055664,
EmbeddingBag:   4.765470743179321
cuda
torch.Size([10000, 512]) cuda:0
torch.Size([10000, 512]) cuda:0
Embedding:      1.6973094940185547,
EmbeddingBag:   1.9424304962158203

Is this code using the embeddingbag wrong?
Thank you</code></pre>