PyTorch torch.nn.EmbeddingBag 関数

PyTorch torch.nn 参考手册PyTorch torch.nn リファレンスマニュアル

torch.nn.EmbeddingBagこれは PyTorch の EmbeddingBag モジュールです。

これは複数の埋め込みベクトルを単一のベクトルに集約し、テキスト分類や高速処理によく使われます。

関数定義

torch.nn.EmbeddingBag(num_embeddings, embedding_dim, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, mode='mean', sparse=False, include_last_offset=False)

パラメータ：

  • mode：集約方法、オプション'mean'、'sum'、'max'

使用例

例 1: 基本的な使い方

サンプル

import torch
import torch.nn as nn

# EmbeddingBag
ebag = nn.EmbeddingBag(1000, 128, mode='mean')

# 単語インデックス
indices = torch.tensor([[1, 2, 3], [4, 5]])

# オフセット（各文の境界を表す）
offsets = torch.tensor([0, 3])

output = ebag(indices, offsets)

print(入力形状：, indices.shape)
print(出力形状：, output.shape)

例 2: 高速テキスト分類

サンプル

import torch
import torch.nn as nn

class FastText(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_classes):
        super(FastText, self).__init__()
        self.embedding = nn.EmbeddingBag(vocab_size, embed_dim, mode='mean')
        self.fc = nn.Linear(embed_dim, num_classes)

    def forward(self, text, offset):
        x = self.embedding(text, offset)
        return self.fc(x)

model = FastText(10000, 128, 2)

# 模擬データ
text = torch.randint(0, 10000, (100,))
offset = torch.tensor([0, 30, 60, 100])

output = model(text, offset)
print(出力形状：, output.shape)

例 3: 異なる集約方法

サンプル

import torch
import torch.nn as nn

for mode in ['mean', 'sum', 'max']:
    ebag = nn.EmbeddingBag(100, 32, mode=mode)
    indices = torch.tensor([1, 2, 3, 4])
    offsets = torch.tensor([0, 2])
    out = ebag(indices, offsets)
    print(f"{mode}:", out.shape)

使用シーン

  • テキスト分類: FastText
  • Bag of Words モデル：高速集約
  • 大規模テキスト：効率的な処理

注意：文の境界を示すために offsets パラメータを提供する必要があります。

PyTorch torch.nn 参考手册PyTorch torch.nn リファレンスマニュアル

その他の拡張