EmbeddingBag(num_embeddings, embedding_dim, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, mode='mean', sparse=False, _weight=None, include_last_offset=False)¶
Computes sums or means of ‘bags’ of embeddings, without instantiating the intermediate embeddings.
For bags of constant length and no
per_sample_weightsand 2D inputs, this class
EmbeddingBagis much more time and memory efficient than using a chain of these operations.
EmbeddingBag also supports per-sample weights as an argument to the forward pass. This scales the output of the Embedding before performing a weighted reduction as specified by
per_sample_weights`is passed, the only supported
"sum", which computes a weighted sum according to
num_embeddings (int) – size of the dictionary of embeddings
embedding_dim (int) – the size of each embedding vector
max_norm (float, optional) – If given, each embedding vector with norm larger than
max_normis renormalized to have norm
norm_type (float, optional) – The p of the p-norm to compute for the
scale_grad_by_freq (boolean, optional) – if given, this will scale gradients by the inverse of frequency of the words in the mini-batch. Default
False. Note: this option is not supported when
mode (string, optional) –
"max". Specifies the way to reduce the bag.
"sum"computes the weighted sum, taking
"mean"computes the average of the values in the bag,
"max"computes the max value over each bag. Default:
sparse (bool, optional) – if
True, gradient w.r.t.
weightmatrix will be a sparse tensor. See Notes for more details regarding sparse gradients. Note: this option is not supported when
include_last_offset (bool, optional) – if
offsetshas one additional element, where the last element is equivalent to the size of indices. This matches the CSR format.
~EmbeddingBag.weight (Tensor) – the learnable weights of the module of shape (num_embeddings, embedding_dim) initialized from .
input(IntTensor or LongTensor),
offsets(IntTensor or LongTensor, optional), and
offsetshave to be of the same type, either int or long
inputis 2D of shape (B, N),
it will be treated as
Bbags (sequences) each of fixed length
N, and this will return
Bvalues aggregated in a way depending on the
offsetsis ignored and required to be
Nonein this case.
inputis 1D of shape (N),
it will be treated as a concatenation of multiple bags (sequences).
offsetsis required to be a 1D tensor containing the starting index positions of each bag in
input. Therefore, for
offsetsof shape (B),
inputwill be viewed as having
Bbags. Empty bags (i.e., having 0-length) will have returned vectors filled by zeros.
- per_sample_weights (Tensor, optional): a tensor of float / double weights, or None
to indicate all weights should be taken to be
1. If specified,
per_sample_weightsmust have exactly the same shape as input and is treated as having the same
offsets, if those are not
None. Only supported for
Output shape: (B, embedding_dim)
>>> # an Embedding module containing 10 tensors of size 3 >>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum') >>> # a batch of 2 samples of 4 indices each >>> input = torch.LongTensor([1,2,4,5,4,3,2,9]) >>> offsets = torch.LongTensor([0,4]) >>> embedding_sum(input, offsets) tensor([[-0.8861, -5.4350, -0.0523], [ 1.1306, -2.5798, -1.0044]])
from_pretrained(embeddings, freeze=True, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, mode='mean', sparse=False, include_last_offset=False)¶
Creates EmbeddingBag instance from given 2-dimensional FloatTensor.
embeddings (Tensor) – FloatTensor containing weights for the EmbeddingBag. First dimension is being passed to EmbeddingBag as ‘num_embeddings’, second as ‘embedding_dim’.
freeze (boolean, optional) – If
True, the tensor does not get updated in the learning process. Equivalent to
embeddingbag.weight.requires_grad = False. Default:
max_norm (float, optional) – See module initialization documentation. Default:
norm_type (float, optional) – See module initialization documentation. Default
scale_grad_by_freq (boolean, optional) – See module initialization documentation. Default
mode (string, optional) – See module initialization documentation. Default:
sparse (bool, optional) – See module initialization documentation. Default:
include_last_offset (bool, optional) – See module initialization documentation. Default:
>>> # FloatTensor containing pretrained weights >>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]]) >>> embeddingbag = nn.EmbeddingBag.from_pretrained(weight) >>> # Get embeddings for index 1 >>> input = torch.LongTensor([[1, 0]]) >>> embeddingbag(input) tensor([[ 2.5000, 3.7000, 4.6500]])