PyTorch Transformer教程：从零开始实现注意力机制

在深度学习领域,Transformer模型已经成为处理序列数据的首选架构,特别是在自然语言处理任务中。本教程将详细介绍如何使用PyTorch从零开始实现一个Transformer模型,并应用于机器翻译任务。

Transformer的优势

与传统的循环神经网络(RNN)相比,Transformer具有以下显著优势:

并行计算能力强: Transformer可以并行处理序列中的所有元素,而RNN必须按顺序处理。
长距离依赖建模能力强: Transformer可以直接访问序列中的任意位置,而RNN需要通过中间状态传递信息。
表征能力强: 多头注意力机制使Transformer能够从多个角度理解输入序列。

Transformer vs RNN

Transformer的核心组件

多头注意力机制(Multi-Head Attention)

多头注意力是Transformer的核心,它允许模型同时关注序列的不同位置,从多个表示子空间学习信息。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        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 forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # Linear projections
        Q = self.W_q(query)
        K = self.W_k(key)  
        V = self.W_v(value)
        
        # Split into multiple heads
        Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # Scaled dot-product attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = F.softmax(scores, dim=-1)
        
        # Apply attention to values
        context = torch.matmul(attn, V)
        
        # Concatenate heads
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        
        # Final linear projection
        output = self.W_o(context)
        
        return output

位置编码(Positional Encoding)

由于Transformer没有固有的序列处理能力,我们需要通过位置编码为模型提供位置信息。

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, 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)
        
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        return x + self.pe[:x.size(0), :]

前馈神经网络(Feed Forward Network)

每个Transformer层都包含一个前馈神经网络,用于进一步处理注意力机制的输出。

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        
    def forward(self, x):
        return self.linear2(F.relu(self.linear1(x)))

构建Transformer模型

将上述组件组合,我们可以构建完整的Transformer模型:

class Transformer(nn.Module):
    def __init__(self, src_vocab, tgt_vocab, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout):
        super().__init__()
        self.encoder_embedding = nn.Embedding(src_vocab, d_model)
        self.decoder_embedding = nn.Embedding(tgt_vocab, 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)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        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 output

应用于机器翻译任务

我们以英语到德语的翻译为例,展示如何使用Transformer进行机器翻译。

数据准备

首先,我们需要准备平行语料库,并使用BPE(Byte Pair Encoding)等技术处理词表。

from torchtext.data import Field, BucketIterator
from torchtext.datasets import Multi30k

SRC = Field(tokenize = "spacy", tokenizer_language="en_core_web_sm", init_token = '<sos>', eos_token = '<eos>', lower = True)
TGT = Field(tokenize = "spacy", tokenizer_language="de_core_news_sm", init_token = '<sos>', eos_token = '<eos>', lower = True)

train_data, valid_data, test_data = Multi30k.splits(exts = ('.en', '.de'), fields = (SRC, TGT))

SRC.build_vocab(train_data, min_freq = 2)
TGT.build_vocab(train_data, min_freq = 2)

模型训练

接下来,我们定义损失函数和优化器,并开始训练过程。

model = Transformer(len(SRC.vocab), len(TGT.vocab), d_model=512, num_heads=8, num_layers=6, d_ff=2048, max_seq_length=100, dropout=0.1)
criterion = nn.CrossEntropyLoss(ignore_index=TGT.vocab.stoi['<pad>'])
optimizer = optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)

def train(model, iterator, optimizer, criterion, clip):
    model.train()
    epoch_loss = 0
    for i, batch in enumerate(iterator):
        src = batch.src
        tgt = batch.trg
        
        optimizer.zero_grad()
        output = model(src, tgt[:,:-1])
        output = output.contiguous().view(-1, output.shape[-1])
        tgt = tgt[:,1:].contiguous().view(-1)
        
        loss = criterion(output, tgt)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        optimizer.step()
        
        epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)

N_EPOCHS = 10
CLIP = 1

for epoch in range(N_EPOCHS):
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
    print(f'Epoch: {epoch+1:02} | Train Loss: {train_loss:.3f}')

翻译推理

最后,我们可以使用训练好的模型进行翻译:

def translate_sentence(sentence, src_field, trg_field, model, device, max_len = 50):
    model.eval()
    
    if isinstance(sentence, str):
        nlp = spacy.load('en_core_web_sm')
        tokens = [token.text.lower() for token in nlp(sentence)]
    else:
        tokens = [token.lower() for token in sentence]
    
    tokens = ['<sos>'] + tokens + ['<eos>']
    src_indexes = [src_field.vocab.stoi[token] for token in tokens]
    src_tensor = torch.LongTensor(src_indexes).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_indexes = [trg_field.vocab.stoi['<sos>']]

    for i in range(max_len):
        trg_tensor = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)
        trg_mask = model.make_trg_mask(trg_tensor)
        
        with torch.no_grad():
            output = model.decoder(trg_tensor, enc_src, trg_mask, src_mask)
        
        pred_token = output.argmax(2)[:,-1].item()
        trg_indexes.append(pred_token)

        if pred_token == trg_field.vocab.stoi['<eos>']:
            break
    
    trg_tokens = [trg_field.vocab.itos[i] for i in trg_indexes]
    
    return trg_tokens[1:]

example_sentence = "The sun is shining brightly today."
translation = translate_sentence(example_sentence, SRC, TGT, model, device)
print(f'Source: {example_sentence}')
print(f'Translation: {" ".join(translation)}')