Repository: pytorch/pytorch
Pull request: Embedding related fixes
This PR aims to 
* fix https://github.com/pytorch/pytorch/issues/2979 (adding functional `EmbeddingBag`
* fix https://github.com/pytorch/pytorch/issues/3104 (missing docstring for keyword argument in Embedding)
* fix a bug in `Embedding` docstring (calling `variable[0].zero_()` results in `RuntimeError`)
* convert `nn.EmbeddingBag` to new style `Function`.

I've moved the core logic out of `EmbeddingBag` `Module` to the functional interface, since we would either have to dulpicate behaviour of `forward`. I hope this decision makes sense.

Problem statement:
No EmbeddingBag Functional
Hi,

It would be nice to have an EmbeddingBag functional implementation.

Thanks,
Shawn 

sparse argument in Embedding and EmbeddingBag is not documented
The docstring for these two layers does not mention about the `sparse=True/False` argument they receive during initialization.

Hints from issue discussion:
Hi, I'm preparing a PR that'll fix that.

Commit message: Docfix: remove torch.manual_seed from example code.

Patch to review:
diff --git a/torch/nn/_functions/thnn/sparse.py b/torch/nn/_functions/thnn/sparse.py
index ecdcf93ae30f..6d75eb9d2d4e 100644
--- a/torch/nn/_functions/thnn/sparse.py
+++ b/torch/nn/_functions/thnn/sparse.py
@@ -120,35 +120,37 @@ def backward(ctx, grad_output):
 
 class EmbeddingBag(Function):
 
-    def __init__(self, max_norm, norm_type, scale_grad_by_freq, mode):
-        super(EmbeddingBag, self).__init__()
-        self.max_norm = max_norm
-        self.norm_type = norm_type
-        self.scale_grad_by_freq = scale_grad_by_freq
-        self._indices = None
-        assert mode is not None
+    @staticmethod
+    def _renorm(ctx, indices, weight, max_norm, norm_type):
+        # clone indices since LookupTable_renorm modifies it in-place
+        ctx._backend.LookupTable_renorm(
+            ctx._backend.library_state,
+            indices.clone().view(-1),
+            weight,
+            max_norm,
+            norm_type
+        )
+
+    @classmethod
+    def forward(cls, ctx, weight, indices, offsets,
+                max_norm, norm_type, scale_grad_by_freq, mode):
+
+        ctx.max_norm = max_norm
+        ctx.norm_type = norm_type
+        ctx.scale_grad_by_freq = scale_grad_by_freq
+
         if mode == 'sum':
-            self.mode = MODE_SUM
+            ctx.mode = MODE_SUM
         elif mode == 'mean':
-            self.mode = MODE_MEAN
+            ctx.mode = MODE_MEAN
         else:
             raise ValueError("mode needs to be 'sum' or 'mean', but got {}"
                              .format(mode))
 
-    def _renorm(self, indices, weight):
-        self._backend.LookupTable_renorm(
-            self._backend.library_state,
-            indices,
-            weight,
-            self.max_norm,
-            self.norm_type
-        )
-
-    def forward(self, weight, indices, offsets):
-        assert not self.needs_input_grad[1], "EmbeddingBag doesn't " \
+        assert not ctx.needs_input_grad[1], "EmbeddingBag doesn't " \
             "compute the gradient w.r.t. the indices"
 
-        assert not self.needs_input_grad[2], "EmbeddingBag doesn't " \
+        assert not ctx.needs_input_grad[2], "EmbeddingBag doesn't " \
             "compute the gradient w.r.t. the offsets"
 
         assert indices.dim() == 1
@@ -164,57 +166,60 @@ def forward(self, weight, indices, offsets):
                              " ({}), but got offsets[-1] of {}"
                              .format(indices.size(0), offsets[-1]))
 
-        self._backend = type2backend[type(weight)]
-        self._weight_size = weight.size()
-        self._offset2bag = offsets.new()
+        ctx._backend = type2backend[type(weight)]
+        ctx._weight_size = weight.size()
+        ctx._offset2bag = offsets.new()
 
-        self.save_for_backward(indices)
+        ctx.save_for_backward(indices)
 
         indices = indices.contiguous().view(-1)
         output = weight.new()
-        if self.max_norm is not None:
-            self._renorm(indices, weight)
+
+        if ctx.max_norm is not None:
+            cls._renorm(ctx, indices, weight, max_norm=max_norm, norm_type=norm_type)
 
         if weight.is_cuda:
-            if self.mode == MODE_MEAN:
-                self.bag_size = offsets.new().resize_(offsets.size())
+            if ctx.mode == MODE_MEAN:
+                ctx.bag_size = offsets.new().resize_(offsets.size())
             else:
-                self.bag_size = None
+                ctx.bag_size = None
 
-            self._backend.LookupTableBag_updateOutput(
-                self._backend.library_state,
+            ctx._backend.LookupTableBag_updateOutput(
+                ctx._backend.library_state,
                 indices,
                 offsets,
                 weight,
                 output,
-                self._offset2bag,
-                self.mode,
-                self.bag_size
+                ctx._offset2bag,
+                ctx.mode,
+                ctx.bag_size
             )
         else:
             # slow CPU implementation
             index_output = torch.index_select(weight, 0, indices)
             # indices = [1, 2, 30, 100, 12], offsets = [0, 2, 3]
-            self._offset2bag.resize_(indices.size(0)).zero_()  # offset2bag = [0 0 0 0 0]
-            self._offset2bag.index_fill_(0, offsets, 1)  # offset2bag = [1 0 1 0 1]
-            self._offset2bag[0] = 0  # offset2bag = [0 0 1 0 1]
-            self._offset2bag = self._offset2bag.cumsum(0)  # offset2bag = [0 0 1 1 2]
+            ctx._offset2bag.resize_(indices.size(0)).zero_()  # offset2bag = [0 0 0 0 0]
+            ctx._offset2bag.index_fill_(0, offsets, 1)  # offset2bag = [1 0 1 0 1]
+            ctx._offset2bag[0] = 0  # offset2bag = [0 0 1 0 1]
+            ctx._offset2bag = ctx._offset2bag.cumsum(0)  # offset2bag = [0 0 1 1 2]
             output.resize_(offsets.size(0), weight.size(1)).zero_()
-            output.index_add_(0, self._offset2bag, index_output)
-            if self.mode == MODE_MEAN:
+            output.index_add_(0, ctx._offset2bag, index_output)
+            if ctx.mode == MODE_MEAN:
                 if offsets.size(0) == 1:
-                    self.bag_size = indices.size(0)
+                    ctx.bag_size = indices.size(0)
                 else:
-                    self.bag_size = weight.new().resize_(offsets.size())
-                    self.bag_size[:-1] = offsets[1:] - offsets[:-1]
-                    self.bag_size[-1] = indices.size(0) - offsets[-1]
-                    self.bag_size = self.bag_size[:, None].expand_as(output)
-                output /= self.bag_size
+                    ctx.bag_size = weight.new().resize_(offsets.size())
+                    ctx.bag_size[:-1] = offsets[1:] - offsets[:-1]
+                    ctx.bag_size[-1] = indices.size(0) - offsets[-1]
+                    ctx.bag_size = ctx.bag_size[:, None].expand_as(output)
+                output /= ctx.bag_size
 
         return output
 
-    def backward(self, grad_output):
-        indices, = self.saved_tensors
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        indices, = ctx.saved_tensors
         indices = indices.contiguous().view(-1)
         grad_output = grad_output.contiguous()
 
@@ -227,44 +232,44 @@ def backward(self, grad_output):
                 _count = torch.IntTensor()
                 _sorted = _indices = None
 
-        grad_weight = grad_output.new(self._weight_size).zero_()
+        grad_weight = grad_output.new(ctx._weight_size).zero_()
 
         if grad_output.is_cuda:
-            self._backend.LookupTableBag_accGradParameters(
-                self._backend.library_state,
+            ctx._backend.LookupTableBag_accGradParameters(
+                ctx._backend.library_state,
                 indices,
                 grad_output,
                 grad_weight,
-                self._offset2bag,
+                ctx._offset2bag,
                 _count,
                 _sorted,
                 _indices,
-                self.scale_grad_by_freq,
-                self.mode,
-                self.bag_size,
+                ctx.scale_grad_by_freq,
+                ctx.mode,
+                ctx.bag_size,
                 1
             )
         else:
             # slow CPU implementation
-            if self.mode == MODE_MEAN:
+            if ctx.mode == MODE_MEAN:
                 # divide by average count
-                grad_output = grad_output / self.bag_size
+                grad_output = grad_output / ctx.bag_size
 
-            index_grad_output = grad_output.index_select(0, self._offset2bag)
-            self._backend.LookupTable_accGradParameters(
-                self._backend.library_state,
+            index_grad_output = grad_output.index_select(0, ctx._offset2bag)
+            ctx._backend.LookupTable_accGradParameters(
+                ctx._backend.library_state,
                 indices,
                 index_grad_output,
                 grad_weight,
                 _count,
                 _sorted,
                 _indices,
-                self.scale_grad_by_freq,
+                ctx.scale_grad_by_freq,
                 -1,
                 1
             )
 
-        return grad_weight, None, None
+        return grad_weight, None, None, None, None, None, None
 
 
 _all_functions.append(EmbeddingBag)
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
index 163ac3d77bb6..a4232d77e21f 100644
--- a/torch/nn/functional.py
+++ b/torch/nn/functional.py
@@ -655,6 +655,7 @@ def embedding(input, embedding_matrix,
         norm_type (float, optional): The p of the p-norm to compute for the max_norm option
         scale_grad_by_freq (boolean, optional): if given, this will scale gradients by the frequency of
                                                 the words in the mini-batch.
+        sparse (boolean, optional): if True, gradient w.r.t. weight matrix will be a sparse tensor.
 
     Shape:
         - Input: LongTensor `(N, W)`, N = mini-batch, W = number of indices to extract per mini-batch
@@ -684,8 +685,9 @@ def embedding(input, embedding_matrix,
         [torch.FloatTensor of size 2x4x3]
 
         >>> # example with padding_idx
-        >>> embedding_matrix = Variable(torch.rand(10, 3))
-        >>> embedding_matrix[0].zero_()
+        >>> weights = torch.rand(10, 3)
+        >>> weights[0, :].zero_()
+        >>> embedding_matrix = Variable(weights)
         >>> input = Variable(torch.LongTensor([[0,2,0,5]]))
         >>> torch.nn.functional.embedding(input, embedding_matrix)
 
@@ -705,6 +707,85 @@ def embedding(input, embedding_matrix,
     )
 
 
+def embedding_bag(embedding_matrix, indices, offsets=None,
+                  max_norm=None, norm_type=2, scale_grad_by_freq=False, mode='mean'):
+    r"""Computes sums or means of 'bags' of embeddings, without instantiating the
+        intermediate embeddings.
+
+        For bags of constant length,
+            * embedding_bag with `mode=sum` is equivalent to nn.functional.embedding followed by `torch.sum(dim=1)`
+            * with `mode=mean` is equivalent to nn.functional.embedding followed by `torch.mean(dim=1)`
+
+        However, embedding_bag is much more time and memory efficient than using a chain of these
+        operations.
+
+        Args:
+            embedding_matrix: FloatTensor, where number of rows should correspond to the maximum possible index + 1,
+                              number of columns is the embedding size
+            indices (N or BxN): LongTensor containing the indices of the embeddings to extract.
+                                When `input` is 1D Tensor of shape `N`, an `offsets` Tensor is given, that contains the
+                                starting position of each new sequence in the mini-batch.
+            offsets (B or None): LongTensor containing the starting positions of each sample in a mini-batch of variable
+                                 length sequences. If `input` is 2D (BxN), then offsets does not need to be given,
+                                 as the `input` is treated as a mini-batch of fixed length sequences of length `N` each.
+            max_norm (float, optional): If given, will renormalize the embeddings to always have a norm lesser than this
+            norm_type (float, optional): The p of the p-norm to compute for the max_norm option
+            scale_grad_by_freq (boolean, optional): if given, this will scale gradients by the frequency of
+                                                    the words in the dictionary.
+            mode (string, optional): 'sum' | 'mean'. Specifies the way to reduce the bag. Default: 'mean'
+
+        Shape:
+            - Embedding_matrix: FloatTensor `(V, embedding_dim)`,
+                                V = number of embeddings, embedding_dim = embedding size
+            - Input: LongTensor `N`, N = number of embeddings to extract
+                     (or) LongTensor `BxN`, B = number of sequences in mini-batch,
+                                            N = number of embeddings per sequence
+            - Offsets: LongTensor `B`, B = number of bags. The values are the
+                       offsets in `input` for each bag, i.e. the cumsum of lengths.
+                       Offsets is not given if Input is 2D `BxN` Tensor,
+                       the input is considered to be of fixed-length sequences
+            - Output: `(B, embedding_dim)`
+
+        Examples::
+
+            >>> # an Embedding module containing 10 tensors of size 3
+            >>> embedding_matrix = Variable(torch.rand(10, 3))
+            >>> # a batch of 2 samples of 4 indices each
+            >>> input = Variable(torch.LongTensor([1,2,4,5,4,3,2,9]))
+            >>> offsets = Variable(torch.LongTensor([0,4]))
+            >>> embedding_bag(embedding_matrix, input, offsets)
+
+            Variable containing:
+            -1.1840 -0.2547 -0.5860
+            -0.7126  0.0002 -0.3411
+            [torch.FloatTensor of size 2x3]
+
+        """
+    if indices.dim() == 2:
+        if offsets is not None:
+            raise ValueError("if input is 2D, then offsets has to be None"
+                             ", as input is treated is a mini-batch of"
+                             " fixed length sequences. However, found "
+                             "offsets of type {}".format(type(offsets)))
+        else:
+            offsets = Variable(torch.arange(0, indices.numel(), indices.size(1),
+                               out=indices.data.new().long()))
+            indices = indices.view(-1)
+
+    elif indices.dim() != 1:
+        raise ValueError("input has to be 1D or 2D Tensor,"
+                         " but got Tensor of dimension {}".format(indices.dim()))
+
+    if offsets is None:
+        raise ValueError("offsets has to be a 1D Tensor but got None")
+
+    return torch.nn.backends.thnn.backend.EmbeddingBag.apply(
+        embedding_matrix, indices, offsets,
+        max_norm, norm_type,
+        scale_grad_by_freq, mode
+    )
+
+
 def batch_norm(input, running_mean, running_var, weight=None, bias=None,
                training=False, momentum=0.1, eps=1e-5):
     if training:
diff --git a/torch/nn/modules/sparse.py b/torch/nn/modules/sparse.py
index 4286320a640d..a7966887198b 100644
--- a/torch/nn/modules/sparse.py
+++ b/torch/nn/modules/sparse.py
@@ -1,8 +1,9 @@
 import torch
-from torch.nn.parameter import Parameter
 from torch.autograd import Variable
+from torch.nn.parameter import Parameter
 
 from .module import Module
+from .. import functional as F
 
 
 class Embedding(Module):
@@ -20,6 +21,7 @@ class Embedding(Module):
         norm_type (float, optional): The p of the p-norm to compute for the max_norm option
         scale_grad_by_freq (boolean, optional): if given, this will scale gradients by the frequency of
                                                 the words in the mini-batch.
+        sparse (boolean, optional): if True, gradient w.r.t. weight matrix will be a sparse tensor.
 
     Attributes:
         weight (Tensor): the learnable weights of the module of shape (num_embeddings, embedding_dim)
@@ -192,26 +194,9 @@ def reset_parameters(self):
         self.weight.data.normal_(0, 1)
 
     def forward(self, input, offsets=None):
-        if input.dim() == 2:
-            if offsets is not None:
-                raise ValueError("if input is 2D, then offsets has to be None"
-                                 ", as input is treated is a mini-batch of"
-                                 " fixed length sequences. However, found "
-                                 "offsets of type {}".format(type(offsets)))
-            else:
-                offsets = Variable(torch.arange(0, input.numel(), input.size(1),
-                                   out=input.data.new().long()))
-                input = input.view(-1)
-        elif input.dim() != 1:
-            raise ValueError("input has to be 1D or 2D Tensor,"
-                             " but got Tensor of dimension {}".format(input.dim()))
-        if offsets is None:
-            raise ValueError("offsets has to be a 1D Tensor but got None")
-
-        return self._backend.EmbeddingBag(
-            self.max_norm, self.norm_type,
-            self.scale_grad_by_freq, mode=self.mode
-        )(self.weight, input, offsets)
+        return F.embedding_bag(self.weight, input, offsets,
+                               self.max_norm, self.norm_type,
+                               self.scale_grad_by_freq, self.mode)
 
     def __repr__(self):
         s = '{name}({num_embeddings}, {embedding_dim}'


Code review comments:
diff_hunk: @@ -655,6 +655,7 @@ def embedding(input, embedding_matrix,
         norm_type (float, optional): The p of the p-norm to compute for the max_norm option
         scale_grad_by_freq (boolean, optional): if given, this will scale gradients by the frequency of
                                                 the words in the mini-batch.
+        sparse (boolean, optional): if True, gradient w.r.t. weight matrix will be a sparse tensor.
path: torch/nn/functional.py
text: @user3:
I think we should have a warning that `sparse` should be used only if weights are a leaf (maybe we should check it?). I'm afraid that doing it in any other case will result in weird errors in backward, since a lot of autograd functions can't handle sparse tensors.

BTW @user1 @user2 can we do sparse optimizers with CUDA? I can't remember if this was finally implemented?

@user2:
I'm not aware of the latest developments but I think at least SGD/Adagrad should work (with the usual momentum/weight decay caveats).

@user3:
We should at least say under what conditions can this be used (i.e. restrict to leaves and mention which optimizers can handle sparse updates)

@author:
Should we specify which optimizers support sparse grad? Someone will have to remember to change it when new support is introduced.

BTW it seems only `sgd` supports `cuda` and `cpu`, and `adagrad` supports `cpu` only

diff_hunk: @@ -705,6 +707,85 @@ def embedding(input, embedding_matrix,
     )
 
 
+def embedding_bag(embedding_matrix, indices, offsets=None,
+                  max_norm=None, norm_type=2, scale_grad_by_freq=False, mode='mean'):
+    r"""Computes sums or means of 'bags' of embeddings, without instantiating the
+        intermediate embeddings.
+
+        For bags of constant length,
+            * embedding_bag with `mode=sum` is equivalent to nn.functional.embedding followed by `torch.sum(dim=1)`
+            * with `mode=mean` is equivalent to nn.functional.embedding followed by `torch.mean(dim=1)`
+
+        However, embedding_bag is much more time and memory efficient than using a chain of these
+        operations.
+
+        Args:
+            embedding_matrix: FloatTensor, where number of rows should correspond to the maximum possible index + 1,
+                              number of columns is the embedding size
+            indices (N or BxN): LongTensor containing the indices of the embeddings to extract.
+                                When `input` is 1D Tensor of shape `N`, an `offsets` Tensor is given, that contains the
+                                starting position of each new sequence in the mini-batch.
+            offsets (B or None): LongTensor containing the starting positions of each sample in a mini-batch of variable
+                                 length sequences. If `input` is 2D (BxN), then offsets does not need to be given,
+                                 as the `input` is treated as a mini-batch of fixed length sequences of length `N` each.
+            max_norm (float, optional): If given, will renormalize the embeddings to always have a norm lesser than this
+            norm_type (float, optional): The p of the p-norm to compute for the max_norm option
+            scale_grad_by_freq (boolean, optional): if given, this will scale gradients by the frequency of
+                                                    the words in the dictionary.
+            mode (string, optional): 'sum' | 'mean'. Specifies the way to reduce the bag. Default: 'mean'
+
+        Shape:
+            - Embedding_matrix: FloatTensor `(V, embedding_dim)`,
+                                V = number of embeddings, embedding_dim = embedding size
+            - Input: LongTensor `N`, N = number of embeddings to extract
+                     (or) LongTensor `BxN`, B = number of sequences in mini-batch,
+                                            N = number of embeddings per sequence
+            - Offsets: LongTensor `B`, B = number of bags. The values are the
+                       offsets in `input` for each bag, i.e. the cumsum of lengths.
+                       Offsets is not given if Input is 2D `BxN` Tensor,
+                       the input is considered to be of fixed-length sequences
+            - Output: `(B, embedding_dim)`
+
+        Examples::
+
+            >>> # an Embedding module containing 10 tensors of size 3
+            >>> embedding_matrix = Variable(torch.rand(10, 3))
+            >>> # a batch of 2 samples of 4 indices each
+            >>> input = Variable(torch.LongTensor([1,2,4,5,4,3,2,9]))
+            >>> offsets = Variable(torch.LongTensor([0,4]))
+            >>> embedding_bag(embedding_matrix, input, offsets)
+
+            Variable containing:
+            -1.1840 -0.2547 -0.5860
+            -0.7126  0.0002 -0.3411
+            [torch.FloatTensor of size 2x3]
+
+        """
+    if indices.dim() == 2:
+        if offsets is not None:
+            raise ValueError("if input is 2D, then offsets has to be None"
+                             ", as input is treated is a mini-batch of"
+                             " fixed length sequences. However, found "
+                             "offsets of type {}".format(type(offsets)))
+        else:
+            offsets = Variable(torch.arange(0, indices.numel(), indices.size(1),
+                               out=indices.data.new().long()))
+            indices = indices.view(-1)
+
+    elif indices.dim() != 1:
+        raise ValueError("input has to be 1D or 2D Tensor,"
+                         " but got Tensor of dimension {}".format(indices.dim()))
+
+    if offsets is None:
+        raise ValueError("offsets has to be a 1D Tensor but got None")
+
+    return torch.nn.backends.thnn.backend.EmbeddingBag.apply(
path: torch/nn/functional.py
text: @user1:
It's simpler to do `_functions.thnn.EmbeddingBag`

@author:
sorry, I got confused since `Embedding` also calls its function via `torch.nn`. 

@user1:
No worries. I think this is valid too, it's just unnecessarily long