PyTorch torch.nn.Embedding 関数
PyTorch torch.nn リファレンスマニュアルtorch.nn.Embeddingこれは PyTorch で単語埋め込みに使用するモジュールです。
離散的な語彙インデックスを連続的なベクトル空間にマッピングするもので、自然言語処理における最も基本的な操作の一つです。
関数定義
torch.nn.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False)
パラメータの説明:
num_embeddings(int): 語彙サイズ、つまり埋め込み行列の行数。embedding_dim(int): 各埋め込みベクトルの次元。padding_idx(int): パディングインデックスを指定します。その埋め込みベクトルはゼロベクトルになります。デフォルトは None です。max_norm(float): None でない場合、埋め込みベクトルはこのノルムに正規化されます。デフォルトは None です。norm_type(float): ノルムを計算する次数。デフォルトは 2.0 です。scale_grad_by_freq(bool): 語彙の頻度に応じて勾配をスケーリングするかどうか。デフォルトは False です。sparse(bool): 重み行列がスパース行列であるかどうか。デフォルトは False です。
属性:
weight(Tensor): 形状が (num_embeddings, embedding_dim) の学習可能な重み。
使用例
例 1: 基本的な使い方
単語埋め込みを作成して使用する:
実例
import torch
import torch.nn as nn
# 埋め込み層を作成: 語彙サイズ 10000、埋め込み次元 256
embedding = nn.Embedding(num_embeddings=10000, embedding_dim=256)
# 単語インデックス（0 から始まる）
# 形状: (batch, seq_len)
input_indices = torch.tensor([[12, 45, 678], [901, 23, 56]])
# ルックアップで埋め込みベクトルを取得
output = embedding(input_indices)
print("入力インデックスの形状:", input_indices.shape)
print("出力埋め込みの形状:", output.shape) # (2, 3, 256)
# 埋め込み行列の形状を確認
print("埋め込み行列の形状:", embedding.weight.shape)
import torch.nn as nn
# 埋め込み層を作成: 語彙サイズ 10000、埋め込み次元 256
embedding = nn.Embedding(num_embeddings=10000, embedding_dim=256)
# 単語インデックス（0 から始まる）
# 形状: (batch, seq_len)
input_indices = torch.tensor([[12, 45, 678], [901, 23, 56]])
# ルックアップで埋め込みベクトルを取得
output = embedding(input_indices)
print("入力インデックスの形状:", input_indices.shape)
print("出力埋め込みの形状:", output.shape) # (2, 3, 256)
# 埋め込み行列の形状を確認
print("埋め込み行列の形状:", embedding.weight.shape)
例 2: padding_idx を使用
パディングインデックスを指定する:
実例
import torch
import torch.nn as nn
# パディング付きの埋め込み層を作成
embedding = nn.Embedding(num_embeddings=1000, embedding_dim=64, padding_idx=0)
# 0 がパディングとして使用される
input_indices = torch.tensor([[1, 2, 3], [4, 0, 0]]) # 2番目の文にはパディングがある
output = embedding(input_indices)
print("入力形状:", input_indices.shape)
print("出力形状:", output.shape)
print("パディングの埋め込みベクトル:", output[1, 1].tolist()) # すべて 0
print("パディングでない埋め込みベクトル:", output[0, 0].tolist()[:5]) # 非ゼロ
import torch.nn as nn
# パディング付きの埋め込み層を作成
embedding = nn.Embedding(num_embeddings=1000, embedding_dim=64, padding_idx=0)
# 0 がパディングとして使用される
input_indices = torch.tensor([[1, 2, 3], [4, 0, 0]]) # 2番目の文にはパディングがある
output = embedding(input_indices)
print("入力形状:", input_indices.shape)
print("出力形状:", output.shape)
print("パディングの埋め込みベクトル:", output[1, 1].tolist()) # すべて 0
print("パディングでない埋め込みベクトル:", output[0, 0].tolist()[:5]) # 非ゼロ
例 3: 事前学習済み単語ベクトル
事前学習済み単語ベクトルを読み込む:
実例
import torch
import torch.nn as nn
import numpy as np
# 事前学習済み単語ベクトルをシミュレート（実際には GloVe、Word2Vec などが使える）
vocab_size = 1000
embedding_dim = 300
# ランダム初期化（実際には事前学習済みベクトルを読み込むべき）
pretrained_weights = np.random.randn(vocab_size, embedding_dim).astype('float32')
# 埋め込み層を作成
embedding = nn.Embedding(vocab_size, embedding_dim)
# 事前学習済みの重みを読み込む
embedding.weight.data = torch.from_numpy(pretrained_weights)
# 埋め込み層を凍結（トレーニングに参加させない）
embedding.weight.requires_grad = False
print("埋め込み層のトレーニング可否:", embedding.weight.requires_grad)
print("埋め込み行列の形状:", embedding.weight.shape)
import torch.nn as nn
import numpy as np
# 事前学習済み単語ベクトルをシミュレート（実際には GloVe、Word2Vec などが使える）
vocab_size = 1000
embedding_dim = 300
# ランダム初期化（実際には事前学習済みベクトルを読み込むべき）
pretrained_weights = np.random.randn(vocab_size, embedding_dim).astype('float32')
# 埋め込み層を作成
embedding = nn.Embedding(vocab_size, embedding_dim)
# 事前学習済みの重みを読み込む
embedding.weight.data = torch.from_numpy(pretrained_weights)
# 埋め込み層を凍結（トレーニングに参加させない）
embedding.weight.requires_grad = False
print("埋め込み層のトレーニング可否:", embedding.weight.requires_grad)
print("埋め込み行列の形状:", embedding.weight.shape)
例 4: 埋め込みベクトルのノルムを制限する
max_norm を使用してベクトルのノルムを制限する:
h2 class="example">実例
import torch
import torch.nn as nn
# 埋め込みベクトルの最大ノルムを 1.0 に制限
embedding = nn.Embedding(1000, 64, max_norm=1.0)
# 入力
input_indices = torch.tensor([1, 2, 3])
# 元の重みのノルム
original_norm = embedding.weight.data.norm(dim=1)[:3]
print("元の重みのノルム:", original_norm.tolist())
# ルックアップ後のベクトルのノルム
output = embedding(input_indices)
output_norm = output.norm(dim=1)
print("出力ベクトルのノルム:", output_norm.tolist())
import torch.nn as nn
# 埋め込みベクトルの最大ノルムを 1.0 に制限
embedding = nn.Embedding(1000, 64, max_norm=1.0)
# 入力
input_indices = torch.tensor([1, 2, 3])
# 元の重みのノルム
original_norm = embedding.weight.data.norm(dim=1)[:3]
print("元の重みのノルム:", original_norm.tolist())
# ルックアップ後のベクトルのノルム
output = embedding(input_indices)
output_norm = output.norm(dim=1)
print("出力ベクトルのノルム:", output_norm.tolist())
例 5: 完全なテキスト分類モデル
埋め込み層を使ったテキスト分類:
実例
import torch
import torch.nn as nn
class TextClassifier(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes=2):
super(TextClassifier, self).__init__()
# 単語埋め込み層
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
# LSTM
self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
# 分類器
self.classifier = nn.Linear(hidden_dim, num_classes)
def forward(self, x):
# x: (batch, seq_len)
embedded = self.embedding(x) # (batch, seq_len, embed_dim)
# LSTM は最後の出力を取る
_, (hidden, _) = self.lstm(embedded)
hidden = hidden[-1] # 最後の層の隠れ状態
# 分類
logits = self.classifier(hidden)
return logits
# パラメータ
VOCAB_SIZE = 10000
EMBED_DIM = 128
HIDDEN_DIM = 256
model = TextClassifier(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM)
# 入力: batch=4、系列長=50
input_ids = torch.randint(1, VOCAB_SIZE, (4, 50))
output = model(input_ids)
print("モデル構造:")
print(model)
print("n入力形状:", input_ids.shape)
print("出力形状:", output.shape)
import torch.nn as nn
class TextClassifier(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes=2):
super(TextClassifier, self).__init__()
# 単語埋め込み層
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
# LSTM
self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
# 分類器
self.classifier = nn.Linear(hidden_dim, num_classes)
def forward(self, x):
# x: (batch, seq_len)
embedded = self.embedding(x) # (batch, seq_len, embed_dim)
# LSTM は最後の出力を取る
_, (hidden, _) = self.lstm(embedded)
hidden = hidden[-1] # 最後の層の隠れ状態
# 分類
logits = self.classifier(hidden)
return logits
# パラメータ
VOCAB_SIZE = 10000
EMBED_DIM = 128
HIDDEN_DIM = 256
model = TextClassifier(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM)
# 入力: batch=4、系列長=50
input_ids = torch.randint(1, VOCAB_SIZE, (4, 50))
output = model(input_ids)
print("モデル構造:")
print(model)
print("n入力形状:", input_ids.shape)
print("出力形状:", output.shape)
Embedding と EmbeddingBag の違い
| 種類 | 入力 | 出力 | 適用シーン |
|---|---|---|---|
nn.Embedding |
単語インデックス | 単語ベクトル系列 | 系列モデル、LSTM、Transformer |
nn.EmbeddingBag |
単語インデックス + オフセット | 集約されたベクトル | テキスト分類、高速処理 |
よくある質問
Q1: 埋め込み次元はどのように選ぶのですか？
- 小規模データセット: 50-100 次元
- 中規模データセット: 100-300 次元
- 大規模データセット: 300-500 次元
Q2: padding_idx の役割は何ですか？
指定されたインデックスの埋め込みベクトルをゼロベクトルに設定し、逆伝播時にその勾配を計算しません。
Q3: 埋め込み層はいつ凍結しますか？
事前学習済み単語ベクトルを使用する場合、通常は最初に凍結して訓練し、その後ファインチューニングします。
使用シーン
nn.Embedding主な応用シーンは以下の通りです:
- 単語ベクトル表現: 単語を密なベクトルに変換する
- テキスト分類: NLP モデルの入力層として
- 系列モデル: LSTM、GRU の入力
- レコメンドシステム: ユーザーとアイテムの埋め込み表現
注意: embedding.weight は学習可能なパラメータであり、オプティマイザーで直接訓練できます。
その他の拡張