Transformer模型自2017年由Vaswani等人在論文《Attention is All You Need》中提出以來，已經成為自然語言處理（NLP）領域的重要基石。Transformer的核心思想是使用自注意力機制（Self-Attention）來捕捉輸入序列中的全局依賴關系，從而避免了傳統RNN和LSTM模型中的序列依賴問題。本文將介紹如何使用PyTorch實現一個簡單的Transformer模型。
Transformer模型由編碼器（Encoder）和解碼器（Decoder）兩部分組成。每個編碼器和解碼器都由多個相同的層堆疊而成。每一層包含兩個主要子層：
此外，每個子層后面都會接一個殘差連接（Residual Connection）和層歸一化（Layer Normalization）。
下面我們將使用PyTorch實現一個簡單的Transformer模型。為了簡化，我們將實現一個只有一層編碼器和一層解碼器的Transformer。
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
self.num_heads = num_heads
self.d_model = d_model
self.depth = d_model // num_heads
self.wq = nn.Linear(d_model, d_model)
self.wk = nn.Linear(d_model, d_model)
self.wv = nn.Linear(d_model, d_model)
self.dense = nn.Linear(d_model, d_model)
def split_heads(self, x, batch_size):
x = x.view(batch_size, -1, self.num_heads, self.depth)
return x.permute(0, 2, 1, 3)
def forward(self, q, k, v, mask):
batch_size = q.size(0)
q = self.wq(q)
k = self.wk(k)
v = self.wv(v)
q = self.split_heads(q, batch_size)
k = self.split_heads(k, batch_size)
v = self.split_heads(v, batch_size)
scaled_attention, attention_weights = self.scaled_dot_product_attention(q, k, v, mask)
scaled_attention = scaled_attention.permute(0, 2, 1, 3)
concat_attention = scaled_attention.reshape(batch_size, -1, self.d_model)
output = self.dense(concat_attention)
return output, attention_weights
def scaled_dot_product_attention(self, q, k, v, mask):
matmul_qk = torch.matmul(q, k.transpose(-2, -1))
dk = torch.tensor(k.size(-1), dtype=torch.float32)
scaled_attention_logits = matmul_qk / torch.sqrt(dk)
if mask is not None:
scaled_attention_logits += (mask * -1e9)
attention_weights = F.softmax(scaled_attention_logits, dim=-1)
output = torch.matmul(attention_weights, v)
return output, attention_weights
class FeedForward(nn.Module):
def __init__(self, d_model, dff):
super(FeedForward, self).__init__()
self.linear1 = nn.Linear(d_model, dff)
self.linear2 = nn.Linear(dff, d_model)
def forward(self, x):
x = F.relu(self.linear1(x))
x = self.linear2(x)
return x
class EncoderLayer(nn.Module):
def __init__(self, d_model, num_heads, dff, rate=0.1):
super(EncoderLayer, self).__init__()
self.mha = MultiHeadAttention(d_model, num_heads)
self.ffn = FeedForward(d_model, dff)
self.layernorm1 = nn.LayerNorm(d_model)
self.layernorm2 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(rate)
self.dropout2 = nn.Dropout(rate)
def forward(self, x, mask):
attn_output, _ = self.mha(x, x, x, mask)
attn_output = self.dropout1(attn_output)
out1 = self.layernorm1(x + attn_output)
ffn_output = self.ffn(out1)
ffn_output = self.dropout2(ffn_output)
out2 = self.layernorm2(out1 + ffn_output)
return out2
class DecoderLayer(nn.Module):
def __init__(self, d_model, num_heads, dff, rate=0.1):
super(DecoderLayer, self).__init__()
self.mha1 = MultiHeadAttention(d_model, num_heads)
self.mha2 = MultiHeadAttention(d_model, num_heads)
self.ffn = FeedForward(d_model, dff)
self.layernorm1 = nn.LayerNorm(d_model)
self.layernorm2 = nn.LayerNorm(d_model)
self.layernorm3 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(rate)
self.dropout2 = nn.Dropout(rate)
self.dropout3 = nn.Dropout(rate)
def forward(self, x, enc_output, look_ahead_mask, padding_mask):
attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)
attn1 = self.dropout1(attn1)
out1 = self.layernorm1(x + attn1)
attn2, attn_weights_block2 = self.mha2(out1, enc_output, enc_output, padding_mask)
attn2 = self.dropout2(attn2)
out2 = self.layernorm2(out1 + attn2)
ffn_output = self.ffn(out2)
ffn_output = self.dropout3(ffn_output)
out3 = self.layernorm3(out2 + ffn_output)
return out3, attn_weights_block1, attn_weights_block2
class Transformer(nn.Module):
def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, rate=0.1):
super(Transformer, self).__init__()
self.encoder = nn.ModuleList([EncoderLayer(d_model, num_heads, dff, rate) for _ in range(num_layers)])
self.decoder = nn.ModuleList([DecoderLayer(d_model, num_heads, dff, rate) for _ in range(num_layers)])
self.embedding = nn.Embedding(input_vocab_size, d_model)
self.pos_encoding = self.positional_encoding(d_model)
self.final_layer = nn.Linear(d_model, target_vocab_size)
def positional_encoding(self, d_model):
position = torch.arange(0, 10000, dtype=torch.float32).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe = torch.zeros(10000, d_model)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
return pe
def forward(self, inp, tar, enc_padding_mask, look_ahead_mask, dec_padding_mask):
seq_len = inp.size(1)
inp = self.embedding(inp) + self.pos_encoding[:, :seq_len, :]
for i in range(self.num_layers):
inp = self.encoder[i](inp, enc_padding_mask)
seq_len = tar.size(1)
tar = self.embedding(tar) + self.pos_encoding[:, :seq_len, :]
for i in range(self.num_layers):
tar, _, _ = self.decoder[i](tar, inp, look_ahead_mask, dec_padding_mask)
final_output = self.final_layer(tar)
return final_output
本文介紹了如何使用PyTorch實現一個簡單的Transformer模型。我們首先實現了多頭自注意力機制和前饋神經網絡，然后構建了編碼器和解碼器層，最后將這些組件組合成一個完整的Transformer模型。雖然這個實現相對簡單，但它涵蓋了Transformer的核心思想，為進一步的優化和擴展提供了基礎。
在實際應用中，Transformer模型通常需要更多的層和更復雜的訓練技巧，例如學習率調度、梯度裁剪等。此外，Transformer模型還可以應用于各種任務，如機器翻譯、文本生成、圖像處理等。希望本文能為讀者提供一個良好的起點，幫助大家更好地理解和應用Transformer模型。
免責聲明：本站發布的內容（圖片、視頻和文字）以原創、轉載和分享為主，文章觀點不代表本網站立場，如果涉及侵權請聯系站長郵箱：is@yisu.com進行舉報，并提供相關證據，一經查實，將立刻刪除涉嫌侵權內容。
