Python PyTorch embedding_bag用法及代码示例

本文简要介绍python语言中 torch.nn.functional.embedding_bag 的用法。

用法:

torch.nn.functional.embedding_bag(input, weight, offsets=None, max_norm=None, norm_type=2, scale_grad_by_freq=False, mode='mean', sparse=False, per_sample_weights=None, include_last_offset=False, padding_idx=None)

参数：

• input(LongTensor) -将包含索引袋的张量包含到嵌入矩阵中

• weight(Tensor) -行数等于最大可能索引 + 1 且列数等于嵌入大小的嵌入矩阵

• offsets(LongTensor,可选的) -仅在input 为 1D 时使用。 offsets 确定 input 中每个包(序列)的起始索引位置。

• max_norm(float,可选的) -如果给定，每个范数大于 max_norm 的嵌入向量将被重新归一化为范数 max_norm 。注意：这将就地修改weight。

• norm_type(float,可选的) -p -norm 中的 p 用于计算 max_norm 选项。默认 2 。

• scale_grad_by_freq(布尔值,可选的) -如果给定，这将通过小批量中单词频率的倒数来缩放梯度。默认 False 。注意：mode="max" 时不支持此选项。

• mode(string,可选的) -"sum" 、"mean" 或 "max" 。指定减少袋子的方式。默认值："mean"

• sparse(bool,可选的) -if True ，梯度 w.r.t. weight 将是一个稀疏张量。有关稀疏梯度的更多详细信息，请参阅 torch.nn.Embedding 下的注释。注意：mode="max" 时不支持此选项。

• per_sample_weights(Tensor,可选的) -浮点/双权重的张量，或 None 表示所有权重应为 1。如果指定，per_sample_weights 必须具有与输入完全相同的形状，并且被视为具有相同的 offsets，如果不是没有。

• include_last_offset(bool,可选的) -如果 True ，则偏移量的大小等于包的数量 + 1。最后一个元素是输入的大小，或者是最后一个包(序列)的结束索引位置。

• padding_idx(int,可选的) -如果指定，padding_idx 处的条目不会影响梯度；因此，padding_idx 处的嵌入向量在训练期间不会更新，即它保持为固定的 “pad”。请注意，padding_idx 处的嵌入向量不包括在缩减范围内。

计算嵌入的 bags 的总和、均值或最大值，而不实例化中间嵌入。

有关详细信息，请参阅 torch.nn.EmbeddingBag 。

注意

当给定 CUDA 设备上的张量时，此操作可能会产生不确定的梯度。有关详细信息，请参阅重现性。

形状：

• input(长张量)和offsets(长张量，可选)

  • 如果 input 是形状 (B, N) 的 2D，它将被视为每个固定长度 N 的 B 袋子(序列)，这将返回 B 值聚合的方式取决于 mode . offsets 被忽略，在这种情况下需要为 None。

  • 如果 input 是形状 (N) 的一维，它将被视为多个包(序列)的串联。 offsets 必须是一维张量，其中包含 input 中每个包的起始索引位置。因此，对于形状为 (B) 的 offsets ，input 将被视为具有 B 包。空包(即长度为 0)将返回由零填充的向量。

• weight(张量)：形状(num_embeddings, embedding_dim)模块的可学习权重

• per_sample_weights(张量，可选)。具有与 input 相同的形状。

• output：形状(B, embedding_dim)的聚合嵌入值

例子：

>>> # an Embedding module containing 10 tensors of size 3
>>> embedding_matrix = torch.rand(10, 3)
>>> # a batch of 2 samples of 4 indices each
>>> input = torch.tensor([1,2,4,5,4,3,2,9])
>>> offsets = torch.tensor([0,4])
>>> F.embedding_bag(input, embedding_matrix, offsets)
tensor([[ 0.3397,  0.3552,  0.5545],
        [ 0.5893,  0.4386,  0.5882]])

>>> # example with padding_idx
>>> embedding_matrix = torch.rand(10, 3)
>>> input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9])
>>> offsets = torch.tensor([0,4])
>>> F.embedding_bag(input, embedding_matrix, offsets, padding_idx=2, mode='sum')
tensor([[ 0.0000,  0.0000,  0.0000],
        [-0.7082,  3.2145, -2.6251]])