秋念

详细复现 Transformer

模型分解

为了实现一个基本的Transformer模型，我们需要构建以下几个主要组件：

1. Multi-Head Attention：多头注意力机制，允许模型同时关注输入的不同部分。

2. Positional Encoding：位置编码，为模型提供序列中单词的位置信息。

3. Encoder and Decoder Layers：编码器和解码器层，分别处理输入数据和输出数据。

4. Feed-Forward Network：前馈网络，在注意力机制后对信息进行进一步处理。

导入必要的库

import torch  
import torch.nn as nn  
import torch.nn.functional as F  
import math  

实现多头注意力

class MultiHeadAttention(nn.Module):  
    def __init__(self, embed_size, heads):  
        super(MultiHeadAttention, self).__init__()  
        self.embed_size = embed_size  
        self.heads = heads  
        self.head_dim = embed_size // heads  
  
        assert (  
            self.head_dim * heads == embed_size  
        ), "Embedding size needs to be divisible by heads"  
  
        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)  
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)  
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)  
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)  
  
    def forward(self, values, keys, query, mask):  
        N = query.shape[0]  
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]  
  
        # Split the embedding into self.heads different pieces  
        values = values.reshape(N, value_len, self.heads, self.head_dim)  
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)  
        queries = query.reshape(N, query_len, self.heads, self.head_dim)  
  
        values = self.values(values)  
        keys = self.keys(keys)  
        queries = self.queries(queries)  
  
        # Einsum does matrix multiplication for query*keys for each training example  
        # with every other key, then scale, mask, and apply softmax  
        attention = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])  
        if mask is not None:  
            attention = attention.masked_fill(mask == 0, float("-1e20"))  
        attention = F.softmax(attention / math.sqrt(self.head_dim), dim=3)  
  
        out = torch.einsum("nhqk,nvhd->nqhd", [attention, values]).reshape(  
            N, query_len, self.heads * self.head_dim  
        )  
        out = self.fc_out(out)  
        return out 

实现位置编码

class PositionalEncoding(nn.Module):  
    def __init__(self, embed_size, max_len, device):  
        super(PositionalEncoding, self).__init__()  
        self.encoding = torch.zeros(max_len, embed_size).to(device)  
        self.encoding.requires_grad = False  # We do not want to update this during training  
  
        pos = torch.arange(0, max_len).unsqueeze(1).float().to(device)  
        _2i = torch.arange(0, embed_size, step=2).float().to(device)  
  
        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / embed_size)))  
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / embed_size)))  
  
    def forward(self, x): 
        # Add positional encoding to the input embeddings  
        seq_len = x.shape[1]  
        x = x + self.encoding[:seq_len, :]  
        return x  

搭建 Transformer 基本单元

class TransformerBlock(nn.Module):  
    def __init__(self, embed_size, heads, dropout, forward_expansion):  
        super(TransformerBlock, self).__init__()  
        self.attention = MultiHeadAttention(embed_size, heads)  
        self.norm1 = nn.LayerNorm(embed_size)  
        self.norm2 = nn.LayerNorm(embed_size)  
          
        self.feed_forward = nn.Sequential(  
            nn.Linear(embed_size, forward_expansion * embed_size),  
            nn.ReLU(),  
            nn.Linear(forward_expansion * embed_size, embed_size)  
        )  
          
        self.dropout = nn.Dropout(dropout)  
  
    def forward(self, value, key, query, mask):  
        attention = self.attention(value, key, query, mask)  
          
        # Add skip connection, followed by layer normalization  
        x = self.dropout(self.norm1(attention + query))  
        forward = self.feed_forward(x)  
          
        # Add skip connection, followed by layer normalization  
        out = self.dropout(self.norm2(forward + x))  
        return out  

搭建编码器

class Encoder(nn.Module):  
    def __init__(self,   
                 src_vocab_size,   
                 embed_size,   
                 num_layers,   
                 heads,   
                 device,   
                 forward_expansion,   
                 dropout,   
                 max_length):  
          
        super(Encoder, self).__init__()  
        self.embed_size = embed_size  
        self.device = device  
        self.word_embedding = nn.Embedding(src_vocab_size, embed_size)  
        self.positional_encoding = PositionalEncoding(embed_size, max_length, device)  
          
        self.layers = nn.ModuleList([  
            TransformerBlock(  
                embed_size,  
                heads,  
                dropout=dropout,  
                forward_expansion=forward_expansion,  
            )  
            for _ in range(num_layers)  
        ])  
          
        self.dropout = nn.Dropout(dropout)  
          
    def forward(self, x, mask):  
        out = self.dropout(self.word_embedding(x))  
        out = self.positional_encoding(out)  
          
        for layer in self.layers:  
            out = layer(out, out, out, mask)  
              
        return out  

搭建解码器块

class DecoderBlock(nn.Module):  
    def __init__(self, embed_size, heads, forward_expansion, dropout, device):  
        super(DecoderBlock, self).__init__()  
        self.attention = MultiHeadAttention(embed_size, heads)  
        self.norm = nn.LayerNorm(embed_size)  
        self.transformer_block = TransformerBlock(  
            embed_size, heads, dropout, forward_expansion  
        )  
        self.dropout = nn.Dropout(dropout)  
  
    def forward(self, x, value, key, src_mask, trg_mask):  
        attention = self.attention(x, x, x, trg_mask)  
        query = self.dropout(self.norm(attention + x))  
        out = self.transformer_block(value, key, query, src_mask)  
        return out  

搭建解码器

class Decoder(nn.Module):  
    def __init__(self,   
                 trg_vocab_size,   
                 embed_size,   
                 num_layers,   
                 heads,   
                 forward_expansion,   
                 dropout,   
                 device,   
                 max_length):  
          
        super(Decoder, self).__init__()  
        self.device = device  
        self.word_embedding = nn.Embedding(trg_vocab_size, embed_size)  
        self.positional_encoding = PositionalEncoding(embed_size, max_length, device)  
          
        self.layers = nn.ModuleList([  
            DecoderBlock(  
                embed_size,   
                heads,   
                forward_expansion,   
                dropout,   
                device  
            )  
            for _ in range(num_layers)  
        ])  
          
        self.fc_out = nn.Linear(embed_size, trg_vocab_size)  
        self.dropout = nn.Dropout(dropout)  
  
    def forward(self, x, enc_out, src_mask, trg_mask):  
        x = self.dropout(self.word_embedding(x))  
        x = self.positional_encoding(x)  
  
        for layer in self.layers:  
            x = layer(x, enc_out, enc_out, src_mask, trg_mask)  
  
        out = self.fc_out(x)  
        return out  

实现完整的 Transformer

class Transformer(nn.Module):  
    def __init__(self,   
                 src_vocab_size,   
                 trg_vocab_size,   
                 src_pad_idx,   
                 trg_pad_idx,   
                 embed_size=256,   
                 num_layers=6,   
                 forward_expansion=4,   
                 heads=8,   
                 dropout=0,   
                 device="cuda",   
                 max_length=100):  
          
        super(Transformer, self).__init__()  
  
        self.encoder = Encoder(  
            src_vocab_size,  
            embed_size,  
            num_layers,  
            heads,  
            device,  
            forward_expansion,  
            dropout,  
            max_length  
        )  
  
        self.decoder = Decoder(  
            trg_vocab_size,  
            embed_size,  
            num_layers,  
            heads,  
            forward_expansion,  
            dropout,  
            device,  
            max_length  
        )  
  
        self.src_pad_idx = src_pad_idx  
        self.trg_pad_idx = trg_pad_idx  
        self.device = device  
  
    def make_src_mask(self, src):  
        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)  
        return src_mask.to(self.device)  
  
    def make_trg_mask(self, trg):  
        N, trg_len = trg.shape  
        trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(  
            N, 1, trg_len, trg_len  
        )  
        return trg_mask.to(self.device)  
  
    def forward(self, src, trg):  
        src_mask = self.make_src_mask(src)  
        trg_mask = self.make_trg_mask(trg)  
  
        enc_src = self.encoder(src, src_mask)  
        out = self.decoder(trg, enc_src, src_mask, trg_mask)  
        return out  

测试 Transformer

# Initialize a transformer model  
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
src_vocab_size = 10000  
trg_vocab_size = 10000  
src_pad_idx = 0  
trg_pad_idx = 0  
  
model = Transformer(src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx).to(device)  
  
# Prepare some sample data (batch size = 1, sequence length = 10)  
src = torch.tensor([[1, 5, 6, 2, 0, 0, 0, 0, 0, 0]], device=device)  
trg = torch.tensor([[1, 7, 4, 3, 2, 0, 0, 0, 0, 0]], device=device)  
  
# Perform a forward pass  
out = model(src, trg[:, :-1])  
print(out.shape)  # Expected shape: (batch_size, trg_seq_length - 1, trg_vocab_size)  

详细复现 Transformer (续)

自定义 Dataset 类

from torch.utils.data import Dataset, DataLoader  
from torch.nn.utils.rnn import pad_sequence  

# 假设有一些数据和词汇表  
src_sentences = ["hello world", "this is an example"]  
trg_sentences = ["hola mundo", "esto es un ejemplo"]  
src_vocab = {'<pad>': 0, '<sos>': 1, '<eos>': 2, 'hello': 3, 'world': 4, 'this': 5, 'is': 6, 'an': 7, 'example': 8}  
trg_vocab = {'<pad>': 0, '<sos>': 1, '<eos>': 2, 'hola': 3, 'mundo': 4, 'esto': 5, 'es': 6, 'un': 7, 'ejemplo': 8}  
src_vocab_stoi = src_vocab  
src_vocab_itos = {index: token for token, index in src_vocab.items()}  
trg_vocab_stoi = trg_vocab  
trg_vocab_itos = {index: token for token, index in trg_vocab.items()}  
 
# 自定义 Dataset 类
class TranslationDataset(Dataset):  
    def __init__(self, src_sentences, trg_sentences, src_vocab, trg_vocab):  
        self.src_sentences = src_sentences  
        self.trg_sentences = trg_sentences  
  
        self.src_vocab = src_vocab  
        self.trg_vocab = trg_vocab  
  
    def __len__(self):  
        return len(self.src_sentences)  
  
    def __getitem__(self, index):  
        src_sentence = self.src_sentences[index]  
        trg_sentence = self.trg_sentences[index]  
  
        # 将句子转换为索引序列  
        src_indices = [self.src_vocab.stoi[token] for token in src_sentence.split()]  
        trg_indices = [self.trg_vocab.stoi[token] for token in trg_sentence.split()]  
  
        # 将索引序列转换为 PyTorch tensors  
        src_tensor = torch.tensor(src_indices)  
        trg_tensor = torch.tensor(trg_indices)  
  
        return src_tensor, trg_tensor  

创建 DataLoader 实例

# 创建 Dataset 实例  
dataset = TranslationDataset(src_sentences, trg_sentences, src_vocab_stoi, trg_vocab_stoi)  
  
# 定义一个函数来进行 padding，使 batch 中的所有句子长度一致  
def collate_fn(batch):  
    src_batch, trg_batch = zip(*batch)  
  
    src_batch = pad_sequence(src_batch, padding_value=src_vocab['<pad>'], batch_first=True)  
    trg_batch = pad_sequence(trg_batch, padding_value=trg_vocab['<pad>'], batch_first=True)  
  
    return src_batch, trg_batch  
  
# 创建 DataLoader 实例  
BATCH_SIZE = 2  
train_loader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, collate_fn=collate_fn) 

训练

# 初始化模型  
model = Transformer(src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, device=device)  
model = model.to(device)  
  
# 定义损失函数和优化器  
criterion = nn.CrossEntropyLoss(ignore_index=trg_pad_idx)  
optimizer = optim.Adam(model.parameters(), lr=0.001)  
  
# 训练模式  
model.train()  
  
num_epochs = 5  
for epoch in range(num_epochs):  
    for batch_idx, (src, trg) in enumerate(train_loader):  # 假设train_loader是你的数据加载器  
  
        src = src.to(device)  
        trg = trg.to(device)  
  
        # 前向传播  
        output = model(src, trg[:, :-1])  
        output = output.reshape(-1, output.shape[2])  
        trg = trg[:, 1:].reshape(-1)  
  
        # 计算损失  
        optimizer.zero_grad()  
        loss = criterion(output, trg)  
  
        # 反向传播  
        loss.backward()  
  
        # 更新模型参数  
        optimizer.step()  
  
        # 打印日志  
        if batch_idx % 100 == 0:  
            print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')  
  
print("Training finished.")  

推理

# 推理过程（生成文本）  
def translate_sentence(model, sentence, src_vocab, trg_vocab, src_pad_idx, device, max_length):  
    model.eval()  
  
    if isinstance(sentence, str):  
        tokens = [token for token in sentence.split(" ")]  # 假设用空格分词  
    else:  
        tokens = [token.lower() for token in sentence]  # 如果输入是token列表  
  
    # 将单词转换为索引  
    text_to_indices = [src_vocab.stoi[token] for token in tokens]  
    src_tensor = torch.LongTensor(text_to_indices).unsqueeze(0).to(device)  
  
    src_mask = model.make_src_mask(src_tensor)  
  
    with torch.no_grad():  
        enc_src = model.encoder(src_tensor, src_mask)  
  
    trg_indices = [trg_vocab.stoi["<sos>"]]  
  
    for i in range(max_length):  
        trg_tensor = torch.LongTensor(trg_indices).unsqueeze(0).to(device)  
  
        trg_mask = model.make_trg_mask(trg_tensor)  
  
        with torch.no_grad():  
            output = model.decoder(trg_tensor, enc_src, src_mask, trg_mask)  
            best_guess = output.argmax(2)[:, -1].item()  
          
        trg_indices.append(best_guess)  
  
        if best_guess == trg_vocab.stoi["<eos>"]:  
            break  
  
    # 将索引转换回单词  
    trg_tokens = [trg_vocab.itos[idx] for idx in trg_indices]  
  
    # 生成翻译结果，去除起始符号<sos>  
    translated_sentence = trg_tokens[1:]  
  
    return translated_sentence  

'''
假设 sentence 是需要翻译的源语言句子，src_vocab 和 trg_vocab 是源语言和目标语言的词汇表，包含stoi（string to index）和itos（index to string）的转换方法，假设max_length是句子的最大长度  
'''
sentence = "replace this with the sentence you want to translate"  
translated_sentence = translate_sentence(model, sentence, src_vocab, trg_vocab, src_pad_idx, device, max_length)  
  
print("Translated sentence:", " ".join(translated_sentence))  

---

参考（Reference）

• Let’s build GPT