I’ve been trying to use the new EmbeddingBag layer to improve the performance of parts of my models where I first perform indexing into an Embedding layer, then sum or mean operations on the resulting embeddings.
Unfortunately, I have found using EmbeddingBag to make my models slower, 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 EmbeddingBag layer is 40% slower than Embedding then sum on a CPU, and about 25% slower on a GPU.
I used the following snippet (on PyTorch 0.2.0):
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))

I generated these timings:

CPU: Embedding: 0.10044050216674805, EmbeddingBag: 0.1490950584411621, ratio: 0.6736675461741425

GPU: Embedding: 0.18432116508483887, EmbeddingBag: 0.2421271800994873, ratio: 0.7612576374494734


Am I doing something that’s obviously wrong?