終于把 Transformer 算法搞懂了?。?！
作者：程序員小寒
我們將首先導(dǎo)入 PyTorch 庫以實現(xiàn)核心功能、導(dǎo)入神經(jīng)網(wǎng)絡(luò)模塊以創(chuàng)建神經(jīng)網(wǎng)絡(luò)、導(dǎo)入優(yōu)化模塊以訓(xùn)練網(wǎng)絡(luò)。
大家好，我是小寒
上次，我們從理論的角度給大家詳細介紹了什么是 Transformer 算法，并對 Transformer 的核心組件進行了完整的解讀。
今天將帶領(lǐng)大家使用 PyTorch 來從頭構(gòu)建一個 Transformer 模型
使用 PyTorch 構(gòu)建 Transformer 模型
1.導(dǎo)入必要的庫和模塊
我們將首先導(dǎo)入 PyTorch 庫以實現(xiàn)核心功能、導(dǎo)入神經(jīng)網(wǎng)絡(luò)模塊以創(chuàng)建神經(jīng)網(wǎng)絡(luò)、導(dǎo)入優(yōu)化模塊以訓(xùn)練網(wǎng)絡(luò)。
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import math
import copy2. 定義基本構(gòu)建塊
多頭注意力機制
多頭注意力機制計算序列中每對位置之間的注意力。
它由多個“注意力頭”組成，用于捕捉輸入序列的不同方面。
圖片
class MultiHeadAttention(nn.Module):
# d_model：輸入的維數(shù)
# num_heads：注意力頭的數(shù)量
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
# Ensure that the model dimension (d_model) is divisible by the number of heads
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
# Initialize dimensions
self.d_model = d_model # Model's dimension
self.num_heads = num_heads # Number of attention heads
self.d_k = d_model // num_heads # Dimension of each head's key, query, and value
# Linear layers for transforming inputs
self.W_q = nn.Linear(d_model, d_model) # Query transformation
self.W_k = nn.Linear(d_model, d_model) # Key transformation
self.W_v = nn.Linear(d_model, d_model) # Value transformation
self.W_o = nn.Linear(d_model, d_model) # Output transformation
# 縮放點積注意力機制
def scaled_dot_product_attention(self, Q, K, V, mask=None):
# Calculate attention scores
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
# Apply mask if provided (useful for preventing attention to certain parts like padding)
if mask is not None:
attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
# Softmax is applied to obtain attention probabilities
attn_probs = torch.softmax(attn_scores, dim=-1)
# Multiply by values to obtain the final output
output = torch.matmul(attn_probs, V)
return output
def split_heads(self, x):
# Reshape the input to have num_heads for multi-head attention
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):
# Combine the multiple heads back to original shape
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):
# Apply linear transformations and split heads
Q = self.split_heads(self.W_q(Q))
K = self.split_heads(self.W_k(K))
V = self.split_heads(self.W_v(V))
# Perform scaled dot-product attention
attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
# Combine heads and apply output transformation
output = self.W_o(self.combine_heads(attn_output))
return outputMultiHeadAttention 類封裝了 Transformer 模型中常用的多頭注意力機制，負責將輸入拆分成多個注意力頭，對每個注意力頭施加注意力，然后將結(jié)果組合起來，這樣模型就可以在不同尺度上捕捉輸入數(shù)據(jù)中的各種關(guān)系，提高模型的表達能力。
前饋網(wǎng)絡(luò)
class PositionWiseFeedForward(nn.Module):
def __init__(self, d_model, d_ff):
super(PositionWiseFeedForward, self).__init__()
self.fc1 = nn.Linear(d_model, d_ff)
self.fc2 = nn.Linear(d_ff, d_model)
self.relu = nn.ReLU()
def forward(self, x):
return self.fc2(self.relu(self.fc1(x)))PositionWiseFeedForward 類定義了一個前饋神經(jīng)網(wǎng)絡(luò)，它由兩個線性層組成，中間有一個 ReLU 激活函數(shù)。
位置編碼
位置編碼用于注入輸入序列中每個 token 的位置信息。它使用不同頻率的正弦和余弦函數(shù)來生成位置編碼。
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)]3.構(gòu)建編碼器模塊
圖片
EncoderLayer 類定義 Transformer 編碼器的單層。
它封裝了一個多頭自注意力機制，隨后是前饋神經(jīng)網(wǎng)絡(luò)，并根據(jù)需要應(yīng)用殘差連接、層規(guī)范化。
這些組件一起允許編碼器捕獲輸入數(shù)據(jù)中的復(fù)雜關(guān)系，并將其轉(zhuǎn)換為下游任務(wù)的有用表示。
通常，多個這樣的編碼器層堆疊在一起以形成 Transformer 模型的完整編碼器部分。
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)
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 x4.構(gòu)建解碼器模塊
圖片
DecoderLayer 類定義 Transformer 解碼器的單層。
它由多頭自注意力機制、多頭交叉注意力機制（關(guān)注編碼器的輸出）、前饋神經(jīng)網(wǎng)絡(luò)以及相應(yīng)的殘差連接、層規(guī)范化組成。
這種組合使解碼器能夠根據(jù)編碼器的表示生成有意義的輸出，同時考慮目標序列和源序列。
與編碼器一樣，多個解碼器層通常堆疊在一起以形成 Transformer 模型的完整解碼器部分。
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)
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 x5.結(jié)合編碼器和解碼器層來創(chuàng)建完整的 Transformer 網(wǎng)絡(luò)
圖片
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)
def generate_mask(self, src, tgt):
src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
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), diagnotallow=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)))
tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
enc_output = src_embedded
for enc_layer in self.encoder_layers:
enc_output = enc_layer(enc_output, src_mask)
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 outputTransformer 類將 Transformer 模型的各個組件整合在一起，包括嵌入、位置編碼、編碼器層和解碼器層。
它提供了一個方便的訓(xùn)練和推理接口，封裝了多頭注意力、前饋網(wǎng)絡(luò)和層規(guī)范化的復(fù)雜性。
訓(xùn)練 PyTorch Transformer 模型
1.樣本數(shù)據(jù)準備
為了便于說明，本例中將制作一個虛擬數(shù)據(jù)集。
但在實際情況下，將使用更大規(guī)模的數(shù)據(jù)集，并且該過程將涉及文本預(yù)處理以及為源語言和目標語言創(chuàng)建詞匯表映射。
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
transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout)
# Generate random sample data
src_data = torch.randint(1, src_vocab_size, (64, max_seq_length)) # (batch_size, seq_length)
tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length)) # (batch_size, seq_length)2.訓(xùn)練模型
接下來，將利用上述樣本數(shù)據(jù)訓(xùn)練模型。
transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout)
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(10):
optimizer.zero_grad()
output = transformer(src_data, tgt_data[:, :-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()}")
圖片
圖片
責任編輯：武曉燕
來源：
程序員學(xué)長
