PyTorch の埋め込みクラス（torch.nn.Embedding）の概略

この記事は、埋め込みクラス（torch.nn.Embedding）処理の具体的な例を見て、理解し易いようにすることが目的となる。機能の詳細や、利用用途などには触れない。

公式のドキュメントの翻訳は、以下のようになっている。

固定された辞書とサイズの埋め込みを格納する単純なルックアップテーブル。

このモジュールは、単語の埋め込みを保存し、インデックスを使用して取得するためによく使用されます。モジュールへの入力はインデックスのリストであり、出力は対応する単語の埋め込みです。

分かる人には、分かるのでしょう。私には、サッパリでした。

Gemini 先生に聞いてみました。

torch.nn.Embedding (埋め込み) は、PyTorchのニューラルネットワークモジュールで、固定サイズの辞書から単語やアイテムの埋め込みベクトルを取得するために使用されます。

役割

単語やアイテムの表現: 自然言語処理 (NLP) や推薦システムなどのタスクでは、単語やアイテムをベクトルとして表現することが重要です。 Embedding は、これらの要素を低次元の密なベクトル空間にマッピングすることで、意味的な関係や特徴を捉えることを可能にします。
効率的な計算: 大規模な語彙やアイテムセットを扱う場合、 one-hotベクトルのような高次元の疎な表現では計算量が大きくなってしまいます。 Embeddingは、低次元の密なベクトルを用いることで、計算効率を向上させます。

仕組み

初期化: Embeddingモジュールは、指定された語彙サイズと埋め込み次元数に基づいて、ランダムな値を持つ埋め込み行列を初期化します。
入力: Embeddingモジュールへの入力は、単語やアイテムのインデックスです。
出力: 指定されたインデックスに対応する埋め込みベクトルが返されます。

他にも、使い方、重要なパラメータ、応用例、補足など回答に含まれていたが割愛する。何となく分かった気になったので、実際の処理結果を確認してゆく。処理は、ログ出力を追加した品詞タグ付け用 LSTM モデルを利用する。

pos_tagger_lstm_test.py

import torch


class PosTaggerLSTM(torch.nn.Module):
    """品詞タグ付け用 LSTM モデル

    Attributes:
        embedding (torch.nn.Embedding): 単語埋め込み層：単語IDをベクトル表現に変換
        lstm (torch.nn.LSTM): LSTM 層
        linear (torch.nn.Linear): 線形変換層
    """

    def __init__(
        self,
        vocab_size: int,
        num_pos_tags: int,
        embedding_dim: int,
        num_layers: int = 1,
        bidirectional: bool = False
    ) -> None:
        """コンストラクタ

        Args:
            vocab_size (int): 単語の種類数 (語彙サイズ)
            num_pos_tags (int): 品詞の種類数
            embedding_dim (int): 単語ベクトルの次元数、LSTM の隠れ層の次元数
            num_layers (int): LSTM の層数
            bidirectional (bool): 双方向化するか否か
        """
        super(PosTaggerLSTM, self).__init__()
        # 単語埋め込み層：単語IDをベクトル表現に変換
        self.embedding = torch.nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        # LSTM 層
        self.lstm = torch.nn.LSTM(
            embedding_dim,
            embedding_dim,
            batch_first=True,
            num_layers=num_layers,
            bidirectional=bidirectional
        )
        # 双方向にする場合、線形変換層の次元を 2 倍にする
        linear_dim = embedding_dim * 2 if bidirectional else embedding_dim
        # 線形変換層：LSTM の出力を品詞数に変換
        self.linear = torch.nn.Linear(linear_dim, num_pos_tags)

    def forward(self, input_sequence: torch.Tensor) -> torch.Tensor:
        """品詞予測を行う

        Args:
            input_sequence (torch.Tensor): 入力系列 (単語IDのバッチ)

        Returns:
            torch.Tensor: 出力系列 (各単語の品詞予測)
        """
        print(f'input_sequence.shape: {input_sequence.shape}')
        print(f'input_sequence.dtype: {input_sequence.dtype}')
        print(f'input_sequence.device: {input_sequence.device}')
        print(f'input_sequence: {input_sequence}')
        # 単語埋め込み層
        embedded_sequence = self.embedding(input_sequence)
        print(f'embedded_sequence.shape: {embedded_sequence.shape}')
        print(f'embedded_sequence.dtype: {embedded_sequence.dtype}')
        print(f'embedded_sequence.device: {embedded_sequence.device}')
        print(f'embedded_sequence: {embedded_sequence}')
        # LSTM 層
        lstm_output, _ = self.lstm(embedded_sequence)
        print(f'lstm_output.shape: {lstm_output.shape}')
        print(f'lstm_output.dtype: {lstm_output.dtype}')
        print(f'lstm_output.device: {lstm_output.device}')
        print(f'lstm_output: {lstm_output}')
        # 線形変換層
        output = self.linear(lstm_output)
        print(f'output.shape: {output.shape}')
        print(f'output.dtype: {output.dtype}')
        print(f'output.device: {output.device}')
        print(f'output: {output}')
        return output


if __name__ == '__main__':
    device = torch.device('cpu')
    torch_model = PosTaggerLSTM(4, 3, 5).to(device)
    x = torch.tensor([[0, 1, 2, 3, 1]], dtype=torch.long, device=device)
    torch_out = torch_model(x)
    print(f'torch_out.shape: {torch_out.shape}')
    print(f'torch_out.dtype: {torch_out.dtype}')
    print(f'torch_out.device: {torch_out.device}')
    print(f'torch_out: {torch_out}')

実行ログから、埋め込み処理関連の部分だけ抜き出した。

input_sequence.shape: torch.Size([1, 5])
input_sequence.dtype: torch.int64
input_sequence.device: cpu
input_sequence: tensor([[0, 1, 2, 3, 1]])
embedded_sequence.shape: torch.Size([1, 5, 5])
embedded_sequence.dtype: torch.float32
embedded_sequence.device: cpu
embedded_sequence: tensor([[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0958,  0.2719,  0.2478,  0.9326, -0.7897],
         [ 1.6113,  0.4165,  0.2865,  0.4671, -0.2739],
         [-0.5610,  0.2716,  0.5785, -1.5057, -0.8792],
         [ 0.0958,  0.2719,  0.2478,  0.9326, -0.7897]]],
       grad_fn=<EmbeddingBackward0>)

埋め込みクラスの初期化は以下となる。

torch.nn.Embedding(4, 5, padding_idx=0)

埋め込み辞書のサイズが 4、各埋め込みベクトルのサイズが 5 となる。クラス内部で生成される各ベクトル値はランダムで、インスタンスを作成する度に別の値となる。辞書インデックス 0 は、パディングデータ（データサイズを合わせるためのデータ）なので、各ベクトル値が、0.0000 となっている。

[
    [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000 ],
    [ 0.0958,  0.2719,  0.2478,  0.9326, -0.7897 ],
    [ 1.6113,  0.4165,  0.2865,  0.4671, -0.2739 ],
    [-0.5610,  0.2716,  0.5785, -1.5057, -0.8792 ]
]

次に埋め込み処理では、以下のような処理となる。

self.embedding( [ [ 0, 1, 2, 3, 1 ] ] )

結果には、入力値のインデックスに対応する辞書ベクトルが返っている。入力のインデックス 1 は二つあり、どちらも同じ辞書ベクトルとなっている。

[
    [
        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000 ],
        [ 0.0958,  0.2719,  0.2478,  0.9326, -0.7897 ],
        [ 1.6113,  0.4165,  0.2865,  0.4671, -0.2739 ],
        [-0.5610,  0.2716,  0.5785, -1.5057, -0.8792 ],
        [ 0.0958,  0.2719,  0.2478,  0.9326, -0.7897 ]
    ]
]

分かってしまえば何てこと無い機能だが、初学者には躓き箇所になりがちなので、この記事で微力ながら役に立てればと願う。

いちりのつれづれ

このブログを検索

PyTorch の埋め込みクラス（torch.nn.Embedding）の概略

役割

仕組み