EmbeddingBag#

classtorch.nn.modules.sparse.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,padding_idx=None,device=None,dtype=None)[source]#

Compute sums or means of ‘bags’ of embeddings, without instantiating the intermediate embeddings.

For bags of constant length, noper_sample_weights, no indices equal topadding_idx,and with 2D inputs, this class

• withmode="sum" is equivalent toEmbedding followed bytorch.sum(dim=1),

• withmode="mean" is equivalent toEmbedding followed bytorch.mean(dim=1),

• withmode="max" is equivalent toEmbedding followed bytorch.max(dim=1).

However,EmbeddingBag is much more time and memory efficient than using a chain of theseoperations.

EmbeddingBag also supports per-sample weights as an argument to the forwardpass. This scales the output of the Embedding before performing a weightedreduction as specified bymode. Ifper_sample_weights is passed, theonly supportedmode is"sum", which computes a weighted sum according toper_sample_weights.

Parameters

• 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 thanmax_normis renormalized to have normmax_norm.

• norm_type (float,optional) – The p of the p-norm to compute for themax_norm option. Default2.

• scale_grad_by_freq (bool,optional) – if given, this will scale gradients by the inverse of frequency ofthe words in the mini-batch. DefaultFalse.Note: this option is not supported whenmode="max".

• mode (str,optional) –"sum","mean" or"max". Specifies the way to reduce the bag."sum" computes the weighted sum, takingper_sample_weightsinto consideration."mean" computes the average of the valuesin the bag,"max" computes the max value over each bag.Default:"mean"

• sparse (bool,optional) – ifTrue, gradient w.r.t.weight matrix will be a sparse tensor. SeeNotes for more details regarding sparse gradients. Note: this option is notsupported whenmode="max".

• include_last_offset (bool,optional) – ifTrue,offsets has one additional element, where the last elementis equivalent to the size ofindices. This matches the CSR format.

• padding_idx (int,optional) – If specified, the entries atpadding_idx do not contribute to thegradient; therefore, the embedding vector atpadding_idx is not updatedduring training, i.e. it remains as a fixed “pad”. For a newly constructedEmbeddingBag, the embedding vector atpadding_idx will default to allzeros, but can be updated to another value to be used as the padding vector.Note that the embedding vector atpadding_idx is excluded from thereduction.

Variables

weight (Tensor) – the learnable weights of the module of shape(num_embeddings, embedding_dim)initialized from$\mathcal{N}(0,1)$N(0,1).

Examples:

>>># an EmbeddingBag 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.tensor([1,2,4,5,4,3,2,9],dtype=torch.long)>>>offsets=torch.tensor([0,4],dtype=torch.long)>>>embedding_sum(input,offsets)tensor([[-0.8861, -5.4350, -0.0523],        [ 1.1306, -2.5798, -1.0044]])>>># Example with padding_idx>>>embedding_sum=nn.EmbeddingBag(10,3,mode='sum',padding_idx=2)>>>input=torch.tensor([2,2,2,2,4,3,2,9],dtype=torch.long)>>>offsets=torch.tensor([0,4],dtype=torch.long)>>>embedding_sum(input,offsets)tensor([[ 0.0000,  0.0000,  0.0000],        [-0.7082,  3.2145, -2.6251]])>>># An EmbeddingBag can be loaded from an Embedding like so>>>embedding=nn.Embedding(10,3,padding_idx=2)>>>embedding_sum=nn.EmbeddingBag.from_pretrained(        embedding.weight,        padding_idx=embedding.padding_idx,        mode='sum')

forward(input,offsets=None,per_sample_weights=None)[source]#

Forward pass of EmbeddingBag.

Parameters

• input (Tensor) – Tensor containing bags of indices into the embedding matrix.

• offsets (Tensor,optional) – Only used wheninput is 1D.offsets determinesthe starting index position of each bag (sequence) ininput.

• per_sample_weights (Tensor,optional) – a tensor of float / double weights, or Noneto indicate all weights should be taken to be1. If specified,per_sample_weightsmust have exactly the same shape as input and is treated as having the sameoffsets, if those are notNone. Only supported formode='sum'.

Returns

Tensor output shape of(B, embedding_dim).

Return type

Tensor

Note

A few notes aboutinput andoffsets:

• input andoffsets have to be of the same type, either int or long

• Ifinput is 2D of shape(B, N), it will be treated asB bags (sequences)each of fixed lengthN, and this will returnB values aggregated in a waydepending on themode.offsets is ignored and required to beNone in this case.

• Ifinput is 1D of shape(N), it will be treated as a concatenation ofmultiple bags (sequences).offsets is required to be a 1D tensor containing thestarting index positions of each bag ininput. Therefore, foroffsets of shape(B),input will be viewed as havingB bags. Empty bags (i.e., having 0-length) will havereturned vectors filled by zeros.

classmethodfrom_pretrained(embeddings,freeze=True,max_norm=None,norm_type=2.0,scale_grad_by_freq=False,mode='mean',sparse=False,include_last_offset=False,padding_idx=None)[source]#

Create EmbeddingBag instance from given 2-dimensional FloatTensor.

Parameters

• embeddings (Tensor) – FloatTensor containing weights for the EmbeddingBag.First dimension is being passed to EmbeddingBag as ‘num_embeddings’, second as ‘embedding_dim’.

• freeze (bool,optional) – IfTrue, the tensor does not get updated in the learning process.Equivalent toembeddingbag.weight.requires_grad=False. Default:True

• max_norm (float,optional) – See module initialization documentation. Default:None

• norm_type (float,optional) – See module initialization documentation. Default2.

• scale_grad_by_freq (bool,optional) – See module initialization documentation. DefaultFalse.

• mode (str,optional) – See module initialization documentation. Default:"mean"

• sparse (bool,optional) – See module initialization documentation. Default:False.

• include_last_offset (bool,optional) – See module initialization documentation. Default:False.

• padding_idx (int,optional) – See module initialization documentation. Default:None.

Return type

EmbeddingBag

Examples:

>>># 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]])