PyTorch で Transformer モデルを構築する

Transformer は現代の機械学習において最も強力なモデルの一つです。

Transformer モデルは、自己注意機構(Self-Attention)に基づくディープラーニングアーキテクチャであり、自然言語処理(NLP)の分野を根本的に変え、BERT、GPT などの現代のディープラーニングモデルの基礎となっています。

Transformer は現代の NLP 分野における中核的なアーキテクチャであり、その強力な長距離依存関係のモデリング能力と効率的な並列計算の利点により、言語翻訳やテキスト要約などのタスクにおいて従来の長短期記憶(LSTM)ネットワークを凌駕しています。

Transformer についてまだ理解していない場合は、以下を参照してください:Transformer モデルの紹介。

PyTorch を使用して Transformer モデルを構築する

Transformer モデルを構築する手順は次のとおりです:

1、必要なライブラリとモジュールをインポートする

PyTorch のコアライブラリ、ニューラルネットワークモジュール、オプティマイザモジュール、データ処理ツール、および数学・オブジェクト複製モジュールをインポートして、モデルアーキテクチャの定義、データ管理、トレーニングプロセスをサポートします。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import math
import copy

説明:

  • torch:PyTorch のコアライブラリ。テンソル操作と自動微分に使用されます。

  • torch.nn:PyTorch のニューラルネットワークモジュール。さまざまな層と損失関数を含みます。

  • torch.optim:最適化アルゴリズムモジュール。Adam、SGD など。

  • math:数学関数ライブラリ。平方根などの計算に使用されます。

  • copy:オブジェクトのディープコピーに使用されます。

基本的な構成要素を定義する:マルチヘッドアテンション、位置単位フィードフォワードネットワーク、位置エンコーディング

マルチヘッドアテンション複数の「アテンションヘッド」を通じてシーケンス内の各位置ペア間の関係を計算し、入力シーケンスのさまざまな特徴やパターンを捉えることができます。

MultiHeadAttention クラスは、Transformer モデルでよく使われるマルチヘッドアテンション機構をカプセル化し、入力を複数のアテンションヘッドに分割し、各ヘッドにアテンションを適用して、その結果を組み合わせます。これにより、モデルは異なるスケールで入力データ内のさまざまな関係を捉え、表現力を高めることができます。

例

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model は num_heads で割り切れる必要があります"
       
        self.d_model = d_model    # モデル次元(例:512)
        self.num_heads = num_heads # アテンションヘッド数(例:8)
        self.d_k = d_model // num_heads # 各ヘッドの次元(例:64)
       
        # 線形変換層を定義(バイアスなし)
        self.W_q = nn.Linear(d_model, d_model) # クエリ変換
        self.W_k = nn.Linear(d_model, d_model) # キー変換
        self.W_v = nn.Linear(d_model, d_model) # バリュー変換
        self.W_o = nn.Linear(d_model, d_model) # 出力変換
       
    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        """
スケーリングドット積アテンションを計算する
入力形状:
            Q: (batch_size, num_heads, seq_length, d_k)
K, V: Q と同じ
出力形状: (batch_size, num_heads, seq_length, d_k)
        """

        # アテンションスコアを計算(Q と K のドット積)
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
       
        # マスクを適用(パディングマスクや将来情報マスクなど)
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
       
        # アテンションの重みを計算(softmax 正規化)
        attn_probs = torch.softmax(attn_scores, dim=-1)
       
        # バリューベクトルを加重合計する
        output = torch.matmul(attn_probs, V)
        return output
       
    def split_heads(self, x):
        """
入力テンソルを複数のヘッドに分割する
入力形状: (batch_size, seq_length, d_model)
出力形状: (batch_size, num_heads, seq_length, d_k)
        """

        batch_size, seq_length, d_model = x.size()
        return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)
       
    def combine_heads(self, x):
        """
複数のヘッドの出力を元の形状に結合する
入力形状: (batch_size, num_heads, seq_length, d_k)
出力形状: (batch_size, seq_length, d_model)
        """

        batch_size, _, seq_length, d_k = x.size()
        return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)
       
    def forward(self, Q, K, V, mask=None):
        """
フォワード伝播
入力形状: Q/K/V: (batch_size, seq_length, d_model)
出力形状: (batch_size, seq_length, d_model)
        """

        # 線形変換してマルチヘッドに分割
        Q = self.split_heads(self.W_q(Q)) # (batch, heads, seq_len, d_k)
        K = self.split_heads(self.W_k(K))
        V = self.split_heads(self.W_v(V))
       
        # アテンションを計算
        attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
       
        # マルチヘッドを結合して出力変換
        output = self.W_o(self.combine_heads(attn_output))
        return output

説明:

  • マルチヘッドアテンション機構:入力を複数のヘッドに分割し、各ヘッドが独立してアテンションを計算し、最後に結果を結合します。

  • スケーリングドット積アテンション:クエリとキーのドット積を計算し、スケーリング後に softmax を使用してアテンションの重みを計算し、最後にバリューを加重合計します。

  • マスク:無効な位置(パディング部分など)をマスクするために使用されます。

位置単位フィードフォワードネットワーク(Position-wise Feed-Forward Network)

例

class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(PositionWiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)  # 第1層全結合
        self.fc2 = nn.Linear(d_ff, d_model)  # 第2層全結合
        self.relu = nn.ReLU()  # 活性化関数

    def forward(self, x):
        # フィードフォワードネットワークの計算
        return self.fc2(self.relu(self.fc1(x)))

フィードフォワードネットワーク:2つの全結合層と1つの ReLU 活性化関数で構成され、アテンション機構の出力をさらに処理するために使用されます。

位置エンコーディング

位置エンコーディングは、入力シーケンス内の各トークンの位置情報を注入するために使用されます。

異なる周波数の正弦関数と余弦関数を使用して位置エンコーディングを生成します。

例

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_length):
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_seq_length, d_model)  # 位置エンコーディング行列を初期化
        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数位置では正弦関数を使用
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数位置では余弦関数を使用
        self.register_buffer('pe', pe.unsqueeze(0))  # バッファとして登録
       
    def forward(self, x):
        # 位置エンコーディングを入力に追加
        return x + self.pe[:, :x.size(1)]

エンコーダーブロック(Encoder Layer)を構築する

エンコーダー層:セルフアテンション機構とフィードフォワードネットワークを含み、各サブレイヤーの後に残差接続とレイヤー正規化が続きます。

例

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)  # セルフアテンション機構
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)  # フィードフォワードネットワーク
        self.norm1 = nn.LayerNorm(d_model)  # レイヤー正規化
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)  # Dropout
       
    def forward(self, x, mask):
        # セルフアテンション機構
        attn_output = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))  # 残差接続とレイヤー正規化
       
        # フィードフォワードネットワーク
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))  # 残差接続とレイヤー正規化
        return x

デコーダーモジュールを構築する

デコーダー層:セルフアテンション機構、クロスアテンション機構、およびフィードフォワードネットワークを含み、各サブレイヤーの後に残差接続とレイヤー正規化が続きます。

例

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)  # セルフアテンション機構
        self.cross_attn = MultiHeadAttention(d_model, num_heads)  # クロスアテンション機構
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)  # フィードフォワードネットワーク
        self.norm1 = nn.LayerNorm(d_model)  # レイヤー正規化
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)  # Dropout
       
    def forward(self, x, enc_output, src_mask, tgt_mask):
        # セルフアテンション機構
        attn_output = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_output))  # 残差接続とレイヤー正規化
       
        # クロスアテンション機構
        attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = self.norm2(x + self.dropout(attn_output))  # 残差接続とレイヤー正規化
       
        # フィードフォワードネットワーク
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout(ff_output))  # 残差接続とレイヤー正規化
        return x

完全な Transformer モデルを構築する

例

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout):
        super(Transformer, self).__init__()
        self.encoder_embedding = nn.Embedding(src_vocab_size, d_model)  # エンコーダー単語埋め込み
        self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model)  # デコーダー単語埋め込み
        self.positional_encoding = PositionalEncoding(d_model, max_seq_length)  # 位置エンコーディング

        # エンコーダー層とデコーダー層
        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])

        self.fc = nn.Linear(d_model, tgt_vocab_size)  # 最終全結合層
        self.dropout = nn.Dropout(dropout)  # Dropout

    def generate_mask(self, src, tgt):
        # ソースマスク:パディングトークンをマスク(パディングトークンのインデックスは0と仮定)
        # 形状:(batch_size, 1, 1, seq_length)
        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
   
        # ターゲットマスク:パディングトークンと将来情報をマスク
        # 形状:(batch_size, 1, seq_length, 1)
        tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3)
        seq_length = tgt.size(1)
        # 上三角行列マスクを生成し、デコード時に将来情報を参照しないようにする
        nopeak_mask = (1 - torch.triu(torch.ones(1, seq_length, seq_length), diagonal=1)).bool()
        tgt_mask = tgt_mask & nopeak_mask  # パディングマスクと将来情報マスクを結合
        return src_mask, tgt_mask

    def forward(self, src, tgt):
        # マスクを生成
        src_mask, tgt_mask = self.generate_mask(src, tgt)
       
        # エンコーダー部分
        src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
        enc_output = src_embedded
        for enc_layer in self.encoder_layers:
            enc_output = enc_layer(enc_output, src_mask)
       
        # デコーダー部分
        tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
        dec_output = tgt_embedded
        for dec_layer in self.decoder_layers:
            dec_output = dec_layer(dec_output, enc_output, src_mask, tgt_mask)
       
        # 最終出力
        output = self.fc(dec_output)
        return output

説明:

  • Transformer モデルエンコーダーとデコーダー部分を含み、各部分は複数の層を積み重ねて構成されます。

  • マスク生成無効な位置や未来の情報をマスクするために使用します。

  • 順伝播エンコーダーとデコーダーを順に通過し、最後に全結合層を通じて出力します。

モデル初期化パラメータの説明:

class Transformer(nn.Module):
    def __init__(
        self, 
        src_vocab_size,  # 源语言词汇表大小(如英文单词数)
        tgt_vocab_size,  # 目标语言词汇表大小(如中文单词数)
        d_model=512,     # 模型维度(每个词向量的长度)
        num_heads=8,     # 多头注意力的头数
        num_layers=6,    # 编码器/解码器的堆叠层数
        d_ff=2048,       # 前馈网络隐藏层维度
        max_seq_length=100, # 最大序列长度(用于位置编码)
        dropout=0.1      # Dropout概率
    ):

PyTorch Transformer モデルをトレーニングする

ランダムデータを使用してモデルをトレーニングし、損失を計算してパラメータを更新します。

例

# ハイパーパラメータ
src_vocab_size = 5000  # ソース語彙サイズ
tgt_vocab_size = 5000  # ターゲット語彙サイズ
d_model = 512  # モデル次元
num_heads = 8  # アテンションヘッド数
num_layers = 6  # エンコーダーとデコーダーの層数
d_ff = 2048  # フィードフォワードネットワークの内部次元
max_seq_length = 100  # 最大シーケンス長
dropout = 0.1  # Dropout 確率

# モデルの初期化
transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout)

# ランダムデータの生成
src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))  # ソースシーケンス
tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))  # ターゲットシーケンス

# 損失関数とオプティマイザの定義
criterion = nn.CrossEntropyLoss(ignore_index=0)  # パディング部分の損失を無視
optimizer = optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)

# トレーニングループ
transformer.train()
for epoch in range(100):
    optimizer.zero_grad()  # 勾配をクリアして蓄積を防ぐ
   
    # ターゲットシーケンスを入力する際、最後の単語を削除(次の単語を予測するため)
    output = transformer(src_data, tgt_data[:, :-1])  
   
    # 損失を計算する際、ターゲットシーケンスは2番目の単語から開始(つまり次の単語を予測)
    # output 形状: (batch_size, seq_length-1, tgt_vocab_size)
    # ターゲット形状: (batch_size, seq_length-1)
    loss = criterion(
        output.contiguous().view(-1, tgt_vocab_size),
        tgt_data[:, 1:].contiguous().view(-1)
    )
   
    loss.backward()        # 逆伝播
    optimizer.step()       # パラメータの更新
    print(f"Epoch: {epoch+1}, Loss: {loss.item()}")

モデル評価

評価プロセス:検証データ上で損失を計算し、モデルの性能を評価します。

例

transformer.eval()
# 検証データの生成
val_src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))
val_tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))
# 入力が英語と対応する中国語翻訳のバッチであると仮定(すでにインデックスに変換済み)
# サンプルデータ:
# src_data: [[3, 14, 25, ..., 0, 0], ...] # 英文学文(0はパディング記号)
# tgt_data: [[5, 20, 36, ..., 0, 0], ...] # 中国語翻訳(0はパディング記号)
# 注意:実際のアプリケーションでは、テキストに対してトークン化、エンコーディング、パディングなどの前処理が必要です
with torch.no_grad():
    val_output = transformer(val_src_data, val_tgt_data[:, :-1])
    val_loss = criterion(val_output.contiguous().view(-1, tgt_vocab_size), val_tgt_data[:, 1:].contiguous().view(-1))
    print(f"Validation Loss: {val_loss.item()}")
その他の拡張